From f13a66984b4a9df5c30eb72bd0cb165c6e9d3e72 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20=C5=81opaciuk?= <szymon@lopaciuk.eu>
Date: Fri, 19 Apr 2024 10:37:16 +0200
Subject: [PATCH 1/5] Add formatters and structure of refs

---
 xdeps/refs.py  | 87 ++++++++++++++++++++++++++++++++++++++++++++++++--
 xdeps/tasks.py | 12 ++++++-
 2 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/xdeps/refs.py b/xdeps/refs.py
index 0766e26..e841046 100644
--- a/xdeps/refs.py
+++ b/xdeps/refs.py
@@ -5,6 +5,7 @@
 
 # cython: language_level=3
 
+import abc
 import builtins
 import cython
 import logging
@@ -91,6 +92,37 @@ def is_cythonized():
     return cython.compiled
 
 
+@cython.cclass
+class DefaultFormatter:
+    """A class used to govern the display of refs and expressions."""
+    @staticmethod
+    def repr_item(owner, key):
+        return f'{owner!r}[{key!r}]'
+
+    @staticmethod
+    def repr_attr(owner, key):
+        return f'{owner}.{key}'
+
+
+@cython.cclass
+class XMadFormatter:
+    """A class used to govern the display of refs and expressions."""
+
+    @staticmethod
+    def repr_item(owner, key):
+        if not isinstance(owner, Ref):
+            raise ValueError(
+                f'Cannot represent `{owner!r}[{key}]` in XMad syntax.'
+            )
+        return key
+
+    @staticmethod
+    def repr_attr(owner, key):
+        if not isinstance(owner, Ref):
+            return f'{owner!r}->{key}'
+        return key
+
+
 @cython.cclass
 class BaseRef:
     """
@@ -173,6 +205,10 @@ def _value(self):
     def _get_dependencies(self, out=None):
         return out or set()
 
+    @cython.ccall
+    def _item_formatter(self, item):
+        return f"{self!r}[{item!r}]"
+
     # order of precedence
     def __call__(self, *args, **kwargs):
         return CallRef(self, args, kwargs)
@@ -361,6 +397,10 @@ def __setattr__(self, attr, value):
                 # set a "built-in" attribute is during __cinit__ or when
                 # unpickling, and both of those cases are handled by Cython
                 # without the usual pythonic call to __setattr__.
+                # However, the side effect of this solution is that extension
+                # type attributes of cclasses only allow visibility='readonly'.
+                # Specifying visibility='public' will have no effect, and an
+                # explicit setter for these properties is necessary.
                 raise AttributeError(f"Attribute {attr} is read-only.")
 
         ref = AttrRef(self, attr, self._manager)
@@ -528,13 +568,27 @@ class Ref(MutableRef):
     """
     A reference in the top-level container.
     """
+    _show_owner: cython.char  # noqa
+
     def __cinit__(self, _owner, _key, _manager):
         # Cython automatically calls __cinit__ in the base classes
         self._hash = hash((type(self).__name__, _key))
+        self._show_owner = True
 
     def __repr__(self):
         return self._key
 
+    @cython.ccall
+    def _show_owner_in_repr(self, show=True):
+        self._show_owner = show
+
+    @cython.ccall
+    def _item_formatter(self, item):
+        if self._show_owner:
+            return super()._item_formatter(item)
+
+        return item
+
     def _get_value(self):
         return BaseRef._mk_value(self._owner)
 
@@ -589,7 +643,8 @@ def _set_value(self, value):
     def __repr__(self):
         assert self._owner is not None
         assert self._key is not None
-        return f"{self._owner}.{self._key}"
+        return self._manager.formatter.repr_attr(self._owner, self._key)
+        # return f"{self._owner}.{self._key}"
 
 
 @cython.cclass
@@ -610,7 +665,8 @@ def _set_value(self, value):
 
     def __repr__(self):
         assert self._owner is not None
-        return f"{self._owner}[{repr(self._key)}]"
+        # return self._owner._item_formatter(self._key)
+        return self._manager.formatter.repr_item(self._owner, self._key)
 
 
 @cython.cclass
@@ -694,6 +750,33 @@ def __repr__(self):
         return f"({self._op_str}{self._arg})"
 
 
+@cython.cclass
+class LiteralExpr(BaseRef):
+    """
+    This class simply holds a literal value. This is beneficial if we want to
+    'preserve' the history of expressions on a literal instead of evaluating
+    immediately.
+    """
+    _arg = cython.declare(object, visibility='readonly')
+
+    def __cinit__(self, arg):
+        self._arg = arg
+        self._hash = hash((self.__class__, self._arg))
+
+    def _get_value(self):
+        return self._arg
+
+    def _get_dependencies(self, out=None):
+        return out or set()
+
+    def __reduce__(self):
+        """Instruct pickle to not pickle the hash."""
+        return type(self), (self._arg,)
+
+    def __repr__(self):
+        return repr(self._arg)
+
+
 @cython.cclass
 class AddExpr(BinOpExpr):
     _op_str = '+'
diff --git a/xdeps/tasks.py b/xdeps/tasks.py
index f3a3bb4..a0fff46 100644
--- a/xdeps/tasks.py
+++ b/xdeps/tasks.py
@@ -8,7 +8,7 @@
 import logging
 from typing import Set, Hashable
 
-from .refs import BaseRef, MutableRef, ObjectAttrRef, Ref, RefCount
+from .refs import BaseRef, MutableRef, ObjectAttrRef, Ref, RefCount, DefaultFormatter
 from .utils import plot_pdot
 from .utils import AttrDict
 from .sorting import toposort
@@ -210,6 +210,9 @@ class Manager:
         Maps ref to all tasks that have ref as direct target.
     containers: dict
         Maps label to the controlled container.
+    structure: dict
+        Maps a ref to the refs that structurally depend on it (array element to
+        the array, attribute to the owner, etc.)
     """
 
     def __init__(self):
@@ -219,7 +222,9 @@ def __init__(self):
         self.rtasks = defaultdict(RefCount)
         self.deptasks = defaultdict(RefCount)
         self.tartasks = defaultdict(RefCount)
+        self.structure = defaultdict(set)
         self._tree_frozen = False
+        self.formatter = DefaultFormatter
 
     def ref(self, container=None, label="_"):
         """Return a ref to an instance (or dict) associated to a label.
@@ -280,6 +285,9 @@ def register(self, task):
                 # logger.info("T:%s modifies deps of T:%s",taskid,deptask)
                 self.rtasks[taskid].append(deptask)
 
+        if isinstance(taskid, MutableRef):
+            self.structure[taskid._owner].add(taskid)
+
     def unregister(self, taskid):
         """Unregister the task identified by taskid"""
         if self._tree_frozen:
@@ -299,6 +307,8 @@ def unregister(self, taskid):
             self.tartasks[tar].remove(taskid)
         if taskid in self.rtasks:
             del self.rtasks[taskid]
+        if isinstance(taskid, MutableRef):
+            self.structure[taskid._owner].remove(taskid)
         del self.tasks[taskid]
 
     def freeze_tree(self):

From 1677c1ea6bef791bf760539750a543e614af5a6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20=C5=81opaciuk?= <szymon@lopaciuk.eu>
Date: Wed, 31 Jul 2024 13:42:26 +0200
Subject: [PATCH 2/5] Explicit _format function on refs, more flexible
 copy_from_expr

---
 xdeps/refs.py  | 120 ++++++++++++++++++++++++++++++++-----------------
 xdeps/tasks.py |  23 +++++-----
 2 files changed, 92 insertions(+), 51 deletions(-)

diff --git a/xdeps/refs.py b/xdeps/refs.py
index e841046..9f52386 100644
--- a/xdeps/refs.py
+++ b/xdeps/refs.py
@@ -5,8 +5,9 @@
 
 # cython: language_level=3
 
-import abc
 import builtins
+from typing import Tuple, List, Any
+
 import cython
 import logging
 import math
@@ -77,6 +78,8 @@
     operator.ixor: "^",
 }
 
+Pair = Tuple[Any, Any]
+
 
 def is_ref(obj):
     return isinstance(obj, BaseRef)
@@ -92,36 +95,57 @@ def is_cythonized():
     return cython.compiled
 
 
-@cython.cclass
-class DefaultFormatter:
-    """A class used to govern the display of refs and expressions."""
-    @staticmethod
-    def repr_item(owner, key):
-        return f'{owner!r}[{key!r}]'
-
-    @staticmethod
-    def repr_attr(owner, key):
-        return f'{owner}.{key}'
-
-
 @cython.cclass
 class XMadFormatter:
     """A class used to govern the display of refs and expressions."""
+    scope = cython.declare('BaseRef', visibility='readonly')
 
-    @staticmethod
-    def repr_item(owner, key):
-        if not isinstance(owner, Ref):
+    def __init__(self, scope):
+        self.scope = scope
+
+    def repr_item(self, owner, key):
+        if not isinstance(owner, Ref) and owner != self.scope:
             raise ValueError(
-                f'Cannot represent `{owner!r}[{key}]` in XMad syntax.'
+                f'Cannot represent `{owner!r}[{key!r}]` in XMad syntax. '
+                f'Only top-level Ref elements are representable, but '
+                f'scope={self.scope}.'
             )
         return key
 
-    @staticmethod
-    def repr_attr(owner, key):
+    def repr_attr(self, owner, key):
         if not isinstance(owner, Ref):
-            return f'{owner!r}->{key}'
+            return f'{owner._formatted(self)}->{key}'
         return key
 
+    def repr_bin_op(self, op, lhs, rhs):
+        if isinstance(lhs, BaseRef):
+            lhs = lhs._formatted(self)
+        if isinstance(rhs, BaseRef):
+            rhs = rhs._formatted(self)
+        return f'({lhs} {op} {rhs})'
+
+    def repr_unary_op(self, op, arg):
+        if isinstance(arg, BaseRef):
+            arg = arg._formatted(self)
+        return f'({op}{arg})'
+
+    def repr_call(self, function, args: Tuple[Any], kwargs: Tuple[Pair, ...]):
+        all_args = []
+        for arg in args:
+            value = arg._formatted(self) if isinstance(arg, BaseRef) else repr(arg)
+            all_args.append(value)
+
+        for k, v in kwargs:
+            value = v._formatted(self) if isinstance(v, BaseRef) else repr(v)
+            all_args.append(f"{k}={value}")
+
+        if isinstance(function, BaseRef):
+            fname = function._formatted(self)
+        else:
+            fname = function.__name__
+
+        return f"{fname}({', '.join(all_args)})"
+
 
 @cython.cclass
 class BaseRef:
@@ -171,6 +195,10 @@ def __eq__(self, other):
         """
         return str(self) == str(other)
 
+    def _formatted(self, formatter: XMadFormatter):
+        """Return a formatted string representation of the ref/expression."""
+        return repr(self)
+
     def _eq(self, other):
         """Deferred expression for the equality of values of `self` and `other`."""
         return EqExpr(self, other)
@@ -205,10 +233,6 @@ def _value(self):
     def _get_dependencies(self, out=None):
         return out or set()
 
-    @cython.ccall
-    def _item_formatter(self, item):
-        return f"{self!r}[{item!r}]"
-
     # order of precedence
     def __call__(self, *args, **kwargs):
         return CallRef(self, args, kwargs)
@@ -568,26 +592,15 @@ class Ref(MutableRef):
     """
     A reference in the top-level container.
     """
-    _show_owner: cython.char  # noqa
-
     def __cinit__(self, _owner, _key, _manager):
         # Cython automatically calls __cinit__ in the base classes
         self._hash = hash((type(self).__name__, _key))
-        self._show_owner = True
 
     def __repr__(self):
         return self._key
 
-    @cython.ccall
-    def _show_owner_in_repr(self, show=True):
-        self._show_owner = show
-
-    @cython.ccall
-    def _item_formatter(self, item):
-        if self._show_owner:
-            return super()._item_formatter(item)
-
-        return item
+    def _formatted(self, _):
+        return repr(self)
 
     def _get_value(self):
         return BaseRef._mk_value(self._owner)
@@ -643,9 +656,12 @@ def _set_value(self, value):
     def __repr__(self):
         assert self._owner is not None
         assert self._key is not None
-        return self._manager.formatter.repr_attr(self._owner, self._key)
-        # return f"{self._owner}.{self._key}"
+        return f"{self._owner}.{self._key}"
 
+    def _formatted(self, formatter):
+        assert self._owner is not None
+        assert self._key is not None
+        return formatter.repr_attr(self._owner, self._key)
 
 @cython.cclass
 class ItemRef(MutableRef):
@@ -665,8 +681,11 @@ def _set_value(self, value):
 
     def __repr__(self):
         assert self._owner is not None
-        # return self._owner._item_formatter(self._key)
-        return self._manager.formatter.repr_item(self._owner, self._key)
+        return f'{self._owner!r}[{self._key!r}]'
+
+    def _formatted(self, formatter):
+        assert self._owner is not None
+        return formatter.repr_item(self._owner, self._key)
 
 
 @cython.cclass
@@ -710,6 +729,9 @@ def __reduce__(self):
     def __repr__(self):
         return f"({self._lhs} {self._op_str} {self._rhs})"
 
+    def _formatted(self, formatter):
+        return formatter.repr_bin_op(self._op_str, self._lhs, self._rhs)
+
 
 @cython.cclass
 class UnaryOpExpr(BaseRef):
@@ -749,6 +771,9 @@ def __reduce__(self):
     def __repr__(self):
         return f"({self._op_str}{self._arg})"
 
+    def _formatted(self, formatter):
+        return formatter.repr_unary_op(self._op_str, self._arg)
+
 
 @cython.cclass
 class LiteralExpr(BaseRef):
@@ -776,6 +801,9 @@ def __reduce__(self):
     def __repr__(self):
         return repr(self._arg)
 
+    def _formatted(self, _):
+        return repr(self)
+
 
 @cython.cclass
 class AddExpr(BinOpExpr):
@@ -1038,6 +1066,13 @@ def __repr__(self):
         op_symbol = OPERATOR_SYMBOLS.get(self._op, self._op.__name__)
         return f"{op_symbol}({self._arg})"
 
+    def _formatted(self, formatter):
+        if self._op in OPERATOR_SYMBOLS:
+            op_symbol = OPERATOR_SYMBOLS[self._op]
+            return formatter.repr_unary_op(self._op, self._arg)
+        else:
+            return formatter.repr_call(self._op, self._arg, ())
+
 
 @cython.cclass
 class CallRef(BaseRef):
@@ -1089,6 +1124,9 @@ def __repr__(self):
 
         return f"{fname}({args})"
 
+    def _formatted(self, formatter):
+        return formatter.repr_call(self._func, self._args, self._kwargs)
+
 
 class RefCount(dict):
     def append(self, item):
diff --git a/xdeps/tasks.py b/xdeps/tasks.py
index a0fff46..57ba59e 100644
--- a/xdeps/tasks.py
+++ b/xdeps/tasks.py
@@ -8,7 +8,7 @@
 import logging
 from typing import Set, Hashable
 
-from .refs import BaseRef, MutableRef, ObjectAttrRef, Ref, RefCount, DefaultFormatter
+from .refs import BaseRef, MutableRef, ObjectAttrRef, Ref, RefCount
 from .utils import plot_pdot
 from .utils import AttrDict
 from .sorting import toposort
@@ -224,7 +224,6 @@ def __init__(self):
         self.tartasks = defaultdict(RefCount)
         self.structure = defaultdict(set)
         self._tree_frozen = False
-        self.formatter = DefaultFormatter
 
     def ref(self, container=None, label="_"):
         """Return a ref to an instance (or dict) associated to a label.
@@ -285,8 +284,9 @@ def register(self, task):
                 # logger.info("T:%s modifies deps of T:%s",taskid,deptask)
                 self.rtasks[taskid].append(deptask)
 
-        if isinstance(taskid, MutableRef):
+        while isinstance(taskid, MutableRef) and not isinstance(taskid, Ref):
             self.structure[taskid._owner].add(taskid)
+            taskid = taskid._owner
 
     def unregister(self, taskid):
         """Unregister the task identified by taskid"""
@@ -360,20 +360,23 @@ def copy_expr_from(self, mgr, name, bindings=None):
         Copy expression from another manager
 
         name: one of toplevel container in mgr
-        bindings: dictionary mapping old container names into new container refs
+        bindings: dictionary mapping old container refs into new container refs
         """
         ref = mgr.containers[name]
-        if bindings is None:
-            cmbdct = self.containers
-        else:
-            cmbdct = dct_merge(self.containers, bindings)
-        self.load(mgr.iter_expr_tasks_owner(ref), cmbdct)
+        bindings = bindings or {}
+        renamed_tasks = []
+        for taskid, expr in mgr.iter_expr_tasks_owner(ref):
+            new_taskid = taskid
+            for source_ref, target_ref in bindings.items():
+                new_taskid = new_taskid.replace(str(source_ref), str(target_ref))
+            renamed_tasks.append((new_taskid, expr))
+        self.load(renamed_tasks, self.containers)
 
     def mk_fun(self, name, **kwargs):
         """Write a python function that executes a set of tasks in order of dependencies:
         name: name of the functions
         kwargs:
-            the keys are used to defined the argument name of the functions
+            the keys are used to define the argument name of the functions
             the values are the refs that will be set
         """
         varlist, start = list(zip(*kwargs.items()))

From 9f811da35a79e3a2591fcd3c8b0c6c1b68067434 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20=C5=81opaciuk?= <szymon@lopaciuk.eu>
Date: Thu, 1 Aug 2024 17:36:39 +0200
Subject: [PATCH 3/5] add a _set_to_expr method to MutableRef

---
 xdeps/refs.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/xdeps/refs.py b/xdeps/refs.py
index 9f52386..9349a76 100644
--- a/xdeps/refs.py
+++ b/xdeps/refs.py
@@ -217,6 +217,9 @@ def _mk_value(value):
     def _get_value(self):
         raise NotImplementedError
 
+    def _set_to_expr(self, expr):
+        self._manager.set_value(self, expr)
+
     @property
     def _value(self):
         try:

From 9d510067120e6955bf46bdd0b11d79c08a46a347 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20=C5=81opaciuk?= <szymon@lopaciuk.eu>
Date: Fri, 9 Aug 2024 15:29:51 +0200
Subject: [PATCH 4/5] Rename the formatted to XldFormatter

---
 xdeps/refs.py  | 4 ++--
 xdeps/tasks.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/xdeps/refs.py b/xdeps/refs.py
index 9349a76..97e26bb 100644
--- a/xdeps/refs.py
+++ b/xdeps/refs.py
@@ -96,7 +96,7 @@ def is_cythonized():
 
 
 @cython.cclass
-class XMadFormatter:
+class XldFormatter:
     """A class used to govern the display of refs and expressions."""
     scope = cython.declare('BaseRef', visibility='readonly')
 
@@ -195,7 +195,7 @@ def __eq__(self, other):
         """
         return str(self) == str(other)
 
-    def _formatted(self, formatter: XMadFormatter):
+    def _formatted(self, formatter: XldFormatter):
         """Return a formatted string representation of the ref/expression."""
         return repr(self)
 
diff --git a/xdeps/tasks.py b/xdeps/tasks.py
index 57ba59e..ccf6f0a 100644
--- a/xdeps/tasks.py
+++ b/xdeps/tasks.py
@@ -222,7 +222,7 @@ def __init__(self):
         self.rtasks = defaultdict(RefCount)
         self.deptasks = defaultdict(RefCount)
         self.tartasks = defaultdict(RefCount)
-        self.structure = defaultdict(set)
+        self.structure = defaultdict(RefCount)
         self._tree_frozen = False
 
     def ref(self, container=None, label="_"):
@@ -285,7 +285,7 @@ def register(self, task):
                 self.rtasks[taskid].append(deptask)
 
         while isinstance(taskid, MutableRef) and not isinstance(taskid, Ref):
-            self.structure[taskid._owner].add(taskid)
+            self.structure[taskid._owner].append(taskid)
             taskid = taskid._owner
 
     def unregister(self, taskid):

From fd694124ab208637d1e221072dfdff6d6f1c5f71 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20=C5=81opaciuk?= <szymon@lopaciuk.eu>
Date: Tue, 13 Aug 2024 09:38:20 +0200
Subject: [PATCH 5/5] Remove Manager.structure container

---
 xdeps/tasks.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/xdeps/tasks.py b/xdeps/tasks.py
index ccf6f0a..8d52f1b 100644
--- a/xdeps/tasks.py
+++ b/xdeps/tasks.py
@@ -222,7 +222,6 @@ def __init__(self):
         self.rtasks = defaultdict(RefCount)
         self.deptasks = defaultdict(RefCount)
         self.tartasks = defaultdict(RefCount)
-        self.structure = defaultdict(RefCount)
         self._tree_frozen = False
 
     def ref(self, container=None, label="_"):
@@ -284,10 +283,6 @@ def register(self, task):
                 # logger.info("T:%s modifies deps of T:%s",taskid,deptask)
                 self.rtasks[taskid].append(deptask)
 
-        while isinstance(taskid, MutableRef) and not isinstance(taskid, Ref):
-            self.structure[taskid._owner].append(taskid)
-            taskid = taskid._owner
-
     def unregister(self, taskid):
         """Unregister the task identified by taskid"""
         if self._tree_frozen:
@@ -307,8 +302,6 @@ def unregister(self, taskid):
             self.tartasks[tar].remove(taskid)
         if taskid in self.rtasks:
             del self.rtasks[taskid]
-        if isinstance(taskid, MutableRef):
-            self.structure[taskid._owner].remove(taskid)
         del self.tasks[taskid]
 
     def freeze_tree(self):