Csl

class PtrKind(StrEnum):
    SINGLE = "single"
    MANY = "many"

class PtrConst(StrEnum):
    CONST = "const"
    VAR = "var"

class ModuleKind(StrEnum):
    LAYOUT = "layout"
    PROGRAM = "program"

class TaskKind(StrEnum):
    LOCAL = "local"
    DATA = "data"
    CONTROL = "control"

class DsdKind(StrEnum):
    mem1d_dsd = "mem1d_dsd"
    mem4d_dsd = "mem4d_dsd"
    fabin_dsd = "fabin_dsd"
    fabout_dsd = "fabout_dsd"

class Direction(StrEnum):
    NORTH = "north"
    SOUTH = "south"
    EAST = "east"
    WEST = "west"

@dataclass(frozen=True)
class InModuleKind(OpTrait):
    """
    Constrain an op to a particular module kind

    Optionally specify if the op has to be a direct child of CslModuleOp
    (default is yes).

    Ops with this trait are always allowed inside a csl_wrapper.module
    """

    kind: ModuleKind = field()
    _: KW_ONLY
    direct_child: bool = field(default=True)

    def verify(self, op: Operation) -> None:
        from xdsl.dialects import csl_wrapper

        kind: ModuleKind = self.kind
        direct_child: bool = self.direct_child

        direct = "direct" if direct_child else "indirect"
        parent_module = op.parent_op()
        if not direct_child:
            while parent_module is not None and not isinstance(
                parent_module, CslModuleOp | csl_wrapper.ModuleOp
            ):
                parent_module = parent_module.parent_op()
        if isinstance(parent_module, csl_wrapper.ModuleOp):
            return
        if not isinstance(parent_module, CslModuleOp):
            raise VerifyException(
                f"'{op.name}' expexts {direct} parent to be {CslModuleOp.name}, got {parent_module}"
            )
        if parent_module.kind.data != kind:
            raise VerifyException(
                f"'{op.name}' expexts {direct} parent to be {CslModuleOp.name} of kind {kind.value}"
            )

@irdl_attr_definition
class ComptimeStructType(ParametrizedAttribute, TypeAttribute):
    """
    Represents a compile time struct.

    The type makes no guarantees on the fields available.
    """

    name = "csl.comptime_struct"

@irdl_attr_definition
class ImportedModuleType(ParametrizedAttribute, TypeAttribute):
    """
    Represents an imported module (behaves the same as a comptime_struct otherwise).

    The type makes no guarantees on the fields available.
    """

    name = "csl.imported_module"

@irdl_attr_definition
class ComptimeStringType(ParametrizedAttribute, TypeAttribute):
    """
    Represents a compile time string.
    """

    name = "csl.comptime_string"

@irdl_attr_definition
class PtrKindAttr(EnumAttribute[PtrKind], SpacedOpaqueSyntaxAttribute):
    """Attribute representing whether a pointer is a single (*) or many ([*]) pointer"""

    name = "csl.ptr_kind"

@irdl_attr_definition
class PtrConstAttr(EnumAttribute[PtrConst], SpacedOpaqueSyntaxAttribute):
    """Attribute representing whether a pointer's mutability"""

    name = "csl.ptr_const"

@irdl_attr_definition
class ModuleKindAttr(EnumAttribute[ModuleKind], SpacedOpaqueSyntaxAttribute):
    """Attribute representing the kind of CSL module, either layout or program"""

    name = "csl.module_kind"

@irdl_attr_definition
class TaskKindAttr(EnumAttribute[TaskKind], SpacedOpaqueSyntaxAttribute):
    name = "csl.task_kind"

    def get_color_bits(self):
        match self.data:
            case TaskKind.LOCAL | TaskKind.DATA:
                return 5
            case TaskKind.CONTROL:
                return 6

@irdl_attr_definition
class DirectionAttr(EnumAttribute[Direction], SpacedOpaqueSyntaxAttribute):
    name = "csl.dir_kind"

@irdl_attr_definition
class DirectionType(ParametrizedAttribute, TypeAttribute):
    name = "csl.direction"

@irdl_attr_definition
class PtrType(ParametrizedAttribute, TypeAttribute, ContainerType):
    """
    Represents a typed pointer in CSL.

    kind refers to CSL having two types of pointers, single `*type` and many `[*]type`.
    """

    name = "csl.ptr"

    type: TypeAttribute
    kind: PtrKindAttr
    constness: PtrConstAttr

    @staticmethod
    def get(typ: Attribute, is_single: bool, is_const: bool):
        assert isinstance(typ, TypeAttribute)
        return PtrType(
            typ,
            PtrKindAttr(PtrKind.SINGLE if is_single else PtrKind.MANY),
            PtrConstAttr(PtrConst.CONST if is_const else PtrConst.VAR),
        )

    def get_element_type(self) -> Attribute:
        return self.type

@irdl_op_definition
class PtrCastOp(IRDLOperation):
    """
    Implements `@ptrcast(destination_ptr_type, ptr)`
    """

    name = "csl.ptrcast"

    ptr = operand_def(PtrType)
    result = result_def(PtrType)

    traits = traits_def(NoMemoryEffect())

    def __init__(self, ptr: Operation | SSAValue, result_type: PtrType):
        super().__init__(operands=[ptr], result_types=[result_type])

@irdl_op_definition
class IntPtrCastOp(IRDLOperation):
    """
    Implements `@bitcast(destination_type, value)` between a pointer and an integer.

    Exactly one of the input and result types must be a pointer, the other must be
    an integer or index. Use `csl.ptrcast` for pointer to pointer casts.
    """

    name = "csl.intptrcast"

    input = operand_def(IntPtrCastTypeConstr)
    result = result_def(IntPtrCastTypeConstr)

    traits = traits_def(NoMemoryEffect())

    assembly_format = "$input attr-dict `:` type($input) `to` type($result)"

    def __init__(
        self,
        input: Operation | SSAValue,
        result_type: PtrType | IntegerType | IndexType,
    ):
        super().__init__(operands=[input], result_types=[result_type])

    def verify_(self) -> None:
        in_ptr = isinstance(self.input.type, PtrType)
        out_ptr = isinstance(self.result.type, PtrType)
        if in_ptr == out_ptr:
            raise VerifyException(
                "Expected exactly one of input and result to be a pointer, got "
                f"{self.input.type} and {self.result.type}"
            )

@irdl_attr_definition
class DsdType(EnumAttribute[DsdKind], TypeAttribute, SpacedOpaqueSyntaxAttribute):
    """
    Represents a DSD in CSL.
    """

    name = "csl.dsd"

@irdl_attr_definition
class ColorType(ParametrizedAttribute, TypeAttribute):
    """
    Type representing a `color` type in CSL
    """

    name = "csl.color"

@irdl_attr_definition
class VarType(ParametrizedAttribute, TypeAttribute, ContainerType):
    name = "csl.var"

    child_type: TypeAttribute

    def get_element_type(self) -> TypeAttribute:
        return self.child_type

@irdl_op_definition
class VariableOp(IRDLOperation):
    """
    Declares a variable.

    The variable cannot be mutated directly. SSA values of the variable have to
    be loaded and stored using LoadVarOp and StoreVarOp.

    This is similar to how `memref` works.
    """

    name = "csl.variable"

    default = opt_prop_def(ParamAttr)
    res = result_def(VarType)

    assembly_format = "`(` $default `)` `:` type($res) attr-dict"

    def get_element_type(self):
        assert isinstance(self.res.type, VarType)
        return self.res.type.get_element_type()

    @staticmethod
    def from_type(child_type: Attribute) -> VariableOp:
        assert isinstance(child_type, TypeAttribute)
        return VariableOp(result_types=[VarType(child_type)])

    @staticmethod
    def from_value(value: ParamAttr) -> VariableOp:
        return VariableOp(
            properties={"default": value},
            result_types=[VarType(value.type)],
        )

    def verify_(self) -> None:
        assert isinstance(self.res.type, VarType)
        if self.default is not None and (
            self.default.type != self.res.type.get_element_type()
        ):
            raise VerifyException(
                "The type of the default value has to be the same as the type of the result, if the former is supplied"
            )
        return super().verify_()

@irdl_op_definition
class LoadVarOp(IRDLOperation):
    """
    Obtain the SSA value of a CSL variable. The obtained value itself is not
    modifiable, but it can be stored in the variable using `StoreVarOp`.
    """

    name = "csl.load_var"
    var = operand_def(VarType)
    res = result_def()

    traits = traits_def(MemoryReadEffect())

    assembly_format = "`(` $var `:` type($var) `)` `:` type($res) attr-dict"

    def __init__(self, var: VariableOp | SSAValue):
        if isinstance(var, SSAValue):
            assert isinstance(var.type, VarType)
            result_t = var.type.get_element_type()
        else:
            result_t = var.get_element_type()
        super().__init__(
            operands=[var],
            result_types=[result_t],
        )

    def verify_(self) -> None:
        assert isinstance(self.var.type, VarType)
        if self.var.type.get_element_type() != self.res.type:
            raise VerifyException(
                "Result type of the load has to match the child type of the variable"
            )
        return super().verify_()

@irdl_op_definition
class StoreVarOp(IRDLOperation):
    """
    Update the value of a variable.
    """

    name = "csl.store_var"
    var = operand_def(VarType)
    new_value = operand_def()

    traits = traits_def(MemoryWriteEffect())

    assembly_format = (
        "$var `:` type($var) `=` $new_value `:` type($new_value) attr-dict"
    )

    def __init__(self, var: VariableOp, new_value: Operation | SSAValue):
        super().__init__(operands=[var, new_value])

    def verify_(self) -> None:
        assert isinstance(self.var.type, VarType)
        if self.var.type.get_element_type() != self.new_value.type:
            raise VerifyException(
                f"New value must match the element type of {self.var.type.name}"
            )
        return super().verify_()

@irdl_op_definition
class DirectionOp(IRDLOperation):
    name = "csl.get_dir"

    dir = prop_def(DirectionAttr)

    res = result_def(DirectionType)

    traits = traits_def(NoMemoryEffect())

    def __init__(self, direction: DirectionAttr | Direction):
        if isinstance(direction, Direction):
            direction = DirectionAttr(direction)
        super().__init__(properties={"dir": direction}, result_types=[DirectionType()])

@irdl_op_definition
class CslModuleOp(IRDLOperation):
    """
    Separates layout module from program module
    """

    name = "csl.module"
    body = region_def("single_block")
    kind = prop_def(ModuleKindAttr)
    sym_name = prop_def(SymbolNameConstraint())

    traits = traits_def(
        HasParent(ModuleOp),
        IsolatedFromAbove(),
        NoTerminator(),
        SymbolOpInterface(),
    )

@irdl_op_definition
class ImportModuleConstOp(IRDLOperation):
    """
    Equivalent to an `const <va_name> = @import_module("<module_name>", <params>)` call.
    """

    name = "csl.import_module"

    traits = traits_def(HasParent(CslModuleOp))

    module = prop_def(StringAttr)

    params = opt_operand_def(StructLikeConstr)

    result = result_def(ImportedModuleType)

    def __init__(
        self, name: str | StringAttr, params: SSAValue | Operation | None = None
    ):
        if isinstance(name, str):
            name = StringAttr(name)
        super().__init__(
            operands=[params],
            result_types=[ImportedModuleType()],
            properties={"module": name},
        )

@irdl_op_definition
class ConstStrOp(IRDLOperation, HasFolderInterface):
    """
    A compile time string literal.
    """

    name = "csl.const_str"

    traits = traits_def(Pure(), ConstantLike())

    value = prop_def(StringAttr)
    res = result_def(ComptimeStringType)

    assembly_format = "$value attr-dict"

    def __init__(self, value: str | StringAttr):
        if isinstance(value, str):
            value = StringAttr(value)
        super().__init__(
            properties={"value": value},
            result_types=[ComptimeStringType()],
        )

    def fold(self) -> Sequence[SSAValue | Attribute] | None:
        return (self.value,)