core: (attr assembly fmt) add optional group directive (pt3)

tests/irdl/test_attr_declarative_assembly_format.py

@@ -9,6 +9,7 @@
 from xdsl.context import Context
 from xdsl.dialects.builtin import (
     IntegerType,
+    NoneAttr,
     i1,
     i32,
     i64,
@@ -130,6 +131,22 @@ class ThreeParamType(ParametrizedAttribute, TypeAttribute):
     assembly_format = "$a `,` $b `,` $c"
 
 
+@irdl_attr_definition
+class OptionalParamType(ParametrizedAttribute, TypeAttribute):
+    name = "test_af.optional_param"
+    value: IntegerType = param_def()
+    opt: IntegerType | NoneAttr = param_def()
+    assembly_format = "$value (`,` $opt^)?"
+
+
+@irdl_attr_definition
+class ElseBranchType(ParametrizedAttribute, TypeAttribute):
+    name = "test_af.else_branch"
+    a: IntegerType | NoneAttr = param_def()
+    b: IntegerType | NoneAttr = param_def()
+    assembly_format = "($a^) : (`fallback` $b)?"
+
+
 # ============================================================================
 # Unit test helpers (call AttrFormatProgram directly, no @irdl_attr_definition)
 # ============================================================================
@@ -229,6 +246,8 @@ def ctx() -> Context:
     ctx.load_attr_or_type(PairType)
     ctx.load_attr_or_type(PairAttr)
     ctx.load_attr_or_type(ThreeParamType)
+    ctx.load_attr_or_type(OptionalParamType)
+    ctx.load_attr_or_type(ElseBranchType)
     return ctx
 
 
@@ -474,6 +493,58 @@ class NewlineType(ParametrizedAttribute, TypeAttribute):
     assert printed == "!test_af.newline<i32\ni64>"
 
 
+# ============================================================================
+# Integration tests — optional groups
+# ============================================================================
+
+
+@pytest.mark.parametrize("body", ["i32, i64", "i32"])
+def test_optional_param_roundtrip(body: str, ctx: Context):
+    check_roundtrip(OptionalParamType, body, ctx)
+
+
+def test_optional_param_present_constructs(ctx: Context):
+    parsed = parse_type("!test_af.optional_param<i32, i64>", ctx)
+    assert isinstance(parsed, OptionalParamType)
+    assert parsed.value == i32
+    assert parsed.opt == i64
+
+
+def test_optional_param_absent_constructs(ctx: Context):
+    parsed = parse_type("!test_af.optional_param<i32>", ctx)
+    assert isinstance(parsed, OptionalParamType)
+    assert parsed.value == i32
+    assert isinstance(parsed.opt, NoneAttr)
+
+
+def test_else_branch_then(ctx: Context):
+    check_roundtrip(ElseBranchType, "i32", ctx)
+
+
+def test_else_branch_else(ctx: Context):
+    check_roundtrip(ElseBranchType, "fallback i64", ctx)
+
+
+def test_else_branch_then_constructs(ctx: Context):
+    parsed = parse_type("!test_af.else_branch<i32>", ctx)
+    assert isinstance(parsed, ElseBranchType)
+    assert parsed.a == i32
+    assert isinstance(parsed.b, NoneAttr)
+
+
+@pytest.mark.parametrize(
+    "attr, expected",
+    [
+        (OptionalParamType(i32, i64), "!test_af.optional_param<i32, i64>"),
+        (OptionalParamType(i32, NoneAttr()), "!test_af.optional_param<i32>"),
+        (ElseBranchType(i32, NoneAttr()), "!test_af.else_branch<i32>"),
+        (ElseBranchType(NoneAttr(), i64), "!test_af.else_branch<fallback i64>"),
+    ],
+)
+def test_optional_print_correctness(attr: Attribute, expected: str):
+    assert print_attr(attr) == expected
+
+
 # ============================================================================
 # Integration tests — printing correctness
 # ============================================================================

xdsl/irdl/declarative_assembly_format.py

@@ -1689,6 +1689,52 @@ def set_empty(self, state: AttrParsingState) -> None:
         state.parameters[self.index] = NoneAttr()
 
 
+@dataclass(frozen=True)
+class AttrOptionalGroupDirective(AttrFormatDirective):
+    """An optional group in attribute assembly format."""
+
+    anchor: AttrFormatDirective
+    then_whitespace: tuple[AttrWhitespaceDirective, ...]
+    then_first: AttrFormatDirective
+    then_elements: tuple[AttrFormatDirective, ...]
+    else_elements: tuple[AttrFormatDirective, ...]
+
+    def parse(self, parser: AttrParser, state: AttrParsingState) -> bool:
+        if ret := self.then_first.parse_optional(parser, state):
+            for element in self.then_elements:
+                element.parse(parser, state)
+            for element in self.else_elements:
+                element.set_empty(state)
+        else:
+            self.then_first.set_empty(state)
+            for element in self.then_elements:
+                element.set_empty(state)
+            for element in self.else_elements:
+                element.parse(parser, state)
+        return ret
+
+    def print(
+        self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, /
+    ) -> None:
+        if self.anchor.is_present(attr):
+            for element in (
+                *self.then_whitespace,
+                self.then_first,
+                *self.then_elements,
+            ):
+                element.print(printer, state, attr)
+        else:
+            for element in self.else_elements:
+                element.print(printer, state, attr)
+
+    def set_empty(self, state: AttrParsingState) -> None:
+        self.then_first.set_empty(state)
+        for element in self.then_elements:
+            element.set_empty(state)
+        for element in self.else_elements:
+            element.set_empty(state)
+
+
 @dataclass(frozen=True)
 class AttrFormatProgram:
     """

xdsl/irdl/declarative_assembly_format_parser.py

@@ -51,6 +51,7 @@
     AttrFormatProgram,
     AttributeVariable,
     AttrKeywordDirective,
+    AttrOptionalGroupDirective,
     AttrPunctuationDirective,
     AttrWhitespaceDirective,
     DenseArrayAttributeVariable,
@@ -940,6 +941,8 @@ def verify_parameters(self) -> None:
     def parse_format_directive(self) -> AttrFormatDirective:
         if self._current_token.text == "`":
             return self.parse_keyword_or_punctuation()
+        if self.parse_optional_punctuation("("):
+            return self.parse_optional_group()
         if self._current_token.text == "$":
             return self.parse_variable()
         self.raise_error(f"unexpected token '{self._current_token.text}'")
@@ -1001,3 +1004,52 @@ def parse_keyword_or_punctuation(self) -> AttrFormatDirective:
 
         self.parse_characters("`")
         return AttrKeywordDirective(ident)
+
+    def parse_optional_group(self) -> AttrFormatDirective:
+        then_elements = tuple[AttrFormatDirective, ...]()
+        else_elements = tuple[AttrFormatDirective, ...]()
+        anchor: AttrFormatDirective | None = None
+
+        while not self.parse_optional_punctuation(")"):
+            then_elements += (self.parse_format_directive(),)
+            if self.parse_optional_keyword("^"):
+                if anchor is not None:
+                    self.raise_error("An optional group can only have one anchor.")
+                anchor = then_elements[-1]
+
+        if self.parse_optional_punctuation(":"):
+            self.parse_punctuation("(")
+            while not self.parse_optional_punctuation(")"):
+                else_elements += (self.parse_format_directive(),)
+
+        self.parse_punctuation("?")
+
+        first_non_whitespace_index = None
+        for i, x in enumerate(then_elements):
+            if not isinstance(x, AttrWhitespaceDirective):
+                first_non_whitespace_index = i
+                break
+
+        if first_non_whitespace_index is None:
+            self.raise_error("An optional group must have a non-whitespace directive")
+        if anchor is None:
+            self.raise_error("Every optional group must have an anchor.")
+        if not then_elements[first_non_whitespace_index].is_optional_like():
+            self.raise_error(
+                "First element of an optional group must be optionally parsable."
+            )
+        if not anchor.is_anchorable():
+            self.raise_error(
+                "An optional group's anchor must be an anchorable directive."
+            )
+
+        return AttrOptionalGroupDirective(
+            anchor,
+            cast(
+                tuple[AttrWhitespaceDirective, ...],
+                then_elements[:first_non_whitespace_index],
+            ),
+            then_elements[first_non_whitespace_index],
+            then_elements[first_non_whitespace_index + 1 :],
+            else_elements,
+        )