Declarative assembly format

@dataclass
class ParsingState:
    """
    State carried during the parsing of an operation using the declarative assembly
    format.
    It contains the elements that have already been parsed.
    """

    operands: list[Sequence[UnresolvedOperand] | None]
    operand_types: list[Sequence[Attribute] | None]
    result_types: list[Sequence[Attribute] | None]
    regions: list[Sequence[Region] | None]
    successors: list[Sequence[Successor] | None]
    attributes: dict[str, Attribute]
    properties: dict[str, Attribute]
    op_type: type[IRDLOperation]
    context: VerificationContext

    def __init__(self, op_type: type[IRDLOperation]):
        op_def = op_type.get_irdl_definition()
        self.op_type = op_type
        self.operands = [None] * len(op_def.operands)
        self.operand_types = [None] * len(op_def.operands)
        self.result_types = [None] * len(op_def.results)
        self.regions = [None] * len(op_def.regions)
        self.successors = [None] * len(op_def.successors)
        self.attributes = {}
        self.properties = {}
        self.context = VerificationContext()

@dataclass
class PrintingState:
    """
    State carried during the printing of an operation using the declarative assembly
    format.
    It contains information on the last token, to know if a space should be emitted.
    """

    last_was_punctuation: bool = field(default=False)
    """Was the last element parsed a punctuation."""
    should_emit_space: bool = field(default=True)
    """
    Should the printer emit a space before the next element.
    Depending on the directive, the space might not be printed
    (for instance for some punctuations).
    """

    def print_whitespace(self, printer: Printer):
        """
        Handles whitespace printing for the majority of format directives.
        """
        if self.should_emit_space or not self.last_was_punctuation:
            printer.print_string(" ")
        self.should_emit_space = True
        self.last_was_punctuation = False

@dataclass(frozen=True)
class FormatProgram:
    """
    The toplevel data structure of a declarative assembly format program.
    It is used to parse and print an operation.
    """

    stmts: tuple[FormatDirective, ...]
    """The statements composing the program. They are executed in order."""

    @staticmethod
    def from_str(input: str, op_def: OpDef) -> FormatProgram:
        """
        Create the assembly format data program from its string representation.
        This might raise a ParseError exception if the string is invalid.
        """
        from xdsl.irdl.declarative_assembly_format_parser import FormatParser

        return FormatParser(input, op_def).parse_format()

    def parse(
        self, parser: Parser, op_type: type[IRDLOperationInvT]
    ) -> IRDLOperationInvT:
        """
        Parse the operation with this format.
        The given operation type is expected to be the operation type represented by
        the operation definition passed to the FormatParser that created this
        FormatProgram.
        """
        # Parse elements one by one
        op_def = op_type.get_irdl_definition()
        state = ParsingState(op_type)
        for stmt in self.stmts:
            stmt.parse(parser, state)

        # Get constraint variables from the parsed operand and result types
        self.resolve_constraint_variables(state, op_def)

        # Infer operand types that should be inferred
        assert None not in state.operands
        unresolved_operands = cast(list[Sequence[UnresolvedOperand]], state.operands)
        operand_types = self.resolve_operand_types(state, op_def)

        # Infer result types that should be inferred
        result_types = self.resolve_result_types(state, op_def)

        # Resolve all operands
        operands = tuple(
            parser.resolve_operands(uo, ot, parser.pos)
            for uo, ot in zip(unresolved_operands, operand_types, strict=True)
        )

        return op_type.build(
            result_types=result_types,
            operands=operands,
            attributes=state.attributes,
            properties=state.properties,
            regions=state.regions,
            successors=state.successors,
        )

    def resolve_constraint_variables(self, state: ParsingState, op_def: OpDef):
        """
        Runs verification on the parsed parts of the operation, adding the resolved value
        of each constraint variable to the `ConstraintContext` `state.context`.
        """
        ctx = state.context

        for operand, operand_type, (_, operand_def) in zip(
            state.operands, state.operand_types, op_def.operands, strict=True
        ):
            length = len(operand) if isinstance(operand, Sequence) else 1
            operand_def.constr.verify_length(length, ctx)
            if operand_type is None:
                continue
            if isinstance(operand_type, Attribute):
                operand_type = (operand_type,)
            operand_def.constr.verify(operand_type, ctx)

        for result_type, (_, result_def) in zip(
            state.result_types, op_def.results, strict=True
        ):
            if result_type is None:
                continue
            if isinstance(result_type, Attribute):
                result_type = (result_type,)
            result_def.constr.verify(result_type, ctx)

        for prop_name, prop_def in op_def.properties.items():
            if isinstance(prop_def, OptionalDef) and prop_def.default_value is None:
                continue
            attr = state.properties.get(prop_name, prop_def.default_value)
            if attr is None:
                continue
            prop_def.constr.verify(attr, ctx)

        for attr_name, attr_def in op_def.attributes.items():
            if isinstance(attr_def, OptionalDef) and attr_def.default_value is None:
                continue
            attr = state.attributes.get(attr_name, attr_def.default_value)
            if attr is None:
                continue
            attr_def.constr.verify(attr, ctx)

    def resolve_operand_types(
        self, state: ParsingState, op_def: OpDef
    ) -> Sequence[Sequence[Attribute]]:
        """
        Use the inferred type resolutions to fill missing operand types from other parsed
        types.
        """
        return tuple(
            operand_def.constr.infer(
                state.context,
                length=len(cast(Sequence[Any], operand)),
            )
            if operand_type is None
            else operand_type
            for operand_type, operand, (_, operand_def) in zip(
                state.operand_types, state.operands, op_def.operands, strict=True
            )
        )

    def resolve_result_types(
        self, state: ParsingState, op_def: OpDef
    ) -> Sequence[Sequence[Attribute]]:
        """
        Use the inferred type resolutions to fill missing result types from other parsed
        types.
        """
        return tuple(
            result_def.constr.infer(
                state.context,
                length=None,
            )
            if result_type is None
            else result_type
            for result_type, (_, result_def) in zip(
                state.result_types, op_def.results, strict=True
            )
        )

    def print(self, printer: Printer, op: IRDLOperation) -> None:
        """
        Print the operation with this format.
        The given operation is expected to be defined using the operation definition
        passed to the FormatParser that created this FormatProgram.
        """
        state = PrintingState()
        for stmt in self.stmts:
            stmt.print(printer, state, op)

@dataclass(frozen=True)
class Directive(ABC):
    """An assembly format directive"""

    def is_present(self, op: IRDLOperation) -> bool:
        """
        Check if the directive is present in the input.
        """
        return True

    def is_anchorable(self) -> bool:
        """
        Can appear as an anchor in an optional group.
        """
        return False

    def is_variadic_like(self) -> bool:
        """
        Variadic-like format directives parse a comma separated list, and cannot be
        followed by `,` directive.
        """
        return False

    def is_optional_like(self) -> bool:
        """
        Directives that successfully parse the empty string.
        The first element in an optional group must be optional-like.
        """
        return self.is_variadic_like()

class FormatDirective(Directive, ABC):
    """A format directive for operation format."""

    @abstractmethod
    def parse(self, parser: Parser, state: ParsingState) -> None:
        """
        Parses the directive.
        """
        ...

    def parse_optional(self, parser: Parser, state: ParsingState) -> None:
        """
        Optionally parses a directive.
        """
        self.parse(parser, state)

    @abstractmethod
    def print(
        self, printer: Printer, state: PrintingState, op: IRDLOperation
    ) -> None: ...

    def set_empty(self, state: ParsingState):
        """
        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

class CustomDirective(FormatDirective, ABC):
    """
    A user defined assembly format directive.
    A custom directive can have multiple parameters, whose types should be
    declared in the `parameters` field.
    """

    parameters: ClassVar[dict[str, type[FormatDirective]]]

class TypeableDirective(Directive, ABC):
    """
    Directives which can be used to set or get types.
    """

    @abstractmethod
    def set_types(self, state: ParsingState, types: Sequence[Attribute]) -> None:
        """
        Set the types for this directive to the given sequence.
        """
        ...

    @abstractmethod
    def parse_types(self, parser: Parser, state: ParsingState) -> None:
        """
        Parses types for the directive.
        """
        ...

    @abstractmethod
    def parse_single_type(self, parser: Parser, state: ParsingState) -> None:
        """
        Parse exactly one type for the directive.
        """

    @abstractmethod
    def get_types(self, op: IRDLOperation) -> Sequence[Attribute]: ...

@dataclass(frozen=True)
class TypeDirective(FormatDirective):
    """
    A directive which parses the type of a typeable directive, with format:
      type-directive ::= type(typeable-directive)
    """

    inner: TypeableDirective

    def set(self, state: ParsingState, types: Sequence[Attribute]):
        self.inner.set_types(state, types)

    def parse(self, parser: Parser, state: ParsingState):
        self.inner.parse_types(parser, state)

    def get(self, op: IRDLOperation) -> Sequence[Attribute]:
        return self.inner.get_types(op)

    def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None:
        types = self.inner.get_types(op)
        if not types:
            return
        state.print_whitespace(printer)
        printer.print_list(types, printer.print_attribute)

    def is_present(self, op: IRDLOperation) -> bool:
        return self.inner.is_present(op)

    def is_anchorable(self) -> bool:
        return self.inner.is_anchorable()

    def is_variadic_like(self) -> bool:
        return self.inner.is_variadic_like()

    def is_optional_like(self) -> bool:
        return self.inner.is_optional_like()

    def set_empty(self, state: ParsingState):
        self.set(state, ())

@dataclass(frozen=True)
class VariableDirective(Directive, ABC):
    """
    A variable directive, with the following format:
      variable-directive ::= dollar-ident
    The directive will request a space to be printed after.
    """

    name: str
    """The variable name. This is only used for error message reporting."""
    index: int
    """Index of the variable(operand or result) definition."""

class VariadicVariable(VariableDirective, ABC):
    def is_present(self, op: IRDLOperation) -> bool:
        return bool(getattr(op, self.name))

    def is_anchorable(self) -> bool:
        return True

    def is_variadic_like(self) -> bool:
        return True

class OptionalVariable(VariableDirective, ABC):
    def is_present(self, op: IRDLOperation) -> bool:
        return getattr(op, self.name) is not None

    def is_anchorable(self) -> bool:
        return True

    def is_optional_like(self) -> bool:
        return True

@dataclass(frozen=True)
class AttrDictDirective(FormatDirective):
    """
    An attribute dictionary directive, with the following format:
       attr-dict-directive ::= attr-dict
       attr-dict-with-format-directive ::= `attributes` attr-dict
    The directive (with and without the keyword) will always print a space before, and
    will not request a space to be printed after.
    """

    with_keyword: bool
    """If this is set, the format starts with the `attributes` keyword."""

    reserved_attr_names: set[str]
    """
    The set of attributes that should not be printed.
    These attributes are printed in other places in the format, and thus would be
    printed twice otherwise.
    """

    expected_properties: set[str]
    """
    Properties that should be printed and parsed as part of this attr-dict.
    This is used to keep compatibility with MLIR which allows that.
    """

    def parse(self, parser: Parser, state: ParsingState):
        if self.with_keyword:
            attrs = parser.parse_optional_attr_dict_with_keyword()
            if attrs is None:
                attrs = {}
            else:
                attrs = dict(attrs.data)
        else:
            attrs = parser.parse_optional_attr_dict()
        defined_reserved_keys = self.reserved_attr_names & attrs.keys()
        if defined_reserved_keys:
            parser.raise_error(
                f"attributes {', '.join(defined_reserved_keys)} are defined in other parts of the "
                "assembly format, and thus should not be defined in the attribute "
                "dictionary."
            )

        props = tuple(k for k in attrs.keys() if k in self.expected_properties)
        for name in props:
            state.properties[name] = attrs.pop(name)
        state.attributes |= attrs

    def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None:
        if not op.attributes.keys().isdisjoint(self.expected_properties):
            raise ValueError(
                "Cannot print attributes and properties with the same name "
                "in a single dictionary"
            )
        op_def = op.get_irdl_definition()
        dictionary = op.attributes | {
            k: v for k, v in op.properties.items() if k in self.expected_properties
        }
        defs = {
            x: op_def.properties[x] for x in self.expected_properties
        } | op_def.attributes

        reserved_or_default = self.reserved_attr_names.union(
            name
            for name, d in defs.items()
            if d.default_value is not None and dictionary.get(name) == d.default_value
        )

        printed = printer.print_op_attributes(
            dictionary,
            reserved_attr_names=reserved_or_default,
            print_keyword=self.with_keyword,
        )

        if printed:
            state.last_was_punctuation = False
            state.should_emit_space = True

    def is_optional_like(self) -> bool:
        return True