feat[eve]: Preserve annex in custom visitors

src/gt4py/eve/visitors.py

@@ -121,6 +121,21 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
         return None
 
 
+def _preserve_annex(
+    node: concepts.Node, new_node: concepts.Node, preserved_annex_attrs: tuple[str, ...]
+) -> None:
+    if preserved_annex_attrs and (old_annex := getattr(node, "__node_annex__", None)):
+        # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
+        new_annex_dict = new_node.annex.__dict__
+        for key in preserved_annex_attrs:
+            if (value := getattr(old_annex, key, NOTHING)) is not NOTHING:
+                # Note: The annex value of the new node might not be equal
+                # (in the sense that an equality comparison returns false),
+                # but in the context of the pass, they are equivalent.
+                # Therefore, we don't assert equality here.
+                new_annex_dict[key] = value
+
+
 class NodeTranslator(NodeVisitor):
     """Special `NodeVisitor` to translate nodes and trees.
 
@@ -158,13 +173,7 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
                     if (new_child := self.visit(child, **kwargs)) is not NOTHING
                 }
             )
-            if self.PRESERVED_ANNEX_ATTRS and (old_annex := getattr(node, "__node_annex__", None)):
-                # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
-                new_annex_dict = new_node.annex.__dict__
-                for key in self.PRESERVED_ANNEX_ATTRS:
-                    if (value := getattr(old_annex, key, NOTHING)) is not NOTHING:
-                        assert key not in new_annex_dict
-                        new_annex_dict[key] = value
+            _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
 
             return new_node
 
@@ -189,3 +198,11 @@ def generic_visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
             )
 
         return copy.deepcopy(node, memo=memo)
+
+    def visit(self, node: concepts.RootNode, **kwargs: Any) -> Any:
+        new_node = super().visit(node, **kwargs)
+
+        if isinstance(node, concepts.Node) and isinstance(new_node, concepts.Node):
+            _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
+
+        return new_node

src/gt4py/next/iterator/transforms/infer_domain.py

@@ -21,7 +21,7 @@
     domain_utils,
     ir_makers as im,
 )
-from gt4py.next.iterator.transforms import trace_shifts
+from gt4py.next.iterator.transforms import constant_folding, trace_shifts
 from gt4py.next.utils import flatten_nested_tuple, tree_map
 
 
@@ -436,8 +436,12 @@ def _infer_stmt(
     **kwargs: Unpack[InferenceOptions],
 ):
     if isinstance(stmt, itir.SetAt):
+        # constant fold once otherwise constant folding after domain inference might create (syntactic) differences
+        # between the domain stored in IR and in the annex
+        domain = constant_folding.ConstantFolding.apply(stmt.domain)
+
         transformed_call, _ = infer_expr(
-            stmt.expr, domain_utils.SymbolicDomain.from_expr(stmt.domain), **kwargs
+            stmt.expr, domain_utils.SymbolicDomain.from_expr(domain), **kwargs
         )
 
         return itir.SetAt(

tests/eve_tests/unit_tests/test_visitors.py

@@ -24,3 +24,20 @@ class SampleTranslator(eve.NodeTranslator):
     assert translated_node.annex.foo == 1
     assert translated_node.annex.bar is None
     assert not hasattr(translated_node.annex, "baz")
+
+
+def test_annex_preservation_translated_node(compound_node: eve.Node):
+    compound_node.annex.foo = 1
+    compound_node.annex.baz = 2
+
+    class SampleTranslator(eve.NodeTranslator):
+        PRESERVED_ANNEX_ATTRS = ("foo",)
+
+        def visit_Node(self, node: eve.Node):
+            # just return an empty node, we care about the annex only anyway
+            return eve.Node()
+
+    translated_node = SampleTranslator().visit(compound_node)
+
+    assert translated_node.annex.foo == 1
+    assert not hasattr(translated_node.annex, "baz")