Affine

@irdl_op_definition
class ApplyOp(IRDLOperation):
    name = "affine.apply"

    mapOperands = var_operand_def(IndexType)
    map = prop_def(AffineMapAttr)
    result = result_def(IndexType)

    traits = traits_def(Pure())

    def __init__(self, map_operands: Sequence[SSAValue], affine_map: AffineMapAttr):
        super().__init__(
            operands=[map_operands],
            properties={"map": affine_map},
            result_types=[IndexType()],
        )

    def verify_(self) -> None:
        if len(self.mapOperands) != self.map.data.num_dims + self.map.data.num_symbols:
            raise VerifyException(
                f"{self.name} expects "
                f"{self.map.data.num_dims + self.map.data.num_symbols} operands, but "
                f"got {len(self.mapOperands)}. The number of map operands must match "
                "the sum of the dimensions and symbols of its map."
            )
        if len(self.map.data.results) != 1:
            raise VerifyException("affine.apply expects a unidimensional map.")

    @classmethod
    def parse(cls, parser: Parser) -> ApplyOp:
        pos = parser.pos
        m = parser.parse_attribute()
        if not isinstance(m, AffineMapAttr):
            parser.raise_error("Expected affine map attr", at_position=pos)
        dims = parser.parse_optional_comma_separated_list(
            parser.Delimiter.PAREN, lambda: parser.parse_operand()
        )
        if dims is None:
            dims = []
        syms = parser.parse_optional_comma_separated_list(
            parser.Delimiter.SQUARE, lambda: parser.parse_operand()
        )
        if syms is None:
            syms = []
        return ApplyOp(dims + syms, m)

    def print(self, printer: Printer):
        m = self.map.data
        operands = tuple(self.mapOperands)
        assert len(operands) == m.num_dims + m.num_symbols, f"{len(operands)} {m}"
        printer.print_string(" ")
        printer.print_attribute(self.map)
        printer.print_string(" (")
        if m.num_dims:
            printer.print_list(
                operands[: m.num_dims], lambda el: printer.print_operand(el)
            )
        printer.print_string(")")

        if m.num_symbols:
            printer.print_string("[")
            printer.print_list(
                operands[m.num_dims :], lambda el: printer.print_operand(el)
            )
            printer.print_string("]")

@irdl_op_definition
class ForOp(IRDLOperation):
    name = "affine.for"

    lowerBoundOperands = var_operand_def(IndexType)
    upperBoundOperands = var_operand_def(IndexType)
    inits = var_operand_def()
    res = var_result_def()

    lowerBoundMap = prop_def(AffineMapAttr)
    upperBoundMap = prop_def(AffineMapAttr)
    step = prop_def(IntegerAttr)

    body = region_def()

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    # TODO this requires the ImplicitAffineTerminator trait instead of
    # NoTerminator
    # gh issue: https://github.com/xdslproject/xdsl/issues/1149

    def verify_(self) -> None:
        if len(self.inits) != len(self.results):
            raise VerifyException("Expected as many init operands as results.")
        if len(self.lowerBoundOperands) != (
            self.lowerBoundMap.data.num_dims + self.lowerBoundMap.data.num_symbols
        ):
            raise VerifyException(
                "Expected as many lower bound operands as lower bound dimensions and symbols."
            )
        if len(self.upperBoundOperands) != (
            self.upperBoundMap.data.num_dims + self.upperBoundMap.data.num_symbols
        ):
            raise VerifyException(
                "Expected as many upper bound operands as upper bound dimensions and symbols."
            )
        iter_types = self.inits.types
        if iter_types != self.result_types:
            raise VerifyException(
                "Expected all operands and result pairs to have matching types"
            )
        entry_block: Block = self.body.blocks[0]
        block_arg_types = (IndexType(), *iter_types)
        arg_types = entry_block.arg_types
        if block_arg_types != arg_types:
            raise VerifyException(
                "Expected BlockArguments to have the same types as the operands"
            )

    @staticmethod
    def from_region(
        lowerBoundOperands: Sequence[Operation | SSAValue],
        upperBoundOperands: Sequence[Operation | SSAValue],
        inits: Sequence[Operation | SSAValue],
        result_types: Sequence[Attribute],
        lower_bound: int | AffineMapAttr,
        upper_bound: int | AffineMapAttr,
        region: Region,
        step: int | IntegerAttr = 1,
    ) -> ForOp:
        if isinstance(lower_bound, int):
            lower_bound = AffineMapAttr(
                AffineMap(0, 0, (AffineExpr.constant(lower_bound),))
            )
        if isinstance(upper_bound, int):
            upper_bound = AffineMapAttr(
                AffineMap(0, 0, (AffineExpr.constant(upper_bound),))
            )
        if isinstance(step, int):
            step = IntegerAttr.from_index_int_value(step)
        properties: dict[str, Attribute] = {
            "lowerBoundMap": lower_bound,
            "upperBoundMap": upper_bound,
            "step": step,
        }
        return ForOp.build(
            operands=[lowerBoundOperands, upperBoundOperands, inits],
            result_types=[result_types],
            properties=properties,
            regions=[region],
        )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        unresolved_indvar = parser.parse_argument(expect_type=False)
        parser.parse_characters("=")

        lower_bound_map, lower_bound_operands = _parse_affine_for_bound(parser, "max")
        parser.parse_characters("to")
        upper_bound_map, upper_bound_operands = _parse_affine_for_bound(parser, "min")

        if parser.parse_optional_characters("step") is not None:
            step_pos = parser.pos
            step = parser.parse_integer(allow_boolean=False)
            if step < 0:
                parser.raise_error(
                    "expected step to be representable as a positive signed integer",
                    step_pos,
                )
        else:
            step = 1

        unresolved_iter_args: Sequence[Parser.UnresolvedArgument] = ()
        iter_arg_operands: Sequence[SSAValue] = ()
        iter_arg_types: Sequence[Attribute] = ()

        if parser.parse_optional_characters("iter_args") is not None:

            def parse_iter_arg() -> tuple[Parser.UnresolvedArgument, SSAValue]:
                arg = parser.parse_argument(expect_type=False)
                parser.parse_characters("=")
                return arg, parser.parse_operand()

            pairs = parser.parse_comma_separated_list(
                Parser.Delimiter.PAREN, parse_iter_arg
            )
            unresolved_iter_args = tuple(arg for arg, _ in pairs)
            iter_arg_operands = tuple(val for _, val in pairs)
            parser.parse_characters("->")

            # MLIR's `parseArrowTypeList` also accepts a single bare type
            # (no parens) when there is only one loop-carried value.
            if parser.parse_optional_punctuation("(") is not None:
                iter_arg_types = parser.parse_comma_separated_list(
                    Parser.Delimiter.NONE, parser.parse_type
                )
                parser.parse_punctuation(")")
            else:
                iter_arg_types = (parser.parse_type(),)

        iter_args = tuple(
            u_arg.resolve(t) for u_arg, t in zip(unresolved_iter_args, iter_arg_types)
        )
        indvar = unresolved_indvar.resolve(IndexType())
        body = parser.parse_region((indvar, *iter_args))

        # affine.for has no implicit-terminator trait (see TODO above), so the
        # terminator omitted from the printed form when there are no iter_args
        # must be re-inserted by hand here, mirroring `ensureTerminator`
        block = body.block
        if block.last_op is None or not isinstance(block.last_op, YieldOp):
            block.add_op(YieldOp.get())

        attributes = parser.parse_optional_attr_dict()

        for_op = cast(
            Self,
            cls.from_region(
                lower_bound_operands,
                upper_bound_operands,
                iter_arg_operands,
                iter_arg_types,
                lower_bound_map,
                upper_bound_map,
                body,
                step,
            ),
        )
        for_op.attributes |= attributes
        return for_op

    def print(self, printer: Printer):
        printer.print_string(" ")
        indvar, *block_iter_args = self.body.block.args
        printer.print_block_argument(indvar, print_type=False)
        printer.print_string(" = ")

        _print_affine_for_bound(
            printer, self.lowerBoundMap, self.lowerBoundOperands, "max"
        )
        printer.print_string(" to ")
        _print_affine_for_bound(
            printer, self.upperBoundMap, self.upperBoundOperands, "min"
        )

        if self.step.value.data != 1:
            printer.print_string(f" step {self.step.value.data}")

        print_block_terminators = False

        if self.inits:
            printer.print_string(" iter_args")

            with printer.in_parens():
                printer.print_list(
                    zip(block_iter_args, self.inits),
                    lambda pair: print_assignment(printer, *pair),
                )

            printer.print_string(" -> ")

            with printer.in_parens():
                printer.print_list(self.result_types, printer.print_attribute)

            print_block_terminators = True

        printer.print_string(" ")
        printer.print_region(
            self.body,
            print_entry_block_args=False,
            print_empty_block=False,
            print_block_terminators=print_block_terminators,
        )
        printer.print_op_attributes(self.attributes)

@irdl_op_definition
class IfOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/Affine/#affineif-affineaffineifop).
    """

    name = "affine.if"

    args = var_operand_def(IndexType)
    res = var_result_def()

    condition = prop_def(AffineSetAttr)

    then_region = region_def("single_block")
    else_region = region_def()

    traits = traits_def(RecursiveMemoryEffect(), RecursivelySpeculatable())

@irdl_op_definition
class ParallelOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/Affine/#affineparallel-affineaffineparallelop).
    """

    name = "affine.parallel"

    map_operands = var_operand_def(IndexType)

    reductions = prop_def(ArrayAttr[StringAttr])
    lowerBoundsMap = prop_def(AffineMapAttr)
    lowerBoundsGroups = prop_def(DenseIntElementsAttr)
    upperBoundsMap = prop_def(AffineMapAttr)
    upperBoundsGroups = prop_def(DenseIntElementsAttr)
    steps = prop_def(ArrayAttr[IntegerAttr[IntegerType]])

    res = var_result_def()

    body = region_def("single_block")

    def verify_(self) -> None:
        if (
            len(self.operands)
            != len(self.results)
            + self.lowerBoundsMap.data.num_dims
            + self.upperBoundsMap.data.num_dims
            + self.lowerBoundsMap.data.num_symbols
            + self.upperBoundsMap.data.num_symbols
        ):
            raise VerifyException(
                "Expected as many operands as results, lower bound args and upper bound args."
            )

        if sum(self.lowerBoundsGroups.get_values()) != len(
            self.lowerBoundsMap.data.results
        ):
            raise VerifyException("Expected a lower bound group for each lower bound")
        if sum(self.upperBoundsGroups.get_values()) != len(
            self.upperBoundsMap.data.results
        ):
            raise VerifyException("Expected an upper bound group for each upper bound")

@irdl_op_definition
class StoreOp(IRDLOperation):
    name = "affine.store"

    T: ClassVar = VarConstraint("T", AnyAttr())

    value = operand_def(T)
    memref = operand_def(MemRefType.constr(T))
    indices = var_operand_def(IndexType)
    map = prop_def(AffineMapAttr)

    def __init__(
        self,
        value: SSAValue,
        memref: SSAValue,
        indices: Sequence[SSAValue],
        map: AffineMapAttr | None = None,
    ):
        if map is None:
            if not isa(memref_type := memref.type, MemRefType):
                raise ValueError(
                    "affine.store memref operand must be of type MemRefType"
                )

            map = _map_or_identity_attr(map, memref_type)

        super().__init__(
            operands=(value, memref, indices),
            properties={"map": map},
        )

    @classmethod
    def parse(cls, parser: Parser) -> StoreOp:
        value = parser.parse_unresolved_operand()
        parser.parse_punctuation(",")
        memref, affine_map, indices, memref_type = _parse_affine_memref_access(parser)
        resolved_value = parser.resolve_operand(value, memref_type.get_element_type())
        return StoreOp(resolved_value, memref, indices, AffineMapAttr(affine_map))

    def print(self, printer: Printer):
        printer.print_string(" ")
        printer.print_ssa_value(self.value)
        printer.print_string(", ")

        _print_affine_memref_access(
            printer, self.memref, self.map.data, self.indices, self.memref.type
        )

@irdl_op_definition
class LoadOp(IRDLOperation):
    name = "affine.load"

    T: ClassVar = VarConstraint("T", AnyAttr())

    memref = operand_def(MemRefType.constr(T))
    indices = var_operand_def(IndexType)

    result = result_def(T)

    map = prop_def(AffineMapAttr)

    def __init__(
        self,
        memref: SSAValue,
        indices: Sequence[SSAValue],
        map: AffineMapAttr | None = None,
        result_type: Attribute | None = None,
    ):
        if map is None:
            # Create identity map for memrefs with at least one dimension or () -> ()
            # for zero-dimensional memrefs.
            if not isinstance(memref.type, ShapedType):
                raise ValueError(
                    "affine.load memref operand must be of type ShapedType"
                )

            map = _map_or_identity_attr(map, memref.type)

        if result_type is None:
            if not isa(memref.type, ContainerType):
                raise ValueError(
                    "affine.load memref operand must be of type ContainerType"
                )

            result_type = memref.type.get_element_type()

        super().__init__(
            operands=(memref, indices),
            properties={"map": map},
            result_types=(result_type,),
        )

    @classmethod
    def parse(cls, parser: Parser) -> LoadOp:
        memref, affine_map, indices, memref_type = _parse_affine_memref_access(parser)
        result_type = memref_type.get_element_type()

        return LoadOp(memref, indices, AffineMapAttr(affine_map), result_type)

    def print(self, printer: Printer):
        printer.print_string(" ")
        _print_affine_memref_access(
            printer, self.memref, self.map.data, self.indices, self.memref.type
        )