diff --git a/tests/irdl/test_attr_declarative_assembly_format.py b/tests/irdl/test_attr_declarative_assembly_format.py new file mode 100644 index 0000000000..b5f5cc8c16 --- /dev/null +++ b/tests/irdl/test_attr_declarative_assembly_format.py @@ -0,0 +1,908 @@ +"""Tests for declarative assembly format for ParametrizedAttribute types.""" + +from __future__ import annotations + +from io import StringIO + +import pytest + +from xdsl.context import Context +from xdsl.dialects.builtin import ( + IntegerType, + NoneAttr, + i1, + i32, + i64, +) +from xdsl.ir import Attribute, ParametrizedAttribute, TypeAttribute +from xdsl.irdl import ParamAttrDef, irdl_attr_definition, param_def +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 + +# ============================================================================ +# Test attribute definitions (without assembly_format — for unit tests) +# ============================================================================ + + +@irdl_attr_definition +class SimpleType(ParametrizedAttribute, TypeAttribute): + name = "test_af.simple" + value: IntegerType = param_def() + + +@irdl_attr_definition +class TwoParamType(ParametrizedAttribute, TypeAttribute): + name = "test_af.two_param" + first: IntegerType = param_def() + second: IntegerType = param_def() + + +@irdl_attr_definition +class KeywordType(ParametrizedAttribute, TypeAttribute): + name = "test_af.keyword" + value: IntegerType = param_def() + + +# ============================================================================ +# Test attribute definitions (with assembly_format — for integration tests) +# ============================================================================ + + +@irdl_attr_definition +class SimpleIntegType(ParametrizedAttribute, TypeAttribute): + name = "test_af.simple_i" + value: IntegerType = param_def() + assembly_format = "$value" + + +@irdl_attr_definition +class TwoParamIntegType(ParametrizedAttribute, TypeAttribute): + name = "test_af.two_param_i" + first: IntegerType = param_def() + second: IntegerType = param_def() + assembly_format = "$first `,` $second" + + +@irdl_attr_definition +class KeywordIntegType(ParametrizedAttribute, TypeAttribute): + name = "test_af.keyword_i" + value: IntegerType = param_def() + assembly_format = "`stride` `=` $value" + + +@irdl_attr_definition +class InnerType(ParametrizedAttribute, TypeAttribute): + """An inner type with its own assembly_format, for composition testing.""" + + name = "test_af.inner" + elem: IntegerType = param_def() + assembly_format = "$elem" + + +@irdl_attr_definition +class OuterType(ParametrizedAttribute, TypeAttribute): + """Contains an InnerType parameter to test sub-attribute composition.""" + + name = "test_af.outer" + wrapped: InnerType = param_def() + tag: IntegerType = param_def() + assembly_format = "$wrapped `,` $tag" + + +@irdl_attr_definition +class WrapperType(ParametrizedAttribute, TypeAttribute): + """Wraps a generic Attribute — accepts any type including complex ones.""" + + name = "test_af.wrapper" + inner: Attribute = param_def() + assembly_format = "$inner" + + +@irdl_attr_definition +class PairType(ParametrizedAttribute, TypeAttribute): + """A pair of generic attributes.""" + + name = "test_af.pair" + left: Attribute = param_def() + right: Attribute = param_def() + assembly_format = "$left `,` $right" + + +@irdl_attr_definition +class PairAttr(ParametrizedAttribute): + """A pair attribute (not a type) — uses # prefix.""" + + name = "test_af.pair_attr" + left: Attribute = param_def() + right: Attribute = param_def() + assembly_format = "$left `,` $right" + + +@irdl_attr_definition +class ThreeParamType(ParametrizedAttribute, TypeAttribute): + """Three params for more complex tests.""" + + name = "test_af.three_param" + a: IntegerType = param_def() + b: IntegerType = param_def() + c: IntegerType = param_def() + 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) +# ============================================================================ + + +def _attr_def_for(cls: type[ParametrizedAttribute]) -> ParamAttrDef: + """Build a ParamAttrDef from a ParametrizedAttribute class.""" + return ParamAttrDef.from_pyrdl(cls) + + +def _print_with_format( + format_str: str, attr: ParametrizedAttribute, attr_def: ParamAttrDef +) -> str: + program = AttrFormatProgram.from_str(format_str, attr_def) + output = StringIO() + printer = Printer(stream=output) + program.print(printer, attr) + return output.getvalue() + + +def _parse_with_format( + format_str: str, + body: str, + attr_def: ParamAttrDef, +) -> list[Attribute]: + ctx = Context(allow_unregistered=True) + parser = Parser(ctx, body) + program = AttrFormatProgram.from_str(format_str, attr_def) + return program.parse(parser, attr_def) + + +# ============================================================================ +# Integration test helpers (parse/print via operation parser) +# ============================================================================ + + +def parse_type(type_str: str, ctx: Context) -> Attribute: + """Parse a type like '!test_af.simple_i' and return the type attribute.""" + program = f'"test.op"() : () -> {type_str}' + parser = Parser(ctx, program) + op = parser.parse_optional_operation() + assert op is not None, f"Failed to parse operation with type {type_str}" + return op.results[0].type + + +def parse_attr(attr_str: str, ctx: Context) -> Attribute: + """Parse an attribute like '#test_af.pair_attr'.""" + program = f'"test.op"() {{attr = {attr_str}}} : () -> ()' + parser = Parser(ctx, program) + op = parser.parse_optional_operation() + assert op is not None, f"Failed to parse operation with attr {attr_str}" + return op.attributes["attr"] + + +def print_attr(attr: Attribute) -> str: + """Print an attribute to a string.""" + output = StringIO() + printer = Printer(stream=output) + printer.print_attribute(attr) + return output.getvalue() + + +def check_roundtrip( + attr_type: type[ParametrizedAttribute], body: str, ctx: Context +) -> None: + """Parse !type, print it, verify the body matches.""" + type_str = f"!{attr_type.name}<{body}>" + parsed = parse_type(type_str, ctx) + printed = print_attr(parsed) + assert printed == type_str, f"Round-trip failed: {printed!r} != {type_str!r}" + + +def check_attr_roundtrip( + attr_type: type[ParametrizedAttribute], body: str, ctx: Context +) -> None: + """Round-trip for #-prefixed attributes.""" + attr_str = f"#{attr_type.name}<{body}>" + parsed = parse_attr(attr_str, ctx) + printed = print_attr(parsed) + assert printed == attr_str, f"Attr round-trip failed: {printed!r} != {attr_str!r}" + + +# ============================================================================ +# Fixture +# ============================================================================ + + +@pytest.fixture +def ctx() -> Context: + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(SimpleIntegType) + ctx.load_attr_or_type(TwoParamIntegType) + ctx.load_attr_or_type(KeywordIntegType) + ctx.load_attr_or_type(InnerType) + ctx.load_attr_or_type(OuterType) + ctx.load_attr_or_type(WrapperType) + 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 + + +# ============================================================================ +# Unit tests — AttrFormatProgram directly (print) +# ============================================================================ + + +def test_print_single_param(): + attr_def = _attr_def_for(SimpleType) + result = _print_with_format("$value", SimpleType(i32), attr_def) + assert result == "i32" + + +def test_print_two_params(): + attr_def = _attr_def_for(TwoParamType) + result = _print_with_format("$first `,` $second", TwoParamType(i32, i64), attr_def) + assert result == "i32, i64" + + +def test_print_keyword(): + attr_def = _attr_def_for(KeywordType) + result = _print_with_format("`stride` `=` $value", KeywordType(i32), attr_def) + assert result == "stride = i32" + + +def test_print_whitespace_suppress(): + attr_def = _attr_def_for(TwoParamType) + result = _print_with_format("$first `` $second", TwoParamType(i32, i64), attr_def) + assert result == "i32i64" + + +# ============================================================================ +# Unit tests — AttrFormatProgram directly (parse) +# ============================================================================ + + +def test_parse_single_param(): + attr_def = _attr_def_for(SimpleType) + params = _parse_with_format("$value", "i32", attr_def) + assert params == [i32] + + +def test_parse_two_params(): + attr_def = _attr_def_for(TwoParamType) + params = _parse_with_format("$first `,` $second", "i32, i64", attr_def) + assert params == [i32, i64] + + +def test_parse_keyword(): + attr_def = _attr_def_for(KeywordType) + params = _parse_with_format("`stride` `=` $value", "stride = i32", attr_def) + assert params == [i32] + + +# ============================================================================ +# Unit tests — round-trip (print then parse) +# ============================================================================ + + +@pytest.mark.parametrize( + "fmt, attr, attr_cls", + [ + ("$value", SimpleType(i32), SimpleType), + ("$value", SimpleType(i64), SimpleType), + ("$first `,` $second", TwoParamType(i32, i64), TwoParamType), + ("`stride` `=` $value", KeywordType(i32), KeywordType), + ], +) +def test_roundtrip( + fmt: str, attr: ParametrizedAttribute, attr_cls: type[ParametrizedAttribute] +): + attr_def = _attr_def_for(attr_cls) + printed = _print_with_format(fmt, attr, attr_def) + parsed = _parse_with_format(fmt, printed, attr_def) + reconstructed = attr_cls(*parsed) + assert reconstructed == attr + + +# ============================================================================ +# Unit tests — format string validation errors +# ============================================================================ + + +def test_error_missing_parameter(): + attr_def = _attr_def_for(TwoParamType) + with pytest.raises(ParseError, match="parameter 'second' not found"): + AttrFormatProgram.from_str("$first", attr_def) + + +def test_error_duplicate_parameter(): + attr_def = _attr_def_for(SimpleType) + with pytest.raises(ParseError, match="is already bound"): + AttrFormatProgram.from_str("$value `,` $value", attr_def) + + +def test_error_unknown_variable(): + attr_def = _attr_def_for(SimpleType) + with pytest.raises(ParseError, match="does not refer to a parameter"): + AttrFormatProgram.from_str("$nonexistent", attr_def) + + +# ============================================================================ +# Integration tests — basic round-trips via operation parser +# ============================================================================ + + +@pytest.mark.parametrize("body", ["i32", "i64"]) +def test_simple_param_roundtrip(body: str, ctx: Context): + check_roundtrip(SimpleIntegType, body, ctx) + + +@pytest.mark.parametrize("body", ["i32, i64", "i1, i32"]) +def test_two_param_roundtrip(body: str, ctx: Context): + check_roundtrip(TwoParamIntegType, body, ctx) + + +@pytest.mark.parametrize("body", ["stride = i32", "stride = i64"]) +def test_keyword_roundtrip(body: str, ctx: Context): + check_roundtrip(KeywordIntegType, body, ctx) + + +def test_three_param_roundtrip(ctx: Context): + check_roundtrip(ThreeParamType, "i1, i32, i64", ctx) + + +# ============================================================================ +# Integration tests — sub-attribute composition +# ============================================================================ + + +@pytest.mark.parametrize( + "body", + [ + "!test_af.inner, i64", + "!test_af.inner, i32", + ], +) +def test_outer_type_roundtrip(body: str, ctx: Context): + check_roundtrip(OuterType, body, ctx) + + +def test_outer_type_constructs_correctly(ctx: Context): + """Verify the outer type correctly nests the inner type.""" + parsed = parse_type("!test_af.outer, i64>", ctx) + assert isinstance(parsed, OuterType) + assert isinstance(parsed.wrapped, InnerType) + assert parsed.wrapped.elem == i32 + assert parsed.tag == i64 + + +# ============================================================================ +# Integration tests — complex nesting +# ============================================================================ + + +@pytest.mark.parametrize( + "body", + [ + "i32", + "f32", + "vector<4xi32>", + "tensor<2x3xf32>", + ], +) +def test_wrapper_complex_types(body: str, ctx: Context): + check_roundtrip(WrapperType, body, ctx) + + +@pytest.mark.parametrize( + "body", + [ + "i32, i64", + "tensor<2x3xf32>, memref<16xi32>", + ], +) +def test_pair_complex_types(body: str, ctx: Context): + check_roundtrip(PairType, body, ctx) + + +def test_deeply_nested_types(ctx: Context): + """Wrapper wrapping a pair wrapping inner types.""" + check_roundtrip( + WrapperType, + "!test_af.pair, !test_af.inner>", + ctx, + ) + + +# ============================================================================ +# Integration tests — attribute parameters (# prefix) +# ============================================================================ + + +def test_pair_attr_roundtrip(ctx: Context): + check_attr_roundtrip(PairAttr, "i32, i64", ctx) + + +def test_pair_attr_constructs_correctly(ctx: Context): + parsed = parse_attr("#test_af.pair_attr", ctx) + assert isinstance(parsed, PairAttr) + assert parsed.left == i32 + assert parsed.right == i64 + + +def test_wrapper_with_attr_parameter(ctx: Context): + """A type parameter that is a #-prefixed attribute.""" + check_roundtrip(WrapperType, "#test_af.pair_attr", ctx) + + +# ============================================================================ +# Integration tests — whitespace directives +# ============================================================================ + + +def test_whitespace_suppress(): + """`` `` (empty backtick) suppresses whitespace.""" + + @irdl_attr_definition + class SuppressType(ParametrizedAttribute, TypeAttribute): + name = "test_af.suppress" + a: IntegerType = param_def() + b: IntegerType = param_def() + assembly_format = "$a `` $b" + + attr = SuppressType(i32, i64) + printed = print_attr(attr) + assert printed == "!test_af.suppress" + + +def test_whitespace_newline(): + r"""`\n` inserts newline.""" + + @irdl_attr_definition + class NewlineType(ParametrizedAttribute, TypeAttribute): + name = "test_af.newline" + a: IntegerType = param_def() + b: IntegerType = param_def() + assembly_format = "$a `` `\\n` `` $b" + + attr = NewlineType(i32, i64) + printed = print_attr(attr) + 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 — qualified directive +# ============================================================================ + + +def test_qualified_roundtrip(): + """qualified($x) should parse and print identically to $x in xdsl.""" + + @irdl_attr_definition + class QualifiedType(ParametrizedAttribute, TypeAttribute): + name = "test_af.qualified" + value: IntegerType = param_def() + assembly_format = "qualified($value)" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(QualifiedType) + check_roundtrip(QualifiedType, "i32", ctx) + + +# ============================================================================ +# Integration tests — params directive +# ============================================================================ + + +def test_params_directive_roundtrip(): + """params captures all parameters, printed comma-separated.""" + + @irdl_attr_definition + class ParamsType(ParametrizedAttribute, TypeAttribute): + name = "test_af.params_all" + a: IntegerType = param_def() + b: IntegerType = param_def() + assembly_format = "params" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(ParamsType) + check_roundtrip(ParamsType, "i32, i64", ctx) + + +def test_params_directive_three_params(): + """params with three parameters.""" + + @irdl_attr_definition + class Params3Type(ParametrizedAttribute, TypeAttribute): + name = "test_af.params3" + a: IntegerType = param_def() + b: IntegerType = param_def() + c: IntegerType = param_def() + assembly_format = "params" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(Params3Type) + check_roundtrip(Params3Type, "i1, i32, i64", ctx) + + +def test_params_with_optional(): + """params with optional parameter omitted.""" + + @irdl_attr_definition + class ParamsOptType(ParametrizedAttribute, TypeAttribute): + name = "test_af.params_opt" + req: IntegerType = param_def() + opt: IntegerType | NoneAttr = param_def() + assembly_format = "params" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(ParamsOptType) + check_roundtrip(ParamsOptType, "i32, i64", ctx) + check_roundtrip(ParamsOptType, "i32", ctx) + + +# ============================================================================ +# Integration tests — struct directive +# ============================================================================ + + +def test_struct_all_params(): + """struct(params) prints all params as key=value pairs.""" + + @irdl_attr_definition + class StructAllType(ParametrizedAttribute, TypeAttribute): + name = "test_af.struct_all" + x: IntegerType = param_def() + y: IntegerType = param_def() + assembly_format = "struct(params)" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(StructAllType) + check_roundtrip(StructAllType, "x = i32, y = i64", ctx) + + +def test_struct_subset(): + """struct($a, $b) with a subset of parameters.""" + + @irdl_attr_definition + class StructSubType(ParametrizedAttribute, TypeAttribute): + name = "test_af.struct_sub" + a: IntegerType = param_def() + b: IntegerType = param_def() + c: IntegerType = param_def() + assembly_format = "struct($a, $b) `,` $c" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(StructSubType) + check_roundtrip(StructSubType, "a = i32, b = i64, i1", ctx) + + +def test_struct_reordered_parse(): + """Struct fields can be parsed in any order.""" + + @irdl_attr_definition + class StructReorderType(ParametrizedAttribute, TypeAttribute): + name = "test_af.struct_reorder" + a: IntegerType = param_def() + b: IntegerType = param_def() + assembly_format = "struct(params)" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(StructReorderType) + parsed = parse_type("!test_af.struct_reorder", ctx) + assert isinstance(parsed, StructReorderType) + assert parsed.a == i32 + assert parsed.b == i64 + + +def test_struct_optional_param(): + """Struct with optional parameter omitted.""" + + @irdl_attr_definition + class StructOptType(ParametrizedAttribute, TypeAttribute): + name = "test_af.struct_opt" + req: IntegerType = param_def() + opt: IntegerType | NoneAttr = param_def() + assembly_format = "struct(params)" + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(StructOptType) + check_roundtrip(StructOptType, "req = i32, opt = i64", ctx) + 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($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" + parsed = parse_type(printed, ctx) + assert parsed == attr + + +# ============================================================================ +# Integration tests — ref directive +# ============================================================================ + + +def test_ref_directive_in_custom(): + """ref($x) allows re-using a parameter without re-binding it.""" + + @irdl_attr_custom_directive + class PrintTwiceDirective(AttrCustomDirective): + """Prints a parameter, then prints a ref of it again.""" + + value: ParameterVariable + value_ref: ParameterVariable + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + self.value.parse(parser, state) + parser.parse_punctuation(",") + parser.parse_attribute() + return True + + def print( + self, + printer: Printer, + state: PrintingState, + attr: ParametrizedAttribute, + /, + ) -> None: + self.value.print(printer, state, attr) + printer.print_string(", ") + state.should_emit_space = False + state.last_was_punctuation = True + self.value_ref.print(printer, state, attr) + + @irdl_attr_definition + class RefType(ParametrizedAttribute, TypeAttribute): + name = "test_af.ref" + a: IntegerType = param_def() + assembly_format = "custom($a, ref($a))" + custom_directives = (PrintTwiceDirective,) + + attr = RefType(i32) + printed = print_attr(attr) + assert printed == "!test_af.ref" + + +# ============================================================================ +# Integration tests — printing correctness +# ============================================================================ + + +@pytest.mark.parametrize( + "attr, expected", + [ + (SimpleIntegType(i32), "!test_af.simple_i"), + (TwoParamIntegType(i32, i64), "!test_af.two_param_i"), + (KeywordIntegType(i32), "!test_af.keyword_i"), + (ThreeParamType(i1, i32, i64), "!test_af.three_param"), + (PairAttr(i32, i64), "#test_af.pair_attr"), + ], +) +def test_print_correctness(attr: Attribute, expected: str): + assert print_attr(attr) == expected + + +# ============================================================================ +# Integration tests — definition-time errors +# ============================================================================ + + +def test_defn_error_missing_parameter(): + """Format string missing a parameter should fail at definition time.""" + with pytest.raises(PyRDLAttrDefinitionError, match="parameter 'second' not found"): + + @irdl_attr_definition + class BadType( # pyright: ignore[reportUnusedClass] + ParametrizedAttribute, TypeAttribute + ): + name = "test_af.bad_missing" + first: IntegerType = param_def() + second: IntegerType = param_def() + assembly_format = "$first" + + +def test_defn_error_duplicate_parameter(): + """Duplicate parameter in format string should fail at definition time.""" + with pytest.raises(PyRDLAttrDefinitionError, match="is already bound"): + + @irdl_attr_definition + class BadType( # pyright: ignore[reportUnusedClass] + ParametrizedAttribute, TypeAttribute + ): + name = "test_af.bad_dup" + value: IntegerType = param_def() + assembly_format = "$value `,` $value" + + +def test_defn_error_unknown_variable(): + """Unknown variable in format string should fail at definition time.""" + with pytest.raises(PyRDLAttrDefinitionError, match="does not refer to a parameter"): + + @irdl_attr_definition + class BadType( # pyright: ignore[reportUnusedClass] + ParametrizedAttribute, TypeAttribute + ): + name = "test_af.bad_unknown" + value: IntegerType = param_def() + assembly_format = "$nonexistent" + + +def test_defn_error_assembly_format_with_parse_parameters(): + """assembly_format with custom parse_parameters should fail.""" + with pytest.raises( + PyRDLAttrDefinitionError, + match="Cannot use assembly_format with custom parse_parameters", + ): + + @irdl_attr_definition + class BadType( # pyright: ignore[reportUnusedClass] + ParametrizedAttribute, TypeAttribute + ): + name = "test_af.bad_custom_parse" + value: IntegerType = param_def() + assembly_format = "$value" + + @classmethod + def parse_parameters(cls, parser: AttrParser) -> list[Attribute]: + return [] + + +def test_defn_error_assembly_format_with_print_parameters(): + """assembly_format with custom print_parameters should fail.""" + with pytest.raises( + PyRDLAttrDefinitionError, + match="Cannot use assembly_format with custom print_parameters", + ): + + @irdl_attr_definition + class BadType( # pyright: ignore[reportUnusedClass] + ParametrizedAttribute, TypeAttribute + ): + name = "test_af.bad_custom_print" + value: IntegerType = param_def() + assembly_format = "$value" + + def print_parameters(self, printer: Printer) -> None: + pass + + +def test_defn_error_assembly_format_on_singleton(): + """assembly_format on attribute with no parameters should fail.""" + with pytest.raises( + PyRDLAttrDefinitionError, + match="Cannot use assembly_format on attribute with no parameters", + ): + + @irdl_attr_definition + class BadType( # pyright: ignore[reportUnusedClass] + ParametrizedAttribute, TypeAttribute + ): + name = "test_af.bad_singleton" + assembly_format = "" diff --git a/xdsl/irdl/attributes.py b/xdsl/irdl/attributes.py index 8606c1b159..b3f13d9bff 100644 --- a/xdsl/irdl/attributes.py +++ b/xdsl/irdl/attributes.py @@ -8,10 +8,11 @@ import sys from abc import ABC, abstractmethod from collections.abc import Callable, Hashable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from inspect import get_annotations, isclass from types import FunctionType, GenericAlias, UnionType from typing import ( + TYPE_CHECKING, Annotated, Any, Generic, @@ -30,11 +31,12 @@ if sys.version_info >= (3, 14, 0): from typing_extensions import TypeForm else: - from typing import TYPE_CHECKING - if TYPE_CHECKING: from typing_extensions import TypeForm +if TYPE_CHECKING: + from xdsl.parser import AttrParser + from xdsl.printer import Printer from xdsl.ir import ( Attribute, @@ -46,7 +48,7 @@ TypedAttribute, ) from xdsl.utils.classvar import is_const_classvar -from xdsl.utils.exceptions import PyRDLAttrDefinitionError, PyRDLTypeError +from xdsl.utils.exceptions import ParseError, PyRDLAttrDefinitionError, PyRDLTypeError from xdsl.utils.hints import ( PropertyType, get_type_var_from_generic_class, @@ -186,6 +188,7 @@ class ParamAttrDef: name: str parameters: list[tuple[str, ParamDef]] + custom_directives: dict[str, type[Any]] = field(default_factory=lambda: {}) @staticmethod def from_pyrdl( @@ -225,6 +228,9 @@ def from_pyrdl( if ( # Ignore name field field_name != "name" + # Ignore assembly_format and custom_directives (handled separately) + and field_name != "assembly_format" + and field_name != "custom_directives" # Ignore functions and not isinstance( field_value, @@ -299,14 +305,17 @@ def from_pyrdl( f"Missing field type for parameter name {field_name}" ) - return ParamAttrDef(name, list(parameters.items())) + custom_dirs = getattr(pyrdl_def, "custom_directives", ()) + custom_directives = {d.__name__: d for d in custom_dirs} + + return ParamAttrDef(name, list(parameters.items()), custom_directives) def verify(self, attr: ParametrizedAttribute): """Verify that `attr` satisfies the invariants.""" constraint_context = ConstraintContext() - for field, param_def in self.parameters: - param_def.constr.verify(getattr(attr, field), constraint_context) + for param_name, param_def in self.parameters: + param_def.constr.verify(getattr(attr, param_name), constraint_context) _PAttrTT = TypeVar("_PAttrTT", bound=type[ParametrizedAttribute]) @@ -326,6 +335,44 @@ def irdl_param_attr_definition(cls: _PAttrTT) -> _PAttrTT: attr_def = ParamAttrDef.from_pyrdl(cls) new_fields = get_accessors_from_param_attr_def(attr_def) + assembly_format = getattr(cls, "assembly_format", None) + if assembly_format is not None: + if not attr_def.parameters: + raise PyRDLAttrDefinitionError( + "Cannot use assembly_format on attribute with no parameters" + ) + if "parse_parameters" in cls.__dict__: + raise PyRDLAttrDefinitionError( + "Cannot use assembly_format with custom parse_parameters" + ) + if "print_parameters" in cls.__dict__: + raise PyRDLAttrDefinitionError( + "Cannot use assembly_format with custom print_parameters" + ) + + from xdsl.irdl.declarative_assembly_format import AttrFormatProgram + + try: + program = AttrFormatProgram.from_str(assembly_format, attr_def) + except ParseError as e: + raise PyRDLAttrDefinitionError( + "Error during the parsing of the assembly format: ", e.args + ) from e + + @classmethod + def parse_with_format( + cls: type[ParametrizedAttribute], parser: "AttrParser" + ) -> list[Attribute]: + with parser.in_angle_brackets(): + return program.parse(parser, attr_def) + + def print_with_format(self: ParametrizedAttribute, printer: "Printer") -> None: + with printer.in_angle_brackets(): + program.print(printer, self) + + new_fields["parse_parameters"] = parse_with_format + new_fields["print_parameters"] = print_with_format + if issubclass(cls, TypedAttribute): type_indexes = tuple( i for i, (p, _) in enumerate(attr_def.parameters) if p == "type" diff --git a/xdsl/irdl/declarative_assembly_format.py b/xdsl/irdl/declarative_assembly_format.py index 72bd0781b3..125e40c749 100644 --- a/xdsl/irdl/declarative_assembly_format.py +++ b/xdsl/irdl/declarative_assembly_format.py @@ -19,6 +19,7 @@ BytesAttr, DenseArrayBase, IntegerType, + NoneAttr, StringAttr, UnitAttr, ) @@ -37,6 +38,7 @@ IRDLOperationInvT, OpDef, OptionalDef, + ParamAttrDef, Successor, VarIRConstruct, VarOperand, @@ -44,7 +46,7 @@ is_const_classvar, verify_variadic_same_size, ) -from xdsl.parser import Parser, UnresolvedOperand +from xdsl.parser import AttrParser, Parser, UnresolvedOperand from xdsl.printer import Printer from xdsl.utils.exceptions import PyRDLError, VerifyException from xdsl.utils.hints import isa @@ -1358,6 +1360,47 @@ def is_optional_like(self) -> bool: return False +# =========================================================================== +# Shared print helpers for structural directives (whitespace, punctuation, +# keyword). These are used by both the op-side and attr-side directive classes. +# =========================================================================== + + +def _print_whitespace(printer: Printer, state: PrintingState, whitespace: str) -> None: + printer.print_string(whitespace) + state.last_was_punctuation = whitespace == "" + state.should_emit_space = False + + +def _print_punctuation( + printer: Printer, state: PrintingState, punctuation: PunctuationSpelling +) -> None: + emit_space = False + if state.should_emit_space: + if state.last_was_punctuation: + if punctuation not in (">", ")", "}", "]", ","): + emit_space = True + elif punctuation not in ("<", ">", "(", ")", "{", "}", "[", "]", ","): + emit_space = True + + if emit_space: + printer.print_string(" ") + + printer.print_string(punctuation) + + state.should_emit_space = punctuation not in ("<", "(", "{", "[") + state.last_was_punctuation = True + + +def _print_keyword(printer: Printer, state: PrintingState, keyword: str) -> None: + if state.should_emit_space: + printer.print_string(" ") + state.should_emit_space = True + state.last_was_punctuation = False + + printer.print_string(keyword) + + @dataclass(frozen=True) class WhitespaceDirective(FormatDirective): """ @@ -1375,9 +1418,7 @@ def parse(self, parser: Parser, state: ParsingState) -> bool: return False def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None: - printer.print_string(self.whitespace) - state.last_was_punctuation = self.whitespace == "" - state.should_emit_space = False + _print_whitespace(printer, state, self.whitespace) @dataclass(frozen=True) @@ -1399,21 +1440,7 @@ def parse(self, parser: Parser, state: ParsingState) -> bool: return parser.parse_optional_punctuation(self.punctuation) is not None def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None: - emit_space = False - if state.should_emit_space: - if state.last_was_punctuation: - if self.punctuation not in (">", ")", "}", "]", ","): - emit_space = True - elif self.punctuation not in ("<", ">", "(", ")", "{", "}", "[", "]", ","): - emit_space = True - - if emit_space: - printer.print_string(" ") - - printer.print_string(self.punctuation) - - state.should_emit_space = self.punctuation not in ("<", "(", "{", "[") - state.last_was_punctuation = True + _print_punctuation(printer, state, self.punctuation) def is_optional_like(self) -> bool: return True @@ -1435,12 +1462,7 @@ def parse(self, parser: Parser, state: ParsingState) -> bool: return parser.parse_optional_keyword(self.keyword) is not None def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None: - if state.should_emit_space: - printer.print_string(" ") - state.should_emit_space = True - state.last_was_punctuation = False - - printer.print_string(self.keyword) + _print_keyword(printer, state, self.keyword) def is_optional_like(self) -> bool: return True @@ -1486,3 +1508,371 @@ def set_empty(self, state: ParsingState) -> None: element.set_empty(state) for element in self.else_elements: element.set_empty(state) + + +# =========================================================================== +# Attribute/Type declarative assembly format +# =========================================================================== + + +@dataclass +class AttrParsingState: + """State during parsing of a ParametrizedAttribute using declarative format.""" + + parameters: list[Attribute | None] + attr_def: ParamAttrDef + + def __init__(self, attr_def: ParamAttrDef): + self.attr_def = attr_def + self.parameters = [None] * len(attr_def.parameters) + + +@dataclass(frozen=True) +class AttrFormatDirective(ABC): + """Base class for attribute/type format directives. + + Separate from FormatDirective because operation formats parse full ops + (operands, results, regions, etc.) using Parser and ParsingState, while + attribute/type formats only fill parameter slots of a ParametrizedAttribute + using AttrParser and AttrParsingState. + """ + + @abstractmethod + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: ... + + def parse_optional(self, parser: AttrParser, state: AttrParsingState) -> bool: + return self.parse(parser, state) + + @abstractmethod + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: ... + + def set_empty(self, state: AttrParsingState) -> None: + """ + Set the appropriate field of the parsing state to be empty. + Used when a variable appears in an optional group which is not parsed. + """ + return + + def is_present(self, attr: ParametrizedAttribute, /) -> bool: + return True + + def is_anchorable(self) -> bool: + return False + + def is_optional_like(self) -> bool: + return False + + +# =========================================================================== +# Attr-side structural directives +# =========================================================================== + + +@dataclass(frozen=True) +class AttrWhitespaceDirective(AttrFormatDirective): + """Whitespace directive for attribute/type assembly format.""" + + whitespace: Literal[" ", "\n", ""] + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + return False + + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: + _print_whitespace(printer, state, self.whitespace) + + +@dataclass(frozen=True) +class AttrPunctuationDirective(AttrFormatDirective): + """Punctuation directive for attribute/type assembly format.""" + + punctuation: PunctuationSpelling + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + return parser.parse_optional_punctuation(self.punctuation) is not None + + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: + _print_punctuation(printer, state, self.punctuation) + + def is_optional_like(self) -> bool: + return True + + +@dataclass(frozen=True) +class AttrKeywordDirective(AttrFormatDirective): + """Keyword directive for attribute/type assembly format.""" + + keyword: str + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + return parser.parse_optional_keyword(self.keyword) is not None + + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: + _print_keyword(printer, state, self.keyword) + + def is_optional_like(self) -> bool: + return True + + +@dataclass(frozen=True) +class ParameterVariable(AttrFormatDirective): + """A parameter variable directive for attribute assembly format.""" + + name: str + index: int + is_optional: bool + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + if self.is_optional: + attr = parser.parse_optional_attribute() + else: + attr = parser.parse_attribute() + if attr is None: + return False + state.parameters[self.index] = attr + return True + + def parse_optional(self, parser: AttrParser, state: AttrParsingState) -> bool: + attr = parser.parse_optional_attribute() + if attr is None: + return False + state.parameters[self.index] = attr + return True + + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: + param = attr.parameters[self.index] + if isinstance(param, NoneAttr): + return + state.print_whitespace(printer) + printer.print_attribute(param) + + def is_present(self, attr: ParametrizedAttribute, /) -> bool: + return not isinstance(attr.parameters[self.index], NoneAttr) + + def is_anchorable(self) -> bool: + return self.is_optional + + def is_optional_like(self) -> bool: + return self.is_optional + + def set_empty(self, state: AttrParsingState) -> None: + state.parameters[self.index] = NoneAttr() + + +@dataclass(frozen=True) +class QualifiedParameterVariable(ParameterVariable): + """Wraps a ParameterVariable to print with full dialect qualification. + + In xdsl, print_attribute() always qualifies, so this is currently + identical to ParameterVariable. Exists for MLIR format compatibility. + """ + + +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.""" + + params: tuple[ParameterVariable, ...] + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + non_optional = [p for p in self.params if not p.is_optional] + optional = [p for p in self.params if p.is_optional] + for i, pv in enumerate(non_optional): + if i: + parser.parse_punctuation(",") + pv.parse(parser, state) + for pv in optional: + if not parser.parse_optional_punctuation(","): + pv.set_empty(state) + continue + pv.parse(parser, state) + return True + + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: + first = True + for pv in self.params: + if not pv.is_present(attr): + continue + if not first: + printer.print_string(", ") + state.should_emit_space = False + state.last_was_punctuation = True + first = False + pv.print(printer, state, attr) + + +@dataclass(frozen=True) +class StructDirective(AttrFormatDirective): + """Prints parameters as key = value pairs, parseable in any order.""" + + params: tuple[ParameterVariable, ...] + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + remaining = {pv.name: pv for pv in self.params} + first = True + while remaining: + if not first: + if not parser.parse_optional_punctuation(","): + break + name = parser.parse_optional_identifier() + if name is None: + if first: + break + parser.raise_error("expected parameter name in struct") + first = False + if name not in remaining: + parser.raise_error( + f"unexpected parameter '{name}', " + f"expected one of {list(remaining.keys())}" + ) + parser.parse_punctuation("=") + remaining[name].parse(parser, state) + del remaining[name] + for pv in remaining.values(): + if pv.is_optional: + pv.set_empty(state) + else: + parser.raise_error(f"missing required parameter '{pv.name}' in struct") + return True + + def print( + self, printer: Printer, state: PrintingState, attr: ParametrizedAttribute, / + ) -> None: + first = True + for pv in self.params: + if not pv.is_present(attr): + continue + if not first: + printer.print_string(", ") + first = False + printer.print_string(f"{pv.name} = ") + state.should_emit_space = False + state.last_was_punctuation = True + pv.print(printer, state, attr) + + +@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: + """ + The toplevel data structure of a declarative assembly format program + for attributes/types. + """ + + stmts: tuple[AttrFormatDirective, ...] + """The statements composing the program. They are executed in order.""" + + @staticmethod + def from_str(format_str: str, attr_def: ParamAttrDef) -> AttrFormatProgram: + from xdsl.irdl.declarative_assembly_format_parser import AttrFormatParser + + return AttrFormatParser(format_str, attr_def).parse_format() + + def parse(self, parser: AttrParser, attr_def: ParamAttrDef) -> list[Attribute]: + """Parse parameter values and check that every slot was assigned. + + IRDL constraint verification runs when the attribute is constructed, + not during this parse step. + """ + state = AttrParsingState(attr_def) + for stmt in self.stmts: + stmt.parse(parser, state) + for i, (name, _) in enumerate(attr_def.parameters): + if state.parameters[i] is None: + parser.raise_error(f"parameter '{name}' was not parsed") + return cast(list[Attribute], state.parameters) + + def print(self, printer: Printer, attr: ParametrizedAttribute) -> None: + state = PrintingState(last_was_punctuation=True, should_emit_space=False) + for stmt in self.stmts: + stmt.print(printer, state, attr) diff --git a/xdsl/irdl/declarative_assembly_format_parser.py b/xdsl/irdl/declarative_assembly_format_parser.py index b67b1a26b1..03fee4b5ab 100644 --- a/xdsl/irdl/declarative_assembly_format_parser.py +++ b/xdsl/irdl/declarative_assembly_format_parser.py @@ -16,6 +16,7 @@ Builtin, DenseArrayBase, IntegerType, + NoneAttr, SymbolNameConstraint, UnitAttr, ) @@ -32,6 +33,7 @@ OptSingleBlockRegionDef, OptSuccessorDef, ParamAttrConstraint, + ParamAttrDef, ParsePropInAttrDict, SameVariadicOperandSize, SameVariadicResultSize, @@ -45,7 +47,13 @@ ) from xdsl.irdl.declarative_assembly_format import ( AttrDictDirective, + AttrFormatDirective, + AttrFormatProgram, AttributeVariable, + AttrKeywordDirective, + AttrOptionalGroupDirective, + AttrPunctuationDirective, + AttrWhitespaceDirective, DenseArrayAttributeVariable, Directive, FormatDirective, @@ -61,11 +69,15 @@ OptionalResultVariable, OptionalSuccessorVariable, OptionalUnitAttrVariable, + ParameterVariable, + ParamsDirective, PunctuationDirective, + QualifiedParameterVariable, RegionDirective, RegionVariable, ResultsDirective, ResultVariable, + StructDirective, SuccessorDirective, SuccessorVariable, SymbolNameAttributeVariable, @@ -900,3 +912,217 @@ def create_results_directive(self, inside_ref: bool) -> ResultsDirective: if not inside_ref: self.seen_result_types = [True] * len(self.seen_result_types) return ResultsDirective() + + +@dataclass(init=False) +class AttrFormatParser(BaseParser): + """Parser for attribute/type declarative assembly format strings.""" + + attr_def: ParamAttrDef + seen_parameters: set[str] + + def __init__(self, input_str: str, attr_def: ParamAttrDef): + super().__init__(ParserState(FormatLexer(Input(input_str, "")))) + self.attr_def = attr_def + self.seen_parameters = set() + + def parse_format(self) -> AttrFormatProgram: + elements: list[AttrFormatDirective] = [] + while self._current_token.kind != MLIRTokenKind.EOF: + elements.append(self.parse_format_directive()) + self.verify_parameters() + return AttrFormatProgram(tuple(elements)) + + def verify_parameters(self) -> None: + for name, _ in self.attr_def.parameters: + if name not in self.seen_parameters: + self.raise_error( + f"parameter '{name}' not found in format string, " + f"consider adding a '${name}' directive" + ) + + 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() + if self.parse_optional_keyword("qualified"): + return self.parse_qualified_directive() + if self.parse_optional_keyword("params"): + 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: + self._consume_token(MLIRTokenKind.BARE_IDENT) + token = self._current_token + name = self.parse_identifier(" after '$'") + + for idx, (param_name, param_def) in enumerate(self.attr_def.parameters): + if name == param_name: + if not inside_ref: + if name in self.seen_parameters: + self.raise_error( + f"parameter '{name}' is already bound", token.span + ) + self.seen_parameters.add(name) + is_optional = param_def.constr.verifies(NoneAttr()) + return ParameterVariable(name, idx, is_optional) + + self.raise_error( + f"'{name}' does not refer to a parameter of this attribute", token.span + ) + + def parse_keyword_or_punctuation(self) -> AttrFormatDirective: + start_token = self._current_token + self.parse_characters("`") + + # \n case + if self.parse_optional_keyword("\\"): + self.parse_keyword("n") + self.parse_characters("`") + return AttrWhitespaceDirective("\n") + + # Space or empty backtick case + end_token = self._current_token + if self.parse_optional_characters("`"): + whitespace = self.lexer.input.content[ + start_token.span.end : end_token.span.start + ] + if whitespace != " " and whitespace != "": + self.raise_error( + "unexpected whitespace in directive, " + "only ` ` or `` whitespace is allowed" + ) + return AttrWhitespaceDirective(whitespace) + + # Punctuation case + if self._current_token.kind.is_punctuation(): + punctuation = self._consume_token().text + self.parse_characters("`") + assert MLIRTokenKind.is_spelling_of_punctuation(punctuation) + return AttrPunctuationDirective(punctuation) + + # Identifier case + ident = self.parse_optional_identifier() + if ident is None or ident == "`": + self.raise_error("punctuation or identifier expected") + + 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, + ) + + def parse_qualified_directive(self) -> QualifiedParameterVariable: + with self.in_parens(): + pv = self.parse_variable() + return QualifiedParameterVariable(pv.name, pv.index, pv.is_optional) + + def _all_param_variables(self) -> tuple[ParameterVariable, ...]: + """Create ParameterVariable instances for all parameters.""" + params: list[ParameterVariable] = [] + for idx, (name, param_def) in enumerate(self.attr_def.parameters): + if name in self.seen_parameters: + self.raise_error(f"parameter '{name}' is already bound") + self.seen_parameters.add(name) + is_optional = param_def.constr.verifies(NoneAttr()) + params.append(ParameterVariable(name, idx, is_optional)) + return tuple(params) + + def parse_params_directive(self) -> ParamsDirective: + return ParamsDirective(self._all_param_variables()) + + def _parse_struct_variables(self) -> tuple[ParameterVariable, ...]: + """Parse the argument list for struct(): either `params` or `$a, $b, ...`.""" + with self.in_parens(): + if self.parse_optional_keyword("params"): + return self._all_param_variables() + params: list[ParameterVariable] = [] + params.append(self.parse_variable()) + while self.parse_optional_punctuation(","): + params.append(self.parse_variable()) + return tuple(params) + + 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_possible_ref_directive() + 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) + + def parse_possible_ref_directive(self) -> AttrFormatDirective: + """Parse a ref directive or variable: `ref` `(` `$var` `)` | `$var`.""" + if self.parse_optional_keyword("ref"): + with self.in_parens(): + return self.parse_variable(inside_ref=True) + return self.parse_variable()