Format

def print_for_op_like(
    printer: Printer,
    start: IntegerAttr | SSAValue,
    stop: IntegerAttr | SSAValue,
    step: IntegerAttr | SSAValue,
    iter_args: Sequence[SSAValue],
    body: Region,
    default_indvar_type: type[TypeAttribute] | None = None,
    bound_words: Sequence[str] = ["to"],
):
    """
    Prints the loop bounds, step, iteration arguments, and body.

    Users can provide a default induction variable type and specific human-readable
    words for bounds (default: "to").

    Note that providing a default induction variable type is required to suggest that
    all loop control variable types (induction, bounds and step) have the same type,
    hence moving the induction variable type printing to the end of the for expression.
    The induction variable type printing is ommited when it matches the expected default
    type (`default_indvar_type`).

    The `start`, `stop`, and `step` may be dynamic SSAValues or static
    `IntegerAttr`s.
    When static, the typed integer literal is printed (value and type), not an SSA
    value reference.
    """

    block = body.block
    indvar, *block_iter_args = block.args

    printer.print_string(" ")

    def print_indvar_type():
        printer.print_string(" : ")
        printer.print_attribute(indvar.type)
        printer.print_string(" ")

    printer.print_ssa_value(indvar)

    if default_indvar_type is None:
        print_indvar_type()

    printer.print_string(" = ")
    if isinstance(start, IntegerAttr):
        start.print_builtin(printer)
    else:
        printer.print_ssa_value(start)

    for word in bound_words:
        printer.print_string(f" {word} ")

    if isinstance(stop, IntegerAttr):
        stop.print_builtin(printer)
    else:
        printer.print_ssa_value(stop)
    printer.print_string(" step ")
    if isinstance(step, IntegerAttr):
        step.print_builtin(printer)
    else:
        printer.print_ssa_value(step)
    printer.print_string(" ")
    if block_iter_args:
        printer.print_string("iter_args(")
        printer.print_list(
            zip(block_iter_args, iter_args),
            lambda pair: print_assignment(printer, *pair),
        )
        printer.print_string(") -> (")
        printer.print_list((a.type for a in block_iter_args), printer.print_attribute)
        printer.print_string(") ")

    if default_indvar_type is not None and not isinstance(
        indvar.type, default_indvar_type
    ):
        print_indvar_type()

    printer.print_region(
        body,
        print_entry_block_args=False,
        print_empty_block=False,
        print_block_terminators=bool(iter_args),
    )

def parse_for_op_like(
    parser: Parser,
    default_indvar_type: TypeAttribute | None = None,
    bound_words: Sequence[str] = ["to"],
    *,
    allow_static_stop: bool = False,
    allow_static_step: bool = False,
) -> tuple[
    SSAValue, IntegerAttr | SSAValue, IntegerAttr | SSAValue, Sequence[SSAValue], Region
]:
    """
    Returns the initial bound, termination bound, step, iteration arguments, and body.
    Bounds are returned in textual order.

    Users can provide a default induction variable type and specific human-readable
    words for bounds (default: "to").
    Note that providing a default induction variable type is required to suggest that
    all loop control variable types (induction, bounds and step) have the same type,
    hence the induction variable type is potentially expected at the end of the for
    expression.

    When `allow_static_stop=True`, the termination bound may be either a static typed
    integer literal or a dynamic SSA value.
    When `allow_static_step=True`, the step may be either a static typed integer literal
    or a dynamic SSA value.
    The default (`False`) only allows dynamic SSA values.
    """

    unresolved_indvar = parser.parse_argument(expect_type=False)

    indvar_type = None

    if default_indvar_type is None:
        parser.parse_characters(":")
        indvar_type = parser.parse_type()

    parser.parse_characters("=")
    start = parser.parse_operand()

    for word in bound_words:
        parser.parse_characters(word)

    stop: IntegerAttr | SSAValue
    if allow_static_stop:
        if (stop_ssa := parser.parse_optional_operand()) is not None:
            stop = stop_ssa
        else:
            pos = parser.pos
            stop_attr = parser.parse_attribute()
            if not isa(stop_attr, IntegerAttr):
                parser.raise_error("Expected IntegerAttr", pos)
            stop = stop_attr
    else:
        stop = parser.parse_operand()
    parser.parse_characters("step")

    step: IntegerAttr | SSAValue
    if allow_static_step:
        if (step_ssa := parser.parse_optional_operand()) is not None:
            step = step_ssa
        else:
            pos = parser.pos
            step_attr = parser.parse_attribute()
            if not isa(step_attr, IntegerAttr):
                parser.raise_error("Expected IntegerAttr", pos)
            step = step_attr
    else:
        step = parser.parse_operand()

    # parse iteration arguments
    pos = parser.pos
    unresolved_iter_args: list[Parser.UnresolvedArgument] = []
    iter_arg_unresolved_operands: list[UnresolvedOperand] = []
    iter_arg_types: list[Attribute] = []
    if parser.parse_optional_characters("iter_args"):
        for iter_arg, iter_arg_operand in parser.parse_comma_separated_list(
            Parser.Delimiter.PAREN, lambda: parse_assignment(parser)
        ):
            unresolved_iter_args.append(iter_arg)
            iter_arg_unresolved_operands.append(iter_arg_operand)
        parser.parse_characters("->")
        iter_arg_types = parser.parse_comma_separated_list(
            Parser.Delimiter.PAREN, parser.parse_attribute
        )

    iter_arg_operands = parser.resolve_operands(
        iter_arg_unresolved_operands, iter_arg_types, pos
    )

    # set block argument types
    iter_args = [
        u_arg.resolve(t) for u_arg, t in zip(unresolved_iter_args, iter_arg_types)
    ]

    if default_indvar_type is not None:
        indvar_type = (
            parser.parse_type()
            if parser.parse_optional_characters(":")
            else default_indvar_type
        )
    assert indvar_type is not None

    # set induction variable type
    indvar = unresolved_indvar.resolve(indvar_type)

    body = parser.parse_region((indvar, *iter_args))

    return start, stop, step, iter_arg_operands, body

def print_func_op_like(
    printer: Printer,
    sym_name: StringAttr,
    function_type: FunctionType,
    body: Region,
    attributes: dict[str, Attribute],
    *,
    arg_attrs: ArrayAttr[DictionaryAttr] | None = None,
    res_attrs: ArrayAttr[DictionaryAttr] | None = None,
    reserved_attr_names: Sequence[str],
    is_variadic: bool = False,
    print_empty_outputs: bool = True,
):
    printer.print_string(" ")
    printer.print_symbol_name(sym_name.data)

    # Non-variadic declaration
    if not body.blocks and not is_variadic:
        if print_empty_outputs:
            printer.print_attribute(function_type)
        else:
            printer.print_string("(")
            printer.print_list(function_type.inputs, printer.print_attribute)
            printer.print_string(")")
            _print_func_outputs(printer, function_type.outputs.data, res_attrs)
        printer.print_op_attributes(
            attributes, reserved_attr_names=reserved_attr_names, print_keyword=True
        )
        return

    # Definition or variadic declaration
    printer.print_string("(")
    if body.blocks:
        block_args = body.blocks[0].args
        if arg_attrs is not None:
            printer.print_list(
                zip(block_args, arg_attrs),
                lambda t: print_func_argument(printer, t[0], t[1]),
            )
        else:
            printer.print_list(block_args, printer.print_block_argument)
        has_args = bool(block_args)
    else:
        printer.print_list(function_type.inputs, printer.print_attribute)
        has_args = bool(function_type.inputs)

    if is_variadic:
        if has_args:
            printer.print_string(", ")
        printer.print_string("...")
    printer.print_string(")")

    _print_func_outputs(printer, function_type.outputs.data, res_attrs)
    printer.print_op_attributes(
        attributes, reserved_attr_names=reserved_attr_names, print_keyword=True
    )

    if body.blocks:
        printer.print_string(" ", indent=0)
        printer.print_region(body, False, False)

def parse_func_op_like(
    parser: Parser,
    *,
    reserved_attr_names: Sequence[str],
    allow_variadic: bool = False,
) -> FuncOpLikeParseResult | FuncOpLikeParseResultWithVariadic:
    """
    Returns the function name, argument types, return types, body, extra args, arg_attrs and res_attrs.
    If allow_variadic=True, also returns is_variadic as the 8th element.
    """
    # Parse function name
    name = parser.parse_symbol_name().data

    # Track variadic state if enabled
    is_variadic = False

    def parse_fun_input() -> (
        Attribute | tuple[Parser.Argument, dict[str, Attribute]] | None
    ):
        def parse_optional_attrs_and_loc() -> tuple[
            dict[str, Attribute], LocationAttr | None
        ]:
            arg_attr_dict = parser.parse_optional_dictionary_attr_dict()
            arg_loc = parser.parse_optional_location()

            # Reject attrs after location, including empty dictionaries.
            if arg_loc is not None:
                arg_attr_pos = parser.pos
                parser.parse_optional_dictionary_attr_dict()
                if parser.pos != arg_attr_pos:
                    parser.raise_error(
                        "Expected function argument attributes before location."
                    )
            return arg_attr_dict, arg_loc

        nonlocal is_variadic
        if allow_variadic and parser.parse_optional_characters("...") is not None:
            is_variadic = True
            return None
        arg = parser.parse_optional_argument()
        if arg is None:
            ret = parser.parse_optional_type()
            if ret is None:
                parser.raise_error("Expected argument or type")
            # Declarative args keep only the type and consume attributes and location.
            parse_optional_attrs_and_loc()
        else:
            arg_attr_dict, arg_loc = parse_optional_attrs_and_loc()
            arg.location = arg_loc
            ret = (arg, arg_attr_dict)
        return ret

    # Parse function arguments
    args_raw = parser.parse_comma_separated_list(
        parser.Delimiter.PAREN,
        parse_fun_input,
    )
    args: list[Attribute | tuple[Parser.Argument, dict[str, Attribute]]]
    args = [arg for arg in args_raw if arg is not None]

    entry_arg_tuples: list[tuple[Parser.Argument, dict[str, Attribute]]] = []
    input_types: list[Attribute] = []
    for arg in args:
        if isinstance(arg, Attribute):
            input_types.append(arg)
        else:
            entry_arg_tuples.append(arg)

    if entry_arg_tuples:
        # Check consistency (They should be either all named or none)
        if input_types:
            parser.raise_error(
                "Expected all arguments to be named or all arguments to be unnamed."
            )

        entry_args = [arg for arg, _ in entry_arg_tuples]
        input_types = [arg.type for arg in entry_args]
    else:
        entry_args = None

    if any(attrs for _, attrs in entry_arg_tuples):
        arg_attrs = ArrayAttr(DictionaryAttr(attrs) for _, attrs in entry_arg_tuples)
    else:
        arg_attrs = None

    # Parse return types
    return_types, res_attrs = _parse_func_outputs(parser)

    extra_attributes = parser.parse_optional_attr_dict_with_keyword(reserved_attr_names)

    # Parse body
    region = parser.parse_optional_region(entry_args)
    if region is None:
        region = Region()

    if allow_variadic:
        return (
            name,
            input_types,
            return_types,
            region,
            extra_attributes,
            arg_attrs,
            res_attrs,
            is_variadic,
        )
    return (
        name,
        input_types,
        return_types,
        region,
        extra_attributes,
        arg_attrs,
        res_attrs,
    )

def print_func_argument(
    printer: Printer, arg: BlockArgument, attrs: DictionaryAttr | None
):
    """
    Keep function-argument syntax compatible with MLIR parser expectations:
    `%arg : type {attrs} loc(...)` (location after attrs).
    """
    printer.print_block_argument(arg, print_type=False)
    printer.print_string(": ")
    printer.print_attribute(arg.type)
    if attrs is not None and attrs.data:
        printer.print_op_attributes(attrs.data)
    if printer.print_debuginfo:
        printer.print_string(" ")
        printer.print_attribute(arg.location)

def print_func_output(
    printer: Printer, out_type: Attribute, attrs: DictionaryAttr | None
):
    printer.print_attribute(out_type)
    if attrs is not None and attrs.data:
        printer.print_op_attributes(attrs.data)

def print_assignment(printer: Printer, arg: BlockArgument, val: SSAValue):
    printer.print_block_argument(arg, print_type=False)
    printer.print_string(" = ")
    printer.print_ssa_value(val)

def parse_assignment(
    parser: Parser,
) -> tuple[Parser.UnresolvedArgument, UnresolvedOperand]:
    arg = parser.parse_argument(expect_type=False)
    parser.parse_characters("=")
    val = parser.parse_unresolved_operand()
    return arg, val