diff --git a/tests/irdl/test_attr_declarative_assembly_format.py b/tests/irdl/test_attr_declarative_assembly_format.py index 841cd5c791..ba5cf20631 100644 --- a/tests/irdl/test_attr_declarative_assembly_format.py +++ b/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" +# ============================================================================ +# 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", 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", 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", 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"), + (OptionalParamType(i32, NoneAttr()), "!test_af.optional_param"), + (ElseBranchType(i32, NoneAttr()), "!test_af.else_branch"), + (ElseBranchType(NoneAttr(), i64), "!test_af.else_branch"), + ], +) +def test_optional_print_correctness(attr: Attribute, expected: str): + assert print_attr(attr) == expected + + # ============================================================================ # Integration tests — printing correctness # ============================================================================ diff --git a/xdsl/irdl/declarative_assembly_format.py b/xdsl/irdl/declarative_assembly_format.py index 32e3cd2a7d..2ca72ee4a8 100644 --- a/xdsl/irdl/declarative_assembly_format.py +++ b/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: """ diff --git a/xdsl/irdl/declarative_assembly_format_parser.py b/xdsl/irdl/declarative_assembly_format_parser.py index 3308f0dd4a..ceff33789c 100644 --- a/xdsl/irdl/declarative_assembly_format_parser.py +++ b/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, + )