Llvm

@irdl_attr_definition
class LLVMStructType(ParametrizedAttribute, TypeAttribute):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/LLVM/#structure-types).
    """

    name = "llvm.struct"

    # An empty string refers to a struct without a name.
    struct_name: StringAttr
    types: ArrayAttr[Attribute]

    # TODO: Add this parameter once xDSL supports the necessary capabilities.
    #  bitmask: StringAttr

    @staticmethod
    def from_type_list(types: Sequence[Attribute]) -> LLVMStructType:
        return LLVMStructType(StringAttr(""), ArrayAttr(types))

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            if self.struct_name.data:
                printer.print_attribute(self.struct_name)
                printer.print_string(", ")
            with printer.in_parens():
                printer.print_list(self.types.data, printer.print_attribute)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> tuple[StringAttr, ArrayAttr]:
        parser.parse_characters("<", " in LLVM struct")
        struct_name = parser.parse_optional_str_literal()
        if struct_name is None:
            struct_name = ""
        else:
            parser.parse_characters(",", " after type")

        params = parser.parse_comma_separated_list(
            parser.Delimiter.PAREN, lambda: parse_llvm_type(parser)
        )
        parser.parse_characters(">", " to close LLVM struct parameters")
        return (StringAttr(struct_name), ArrayAttr(params))

@irdl_attr_definition
class LLVMPointerType(ParametrizedAttribute, TypeAttribute):
    name = "llvm.ptr"

    addr_space: IntAttr | NoneAttr

    def print_parameters(self, printer: Printer) -> None:
        if isinstance(self.addr_space, IntAttr):
            with printer.in_angle_brackets():
                printer.print_int(self.addr_space.data)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> tuple[IntAttr | NoneAttr]:
        if parser.parse_optional_characters("<") is None:
            return (NoneAttr(),)
        addr_space = parser.parse_integer()
        parser.parse_characters(">", " to end llvm.ptr parameters")
        return (IntAttr(addr_space),)

    def __init__(self, addr_space: IntAttr | NoneAttr = NoneAttr()):
        super().__init__(addr_space)

@irdl_attr_definition
class LLVMArrayType(ParametrizedAttribute, TypeAttribute):
    name = "llvm.array"

    size: IntAttr = param_def(converter=IntAttr.get)
    type: Attribute

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_int(self.size.data)
            printer.print_string(" x ")
            printer.print_attribute(self.type)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> tuple[IntAttr, Attribute]:
        with parser.in_angle_brackets():
            size = IntAttr(parser.parse_integer())
            parser.parse_shape_delimiter()
            type = parse_llvm_type(parser)
        return (size, type)

@irdl_attr_definition
class LLVMVoidType(ParametrizedAttribute, TypeAttribute):
    name = "llvm.void"

@irdl_attr_definition
class LLVMFunctionType(ParametrizedAttribute, TypeAttribute):
    """
    Currently does not support variadics.

    See external [documentation](https://mlir.llvm.org/docs/Dialects/LLVM/#function-types).
    """

    name = "llvm.func"

    inputs: ArrayAttr[Attribute]
    output: Attribute
    variadic: UnitAttr | NoneAttr

    def __init__(
        self,
        inputs: Sequence[Attribute] | ArrayAttr[Attribute],
        output: Attribute | None = None,
        is_variadic: bool = False,
    ) -> None:
        if not isinstance(inputs, ArrayAttr):
            inputs = ArrayAttr(inputs)
        if output is None:
            output = LLVMVoidType()
        variad_attr = UnitAttr() if is_variadic else NoneAttr()
        super().__init__(inputs, output, variad_attr)

    @property
    def is_variadic(self) -> bool:
        return isinstance(self.variadic, UnitAttr)

    def verify(self) -> None:
        for i, inp in enumerate(self.inputs):
            if not is_compatible_type(inp):
                raise VerifyException(
                    f"LLVM function argument #{i} has incompatible type '{inp}'"
                )
        if not isinstance(self.output, LLVMVoidType) and not is_compatible_type(
            self.output
        ):
            raise VerifyException(
                f"LLVM function result has incompatible type '{self.output}'"
            )

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            if isinstance(self.output, LLVMVoidType):
                printer.print_string("void")
            else:
                printer.print_attribute(self.output)

            printer.print_string(" ")
            with printer.in_parens():
                printer.print_list(self.inputs, printer.print_attribute)
                if self.is_variadic:
                    if self.inputs:
                        printer.print_string(", ")
                    printer.print_string("...")

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        parser.parse_characters("<", " in llvm.func parameters")
        output = parse_llvm_type(parser)

        # save pos before args for error message printing
        pos = parser.pos

        def _parse_attr_or_variadic() -> Attribute | EllipsisType:
            """
            This returns either an attribute, or Ellipsis if a
            varargs specifier (`...`) was parsed.
            """
            if parser.parse_optional_characters("...") is not None:
                return ...
            return parse_llvm_type(parser)

        inputs = parser.parse_comma_separated_list(
            Parser.Delimiter.PAREN, _parse_attr_or_variadic
        )
        is_varargs: NoneAttr | UnitAttr = NoneAttr()
        if inputs and inputs[-1] is Ellipsis:
            is_varargs = UnitAttr()
            inputs = inputs[:-1]

        if not isa(inputs, list[Attribute]):
            parser.raise_error(
                "Varargs specifier `...` must be at the end of the argument definition",
                pos,
                parser.pos,
            )

        parser.parse_characters(">", " in llvm.func parameters")

        return [ArrayAttr(inputs), output, is_varargs]

@irdl_attr_definition
class LinkageAttr(ParametrizedAttribute):
    name = "llvm.linkage"

    linkage: StringAttr

    def __init__(self, linkage: str | StringAttr) -> None:
        if isinstance(linkage, str):
            linkage = StringAttr(linkage)
        super().__init__(linkage)

    def print_parameters(self, printer: Printer) -> None:
        printer.print_string("<")
        printer.print_attribute(self.linkage)
        printer.print_string(">")

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        parser.parse_characters("<", "llvm.linkage parameter expected")
        # The linkage string is output from xDSL as a string (and accepted by MLIR as such)
        # however it is always output from MLIR without quotes. Therefore need to determine
        # whether this is a string or not and slightly change how we parse based upon that
        linkage_str = parser.parse_optional_str_literal()
        if linkage_str is None:
            linkage_str = parser.parse_identifier()
        linkage = StringAttr(linkage_str)
        parser.parse_characters(">", " to end llvm.linkage parameters")
        return [linkage]

    def verify(self):
        if self.linkage.data not in _LINKAGE_OPTIONS:
            raise VerifyException(f"Specified linkage '{self.linkage.data}' is unknown")

class ArithmeticBinOperation(IRDLOperation, ABC):
    """Class for arithmetic binary operations."""

    T: ClassVar = VarConstraint(
        "T", SignlessIntegerConstraint | VectorType.constr(SignlessIntegerConstraint)
    )

    lhs = operand_def(T)
    rhs = operand_def(T)
    res = result_def(T)

    traits = traits_def(NoMemoryEffect())

    def __init__(
        self,
        lhs: SSAValue,
        rhs: SSAValue,
        attributes: dict[str, Attribute] = {},
    ):
        super().__init__(
            operands=[lhs, rhs],
            attributes=attributes,
            result_types=[lhs.type],
        )

    @classmethod
    def parse(cls, parser: Parser):
        lhs = parser.parse_unresolved_operand()
        parser.parse_characters(",")
        rhs = parser.parse_unresolved_operand()
        attributes = parser.parse_optional_attr_dict()
        parser.parse_characters(":")
        type = parser.parse_type()
        operands = parser.resolve_operands([lhs, rhs], [type, type], parser.pos)
        return cls(operands[0], operands[1], attributes)

    def print(self, printer: Printer) -> None:
        printer.print_string(" ")
        printer.print_ssa_value(self.lhs)
        printer.print_string(", ")
        printer.print_ssa_value(self.rhs)
        printer.print_op_attributes(self.attributes)
        printer.print_string(" : ")
        printer.print_attribute(self.lhs.type)

class OverflowFlag(StrEnum):
    NO_SIGNED_WRAP = "nsw"
    NO_UNSIGNED_WRAP = "nuw"

@dataclass(frozen=True, init=False)
class OverflowAttrBase(BitEnumAttribute[OverflowFlag]):
    none_value = "none"

@irdl_attr_definition(init=False)
class OverflowAttr(OverflowAttrBase):
    name = "llvm.overflow"

    @classmethod
    def parse(cls, parser: Parser) -> OverflowAttr:
        if parser.parse_optional_keyword("overflow") is not None:
            return OverflowAttr(OverflowAttr.parse_parameter(parser))
        return OverflowAttr("none")

    def print(self, printer: Printer):
        if self.data:
            printer.print_string(" overflow")
            self.print_parameter(printer)

    def to_int(self) -> int:
        res = 0
        if OverflowFlag.NO_SIGNED_WRAP in self.data:
            res |= 1
        if OverflowFlag.NO_UNSIGNED_WRAP in self.data:
            res |= 2
        return res

    @staticmethod
    def from_int(i: int) -> OverflowAttr:
        match i:
            case 0:
                return OverflowAttr("none")
            case 1:
                return OverflowAttr((OverflowFlag.NO_SIGNED_WRAP,))
            case 2:
                return OverflowAttr((OverflowFlag.NO_UNSIGNED_WRAP,))
            case 3:
                return OverflowAttr(
                    (OverflowFlag.NO_SIGNED_WRAP, OverflowFlag.NO_UNSIGNED_WRAP)
                )
            case _:
                raise ValueError("OverflowAttr given out of bounds integer.")

class ArithmeticBinOpOverflow(IRDLOperation, ABC):
    """Class for arithmetic binary operations that use overflow flags."""

    T: ClassVar = VarConstraint(
        "T", SignlessIntegerConstraint | VectorType.constr(SignlessIntegerConstraint)
    )

    lhs = operand_def(T)
    rhs = operand_def(T)
    res = result_def(T)
    overflowFlags = opt_prop_def(IntegerAttr[I32])

    traits = traits_def(NoMemoryEffect())

    def __init__(
        self,
        lhs: SSAValue,
        rhs: SSAValue,
        attributes: dict[str, Attribute] = {},
        overflow: OverflowAttr | IntegerAttr = IntegerAttr(0, 32),
    ):
        if isinstance(overflow, OverflowAttr):
            overflow = IntegerAttr(overflow.to_int(), 32)
        super().__init__(
            operands=[lhs, rhs],
            attributes=attributes,
            result_types=[lhs.type],
            properties={
                "overflowFlags": overflow,
            },
        )

    @classmethod
    def parse(cls, parser: Parser):
        lhs = parser.parse_unresolved_operand()
        parser.parse_characters(",")
        rhs = parser.parse_unresolved_operand()
        overflowFlags = OverflowAttr.parse(parser)
        attributes = parser.parse_optional_attr_dict()
        parser.parse_characters(":")
        type = parser.parse_type()
        operands = parser.resolve_operands([lhs, rhs], [type, type], parser.pos)
        return cls(operands[0], operands[1], attributes, overflowFlags)

    def print(self, printer: Printer) -> None:
        printer.print_string(" ")
        printer.print_ssa_value(self.lhs)
        printer.print_string(", ")
        printer.print_ssa_value(self.rhs)
        if self.overflowFlags:
            OverflowAttr.from_int(self.overflowFlags.value.data).print(printer)
        printer.print_op_attributes(self.attributes)
        printer.print_string(" : ")
        printer.print_attribute(self.lhs.type)