core: (attr assembly fmt) add custom directive (pt7)

tests/irdl/test_attr_declarative_assembly_format.py

@@ -16,7 +16,14 @@
 )
 from xdsl.ir import Attribute, ParametrizedAttribute, TypeAttribute
 from xdsl.irdl import ParamAttrDef, irdl_attr_definition, param_def
-from xdsl.irdl.declarative_assembly_format import AttrFormatProgram
+from xdsl.irdl.declarative_assembly_format import (
+    AttrCustomDirective,
+    AttrFormatProgram,
+    AttrParsingState,
+    ParameterVariable,
+    PrintingState,
+    irdl_attr_custom_directive,
+)
 from xdsl.parser import AttrParser, Parser
 from xdsl.printer import Printer
 from xdsl.utils.exceptions import ParseError, PyRDLAttrDefinitionError
@@ -686,6 +693,57 @@ class StructOptType(ParametrizedAttribute, TypeAttribute):
     check_roundtrip(StructOptType, "req = i32", ctx)
 
 
+# ============================================================================
+# Integration tests — custom directive
+# ============================================================================
+
+
+def test_custom_directive():
+    """Custom directive with two parameters."""
+
+    @irdl_attr_custom_directive
+    class SwapDirective(AttrCustomDirective):
+        """Prints two parameters in reversed order."""
+
+        first: ParameterVariable
+        second: ParameterVariable
+
+        def parse(self, parser: AttrParser, state: AttrParsingState) -> bool:
+            self.second.parse(parser, state)
+            parser.parse_punctuation(",")
+            self.first.parse(parser, state)
+            return True
+
+        def print(
+            self,
+            printer: Printer,
+            state: PrintingState,
+            attr: ParametrizedAttribute,
+            /,
+        ) -> None:
+            self.second.print(printer, state, attr)
+            printer.print_string(", ")
+            state.should_emit_space = False
+            state.last_was_punctuation = True
+            self.first.print(printer, state, attr)
+
+    @irdl_attr_definition
+    class CustomType(ParametrizedAttribute, TypeAttribute):
+        name = "test_af.custom"
+        a: IntegerType = param_def()
+        b: IntegerType = param_def()
+        assembly_format = "custom<SwapDirective>($a, $b)"
+        custom_directives = (SwapDirective,)
+
+    ctx = Context(allow_unregistered=True)
+    ctx.load_attr_or_type(CustomType)
+    attr = CustomType(i32, i64)
+    printed = print_attr(attr)
+    assert printed == "!test_af.custom<i64, i32>"
+    parsed = parse_type(printed, ctx)
+    assert parsed == attr
+
+
 # ============================================================================
 # Integration tests — printing correctness
 # ============================================================================

xdsl/irdl/declarative_assembly_format.py

@@ -1698,6 +1698,40 @@ class QualifiedParameterVariable(ParameterVariable):
     """
 
 
+class AttrCustomDirective(AttrFormatDirective, ABC):
+    """A user defined assembly format directive for attribute/type formats.
+
+    Mirrors CustomDirective for operations. Custom directives can have
+    multiple parameters, whose types should be declared in the
+    `parameters` field.
+    """
+
+    parameters: ClassVar[dict[str, type[AttrFormatDirective]]]
+
+
+AttrCustomDirectiveInvT = TypeVar("AttrCustomDirectiveInvT", bound=AttrCustomDirective)
+
+
+def irdl_attr_custom_directive(
+    cls: type[AttrCustomDirectiveInvT],
+) -> type[AttrCustomDirectiveInvT]:
+    """Decorator used on custom attr directives to define the `parameters`
+    class variable."""
+
+    cls.parameters = {}
+    param_types = inspect.get_annotations(cls, eval_str=True)
+    for field_name, ty in param_types.items():
+        if is_const_classvar(field_name, ty, PyRDLError):
+            continue
+        if not issubclass(ty, AttrFormatDirective):
+            raise PyRDLError(
+                f"Custom attr directive {cls.__name__} has parameter "
+                f"{field_name} which is not an attr format directive."
+            )
+        cls.parameters[field_name] = ty
+    return dataclass(frozen=True)(cls)
+
+
 @dataclass(frozen=True)
 class ParamsDirective(AttrFormatDirective):
     """Captures all parameters of an attribute, printed comma-separated."""

xdsl/irdl/declarative_assembly_format_parser.py

@@ -954,6 +954,8 @@ def parse_format_directive(self) -> AttrFormatDirective:
             return self.parse_params_directive()
         if self.parse_optional_keyword("struct"):
             return self.parse_struct_directive()
+        if self.parse_optional_keyword("custom"):
+            return self.parse_attr_custom_directive()
         self.raise_error(f"unexpected token '{self._current_token.text}'")
 
     def parse_variable(self, inside_ref: bool = False) -> ParameterVariable:
@@ -1095,3 +1097,25 @@ def _parse_struct_variables(self) -> tuple[ParameterVariable, ...]:
 
     def parse_struct_directive(self) -> StructDirective:
         return StructDirective(self._parse_struct_variables())
+
+    def parse_attr_custom_directive(self) -> AttrFormatDirective:
+        """Parse `custom` `<` name `>` `(` args `)` for attr format."""
+        with self.in_angle_brackets():
+            name = self.parse_identifier()
+            if name not in self.attr_def.custom_directives:
+                self.raise_error(f"Custom attr directive '{name}' cannot be found.")
+        directive = self.attr_def.custom_directives[name]
+        param_types = directive.parameters
+        params: list[AttrFormatDirective] = []
+        with self.in_parens():
+            for i, (field_name, ty) in enumerate(param_types.items()):
+                if i:
+                    self.parse_punctuation(",")
+                param = self.parse_variable()
+                if not isinstance(param, ty):
+                    self.raise_error(
+                        f"{name}.{field_name} was expected to be of type "
+                        f"{ty.__name__}, but got {type(param).__name__}"
+                    )
+                params.append(param)
+        return directive(*params)