Format

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

    The induction variable and each SSA operand are printed with their own types.
    The `stop` and `step` may also be typed integer literals.
    Users can provide specific words for bounds (default: "to").
    """

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

    printer.print_string(" ")
    printer.print_ssa_value(indvar)
    printer.print_string(" : ")
    printer.print_attribute(indvar.type)
    printer.print_string(" = ")
    _print_bound(printer, start)

    printer.print_string(f" {' '.join(bound_words)} ")

    _print_bound(printer, stop)
    printer.print_string(" step ")
    _print_bound(printer, 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(") ")

    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,
    bound_words: Sequence[str] = ["to"],
    *,
    allow_static_stop: bool = False,
    allow_static_step: bool = False,
) -> tuple[
    SSAValue, IntegerAttr | SSAValue, IntegerAttr | SSAValue, Sequence[SSAValue], Region
]:
    """
    Parses backend loop bounds, step, iteration arguments, and body.
    Bounds are returned in textual order. The induction variable and each SSA
    operand must be followed by their types.
    Users can provide specific words for bounds (default: "to").

    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.
    """

    indvar = parser.parse_argument()
    parser.parse_characters("=")
    start = _parse_bound(parser)

    for word in bound_words:
        parser.parse_characters(word)

    stop = _parse_bound(parser, allow_static=allow_static_stop)
    parser.parse_characters("step")
    step = _parse_bound(parser, allow_static=allow_static_step)

    # 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)
    ]

    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)

    # 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