Ops

class X86AsmOperation(IRDLOperation, OneLineAssemblyPrintable, ABC):
    """
    Base class for operations that can be a part of x86 assembly printing.
    """

class X86RegisterAllocatableOperation(IRDLOperation, RegisterAllocatableOperation, ABC):
    """
    Base class for x86 operations that can take part in register allocation.
    """

class X86HasRegisterConstraints(
    X86RegisterAllocatableOperation, HasRegisterConstraints, ABC
):
    """
    Base class for x86 operations with register constraints.
    By default, all operands are "in", and all results are "out", subclasses must
    override `get_register_constraints` if some of the results must be in the same
    registers as operands.
    """

    def get_register_constraints(self) -> RegisterConstraints:
        # The default register constraints are that all operands are "in", and all
        # results are "out" registers.
        # If some registers are "inout" then this function must be overridden.
        return RegisterConstraints(self.operands, self.results, ())

class X86CustomFormatOperation(IRDLOperation, ABC):
    @classmethod
    def parse(cls, parser: Parser) -> Self:
        args = cls.parse_unresolved_operands(parser)
        custom_attributes = cls.custom_parse_attributes(parser)
        remaining_attributes = parser.parse_optional_attr_dict()
        # TODO ensure distinct keys for attributes
        attributes = custom_attributes | remaining_attributes
        regions = parser.parse_region_list()
        pos = parser.pos
        operand_types, result_types = cls.parse_op_type(parser)
        operands = parser.resolve_operands(args, operand_types, pos)
        return cls.build(
            operands=operands,
            result_types=result_types,
            attributes=attributes,
            regions=regions,
        )

    @classmethod
    def parse_unresolved_operands(cls, parser: Parser) -> list[UnresolvedOperand]:
        """
        Parse a list of comma separated unresolved operands.
        Notice that this method will consume trailing comma.
        """
        if operand := parser.parse_optional_unresolved_operand():
            operands = [operand]
            while parser.parse_optional_punctuation(",") and (
                operand := parser.parse_optional_unresolved_operand()
            ):
                operands.append(operand)
            return operands
        return []

    @classmethod
    def custom_parse_attributes(cls, parser: Parser) -> dict[str, Attribute]:
        """
        Parse attributes with custom syntax. Subclasses may override this method.
        """
        return parser.parse_optional_attr_dict()

    @classmethod
    def parse_op_type(
        cls, parser: Parser
    ) -> tuple[Sequence[Attribute], Sequence[Attribute]]:
        parser.parse_punctuation(":")
        func_type = parser.parse_function_type()
        return func_type.inputs.data, func_type.outputs.data

    def print(self, printer: Printer) -> None:
        if self.operands:
            printer.print_string(" ")
            printer.print_list(self.operands, printer.print_operand)
        printed_attributes = self.custom_print_attributes(printer)
        unprinted_attributes = {
            name: attr
            for name, attr in self.attributes.items()
            if name not in printed_attributes
        }
        printer.print_op_attributes(unprinted_attributes)
        printer.print_regions(self.regions)
        self.print_op_type(printer)

    def custom_print_attributes(self, printer: Printer) -> AbstractSet[str]:
        """
        Print attributes with custom syntax. Return the names of the attributes printed. Subclasses may override this method.
        """
        printer.print_op_attributes(self.attributes)
        return self.attributes.keys()

    def print_op_type(self, printer: Printer) -> None:
        printer.print_string(" : ")
        printer.print_operation_type(self)

class X86Instruction(X86AsmOperation, X86HasRegisterConstraints):
    """
    Base class for operations that can be a part of x86 assembly printing. Must
    represent an instruction in the x86 instruction set.
    The name of the operation will be used as the x86 assembly instruction name.
    """

    traits = traits_def(RegisterAllocatedMemoryEffect())

    comment = opt_attr_def(StringAttr)
    """
    An optional comment that will be printed along with the instruction.
    """

    @abstractmethod
    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        """
        The arguments to the instruction, in the order they should be printed in the
        assembly.
        """
        raise NotImplementedError()

    def assembly_instruction_name(self) -> str:
        """
        By default, the name of the instruction is the same as the name of the operation.
        """

        return self.name.split(".")[-1]

    def assembly_line(self) -> str | None:
        # default assembly code generator
        instruction_name = self.assembly_instruction_name()
        arg_str = ", ".join(
            assembly_arg_str(arg)
            for arg in self.assembly_line_args()
            if arg is not None
        )
        return AssemblyPrinter.assembly_line(instruction_name, arg_str, self.comment)

class RS_Operation(X86Instruction, ABC, Generic[R1InvT, R2InvT]):
    """
    A base class for x86 operations that have one register that is read and written to,
    and one source register.
    """

    register_in = operand_def(R1InvT)
    register_out: OpResult[R1InvT] = result_def(R1InvT)

    source = operand_def(R2InvT)

    assembly_format = (
        "$register_in `,` $source attr-dict `:` "
        "`(` type($register_in) `,` type($source) `)` `->` type($register_out)"
    )

    def __init__(
        self,
        register_in: Operation | SSAValue,
        source: Operation | SSAValue,
        *,
        comment: str | StringAttr | None = None,
        register_out: R1InvT | None = None,
    ):
        if isinstance(comment, str):
            comment = StringAttr(comment)
        register_in = SSAValue.get(register_in)
        if register_out is None:
            register_out = cast(R1InvT, register_in.type)

        super().__init__(
            operands=[register_in, source],
            attributes={
                "comment": comment,
            },
            result_types=[register_out],
        )

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        return reg(self.register_in), reg(self.source)

    def get_register_constraints(self) -> RegisterConstraints:
        return RegisterConstraints(
            (self.source,), (), ((self.register_in, self.register_out),)
        )

class DS_Operation(X86Instruction, ABC, Generic[R1InvT, R2InvT]):
    """
    A base class for x86 operations that have one destination register and one source
    register.
    """

    destination: OpResult[R1InvT] = result_def(R1InvT)
    source = operand_def(R2InvT)

    assembly_format = (
        "$source attr-dict `:` `(` type($source) `)` `->` type($destination)"
    )

    def __init__(
        self,
        source: Operation | SSAValue,
        *,
        comment: str | StringAttr | None = None,
        destination: R1InvT,
    ):
        if isinstance(comment, str):
            comment = StringAttr(comment)

        super().__init__(
            operands=[source],
            attributes={
                "comment": comment,
            },
            result_types=[destination],
        )

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        return (reg(self.destination), reg(self.source))

class DSK_Operation(X86Instruction, ABC):
    """
    A base class for x86 operations that have one destination register and one source
    register.
    """

    destination: OpResult[X86VectorRegisterType] = result_def(X86VectorRegisterType)
    source = operand_def(X86VectorRegisterType)
    mask_reg = operand_def(AVX512MaskRegisterType)
    z = opt_attr_def(UnitAttr)

    assembly_format = (
        "$source `,` $mask_reg attr-dict `:` "
        "`(` type($source) `,` type($mask_reg) `)` `->` type($destination)"
    )

    def __init__(
        self,
        source: Operation | SSAValue,
        mask_reg: Operation | SSAValue,
        *,
        z: bool = False,
        comment: str | StringAttr | None = None,
        destination: X86VectorRegisterType,
    ):
        if isinstance(comment, str):
            comment = StringAttr(comment)

        super().__init__(
            operands=[source, mask_reg],
            attributes={
                "z": UnitAttr() if z else None,
                "comment": comment,
            },
            result_types=[destination],
        )

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        register_out = masked_source_str(self.destination, self.mask_reg, self.z)
        return register_out, reg(self.source)

class DK_Operation(
    X86Instruction, X86CustomFormatOperation, HasRegisterConstraints, ABC
):
    """
    A base class for x86 operations that have one general purpose destination register
    and one writemask source register.
    """

    destination: OpResult[GeneralRegisterType] = result_def(GeneralRegisterType)
    source = operand_def(AVX512MaskRegisterType)

    def __init__(
        self,
        source: Operation | SSAValue,
        *,
        comment: str | StringAttr | None = None,
        destination: GeneralRegisterType,
    ):
        if isinstance(comment, str):
            comment = StringAttr(comment)

        super().__init__(
            operands=(source,),
            attributes={
                "comment": comment,
            },
            result_types=(destination,),
        )

    def get_register_constraints(self) -> RegisterConstraints:
        return RegisterConstraints((self.source,), (self.destination,), ())

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        return reg(self.destination), reg(self.source)

class KS_Operation(
    X86Instruction, X86CustomFormatOperation, HasRegisterConstraints, ABC
):
    """
    A base class for x86 operations that have one general purpose destination register
    and one writemask source register.
    """

    destination: OpResult[AVX512MaskRegisterType] = result_def(AVX512MaskRegisterType)
    source = operand_def(GeneralRegisterType)

    def __init__(
        self,
        source: Operation | SSAValue,
        *,
        comment: str | StringAttr | None = None,
        destination: AVX512MaskRegisterType,
    ):
        if isinstance(comment, str):
            comment = StringAttr(comment)

        super().__init__(
            operands=(source,),
            attributes={
                "comment": comment,
            },
            result_types=(destination,),
        )

    def get_register_constraints(self) -> RegisterConstraints:
        return RegisterConstraints((self.source,), (self.destination,), ())

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        return reg(self.destination), reg(self.source)

class R_Operation(X86Instruction, ABC, Generic[R1InvT]):
    """
    A base class for x86 operations that have one register that is read and written to.
    """

    register_in = operand_def(R1InvT)
    register_out: OpResult[R1InvT] = result_def(R1InvT)

    assembly_format = (
        "$register_in attr-dict `:` `(` type($register_in) `)` `->` type($register_out)"
    )

    def __init__(
        self,
        register_in: SSAValue[R1InvT],
        *,
        comment: str | StringAttr | None = None,
        register_out: R1InvT | None = None,
    ):
        if isinstance(comment, str):
            comment = StringAttr(comment)
        if register_out is None:
            register_out = register_in.type
        super().__init__(
            operands=[register_in],
            attributes={
                "comment": comment,
            },
            result_types=[register_out],
        )

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        return (reg(self.register_in),)

    def get_register_constraints(self) -> RegisterConstraints:
        return RegisterConstraints((), (), ((self.register_in, self.register_out),))

class RM_Operation(X86Instruction, ABC, Generic[R1InvT, R2InvT]):
    """
    A base class for x86 operations that have one register read and written to and one
    memory access with an optional offset.
    """

    register_in = operand_def(R1InvT)
    register_out: OpResult[R1InvT] = result_def(R1InvT)

    memory = operand_def(R2InvT)
    memory_offset = attr_def(IntegerAttr[I64], default_value=IntegerAttr(0, i64))

    traits = traits_def(MemoryReadEffect())

    assembly_format = (
        "$register_in `,` `[` $memory (`+` $memory_offset^)? `]` attr-dict `:` "
        "`(` type($register_in) `,` type($memory) `)` `->` type($register_out)"
    )

    def __init__(
        self,
        register_in: Operation | SSAValue,
        memory: Operation | SSAValue,
        memory_offset: int | IntegerAttr[I64],
        *,
        comment: str | StringAttr | None = None,
        register_out: R1InvT | None = None,
    ):
        if isinstance(memory_offset, int):
            memory_offset = IntegerAttr(memory_offset, i64)
        if isinstance(comment, str):
            comment = StringAttr(comment)
        register_in = SSAValue.get(register_in)
        if register_out is None:
            register_out = cast(R1InvT, register_in.type)

        super().__init__(
            operands=[register_in, memory],
            attributes={
                "memory_offset": memory_offset,
                "comment": comment,
            },
            result_types=[register_out],
        )

    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
        memory_access = memory_access_str(self.memory, self.memory_offset)
        destination = assembly_arg_str(reg(self.register_in))
        return (destination, memory_access)

    def get_register_constraints(self) -> RegisterConstraints:
        return RegisterConstraints(
            (self.memory,), (), ((self.register_in, self.register_out),)
        )