core: Implement HasAncestor.

tests/filecheck/dialects/stencil/invalid.mlir

@@ -115,7 +115,7 @@ builtin.module {
   }
 }
 
- // CHECK: 'stencil.access' expects parent op 'stencil.apply'
+ // CHECK: 'stencil.access' expects ancestor op 'stencil.apply'
 
 // -----
 

tests/test_traits.py

@@ -37,6 +37,7 @@
     traits_def,
 )
 from xdsl.traits import (
+    HasAncestor,
     HasParent,
     OptionalSymbolOpInterface,
     SymbolOpInterface,
@@ -449,6 +450,25 @@ def test_lazy_parent():
     assert op.traits == frozenset([HasParent(TestOp)])
 
 
+@irdl_op_definition
+class AncestorOp(IRDLOperation):
+    name = "test.ancestor"
+
+    traits = frozenset((HasAncestor(TestOp),))
+
+
+def test_has_ancestor():
+    op = AncestorOp()
+
+    assert op.get_traits_of_type(HasAncestor) == [HasAncestor(TestOp)]
+    assert op.has_trait(HasAncestor, (TestOp,))
+
+    with pytest.raises(
+        VerifyException, match="'test.ancestor' expects ancestor op 'test.test'"
+    ):
+        op.verify()
+
+
 def test_insert_or_update():
     @irdl_op_definition
     class SymbolTableOp(IRDLOperation):

xdsl/dialects/stencil.py

@@ -53,7 +53,7 @@
 )
 from xdsl.parser import AttrParser, Parser
 from xdsl.printer import Printer
-from xdsl.traits import HasParent, IsolatedFromAbove, IsTerminator
+from xdsl.traits import HasAncestor, HasParent, IsolatedFromAbove, IsTerminator
 from xdsl.utils.exceptions import VerifyException
 from xdsl.utils.hints import isa
 
@@ -728,7 +728,7 @@ class AccessOp(IRDLOperation):
         )
     )
 
-    traits = frozenset([HasParent(ApplyOp)])
+    traits = frozenset([HasAncestor(ApplyOp)])
 
     def print(self, printer: Printer):
         printer.print(" ")
@@ -740,7 +740,8 @@ def print(self, printer: Printer):
         )
 
         # IRDL-enforced, not supposed to use custom syntax if not veriied
-        apply = cast(ApplyOp, self.parent_op())
+        trait = cast(HasAncestor, AccessOp.get_trait(HasAncestor, (ApplyOp,)))
+        apply = cast(ApplyOp, trait.get_ancestor(self))
 
         mapping = self.offset_mapping
         if mapping is None:
@@ -825,8 +826,9 @@ def get(
         )
 
     def verify_(self) -> None:
-        apply = self.parent_op()
-        # As promised by HasParent(ApplyOp)
+        # As promised by HasAncestor(ApplyOp)
+        trait = cast(HasAncestor, AccessOp.get_trait(HasAncestor, (ApplyOp,)))
+        apply = trait.get_ancestor(self)
         assert isinstance(apply, ApplyOp)
 
         # TODO This should be handled by infra, having a way to verify things on ApplyOp

xdsl/traits.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+from collections.abc import Iterator
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, TypeVar
 
@@ -68,6 +69,42 @@ def verify(self, op: Operation) -> None:
         raise VerifyException(f"'{op.name}' expects parent op to be one of {names}")
 
 
+@dataclass(frozen=True)
+class HasAncestor(OpTrait):
+    """
+    Constraint the operation to have a specific operation as ancestor, i.e. transitive
+    parent.
+    """
+
+    parameters: tuple[type[Operation], ...]
+
+    def __init__(self, head_param: type[Operation], *tail_params: type[Operation]):
+        super().__init__((head_param, *tail_params))
+
+    def verify(self, op: Operation) -> None:
+        if self.get_ancestor(op) is None:
+            if len(self.parameters) == 1:
+                raise VerifyException(
+                    f"'{op.name}' expects ancestor op '{self.parameters[0].name}'"
+                )
+            names = ", ".join(f"'{p.name}'" for p in self.parameters)
+            raise VerifyException(
+                f"'{op.name}' expects ancestor op to be one of {names}"
+            )
+
+    def walk_ancestors(self, op: Operation) -> Iterator[Operation]:
+        """Iterates over the ancestors of an operation, including the input"""
+        curr = op
+        yield curr
+        while (curr := curr.parent_op()) is not None:
+            yield curr
+
+    def get_ancestor(self, op: Operation) -> Operation | None:
+        ancestors = self.walk_ancestors(op)
+        matching_ancestors = (a for a in ancestors if isinstance(a, self.parameters))
+        return next(matching_ancestors, None)
+
+
 class IsTerminator(OpTrait):
     """
     This trait provides verification and functionality for operations that are