Lowering

@dataclass
class PyLowering:
    """
    Lowers the `py` operations generated from a Python function to the operations
    registered for them.

    Operations are lowered in order, so the types of the operands of an operation are
    known when it is lowered: every value gets the type of the operation it becomes.
    """

    type_converter: TypeConverter
    """The registries, and the globals in which names in the program are resolved."""

    file: str | None
    """Path of the file containing the program being processed."""

    function_definition_constructor: FunctionDefinitionConstructor = build_func
    """Construct the operation for each function."""

    return_constructor: ReturnConstructor = build_return
    """Construct return operations, optionally omitting the terminator."""

    variables: dict[str, Attribute] = field(
        default_factory=dict[str, Attribute], init=False
    )
    """The types of the variables of the function being lowered."""

    packs: dict[SSAValue, tuple[tuple[SSAValue, ...], Operation]] = field(
        default_factory=dict[SSAValue, tuple[tuple[SSAValue, ...], Operation]],
        init=False,
    )
    """
    The `py` values that stand for a number of values other than one, such as the
    results of a call or the elements of a tuple, with the operation that produced
    them. These values can only be unpacked.
    """

    def lower(self, op: Operation) -> None:
        """Lower all `py` operations nested in `op`."""
        self._lower_regions(op)
        for _, producer in self.packs.values():
            Rewriter.erase_op(producer)
        self.packs.clear()

    def _lower_regions(self, op: Operation) -> None:
        for region in op.regions:
            for block in region.blocks:
                self._lower_block(block)

    def _lower_block(self, block: Block) -> None:
        for op in list(block.ops):
            self._lower_op(op)

    def _lower_op(self, op: Operation) -> None:
        if not isinstance(op, py.PyUnpackOp):
            for operand in op.operands:
                if operand in self.packs:
                    producer = self.packs[operand][1]
                    raise self._error(
                        producer,
                        "Tuples are only supported when unpacked."
                        if isinstance(producer, py.PyBuildTupleOp)
                        else "Expected an expression with exactly one result.",
                    )

        match op:
            case py.PyFuncOp():
                self._lower_func(op)
            case py.PyReturnOp():
                self._lower_return(op)
            case py.PyConstOp():
                self._lower_const(op)
            case py.PyCastOp():
                # Casts only change annotations, which are resolved by now.
                Rewriter.replace_op(op, (), (op.value,))
            case py.PyBinOp():
                self._lower_binop(op)
            case py.PyCompareOp():
                self._lower_compare(op)
            case py.PyCallOp():
                self._lower_call(op)
            case py.PyBuildTupleOp():
                # A tuple stands for its elements, until it is unpacked.
                self.packs[op.res] = (tuple(op.elements), op)
            case py.PyUnpackOp():
                self._lower_unpack(op)
            case symref.FetchOp():
                name = op.symbol.root_reference.data
                Rewriter.replace_value_with_new_type(op.value, self.variables[name])
            case symref.UpdateOp():
                self._lower_update(op)
            case _:
                self._lower_regions(op)

    def _lower_func(self, op: py.PyFuncOp) -> None:
        block = op.body.block
        arg_types = [
            self._resolve_annotation(arg, arg_type, "argument")
            for arg, arg_type in zip(
                block.args, op.function_type.inputs.data, strict=True
            )
        ]
        (result_type,) = op.function_type.outputs.data
        return_types = (
            []
            if self._annotation(result_type) == "None"
            else [self._resolve_annotation(op, result_type, "return")]
        )

        for arg, arg_type in zip(block.args, arg_types, strict=True):
            location = arg.location
            Rewriter.replace_value_with_new_type(arg, arg_type)
            block.args[arg.index].location = location
        self.variables = {}
        self._lower_block(block)

        function_op = self.function_definition_constructor(
            op.func_name.data,
            FunctionType.from_lists(arg_types, return_types),
            Rewriter.move_region_contents_to_new_regions(op.body),
            self._location(op),
        )
        Rewriter.replace_op(op, function_op, ())

    def _lower_return(self, op: py.PyReturnOp) -> None:
        value = op.value
        none = value.owner
        if not (isinstance(none, py.PyConstOp) and none.value is None):
            none = None

        # Returning None returns no values.
        values = () if none else (value,)
        return_op = self.return_constructor(values, self._location(op))
        if return_op is None:
            Rewriter.erase_op(op)
        else:
            Rewriter.replace_op(op, return_op, ())
        if none and not none.res.uses:
            Rewriter.erase_op(none)

    def _lower_const(self, op: py.PyConstOp) -> None:
        if self._is_returned_none(op):
            # Returning None returns no values, so the return lowers it.
            return
        value = op.value
        literal_op = self.type_converter.literal_registry.resolve_operation(value)
        if literal_op is None:
            raise self._error(
                op,
                f"Unsupported constant '{value}' of type '{type(value).__qualname__}'.",
            )
        self._replace(op, literal_op)

    @staticmethod
    def _is_returned_none(op: py.PyConstOp) -> bool:
        """Whether `op` is a `None` that is only returned, through a cast."""
        if op.value is not None or not op.res.uses:
            return False
        return all(
            isinstance(cast_op := use.operation, py.PyCastOp)
            and all(isinstance(u.operation, py.PyReturnOp) for u in cast_op.res.uses)
            for use in op.res.uses
        )

    def _lower_binop(self, op: py.PyBinOp) -> None:
        op_name = _BINARY_OPERATORS[op.op.data]
        method_op = resolve_binary_operator(
            self.type_converter, op_name, op.lhs, op.rhs, self._error_at(op)
        )
        self._replace(op, method_op)

    def _lower_compare(self, op: py.PyCompareOp) -> None:
        op_name = _COMPARISON_OPERATORS[op.op.data]
        if op_name in UNSUPPORTED_COMPARISON_OPERATORS:
            raise self._error(op, f"Unsupported comparison operation '{op_name}'.")
        method_op = resolve_comparison_operator(
            self.type_converter, op_name, op.lhs, op.rhs, self._error_at(op)
        )
        self._replace(op, method_op)

    def _lower_call(self, op: py.PyCallOp) -> None:
        ir_op = resolve_callee(self.type_converter, op.callee.data, self._error_at(op))
        keywords = op.keyword_names
        num_positional = len(op.args) - len(keywords)
        args = op.args[:num_positional]
        kwargs = dict(zip(keywords, op.args[num_positional:], strict=True))
        self._replace(op, ir_op(*args, **kwargs))

    def _lower_unpack(self, op: py.PyUnpackOp) -> None:
        if op.value in self.packs:
            values, _ = self.packs[op.value]
        else:
            values = (op.value,)
        if len(values) != len(op.res):
            raise self._error(
                op, f"Expected {len(op.res)} values to unpack, got {len(values)}."
            )
        Rewriter.replace_op(op, (), values)

    def _lower_update(self, op: symref.UpdateOp) -> None:
        name = op.symbol.root_reference.data
        if name not in self.variables:
            self.variables[name] = op.value.type
        elif self.variables[name] != op.value.type:
            raise self._error(op, f"Cannot change the type of variable '{name}'.")

    def _replace(self, op: Operation, new_op: Operation) -> None:
        """Replace a `py` operation by `new_op`, which may have any number of results."""
        new_op.location = op.location
        if len(new_op.results) == 1:
            Rewriter.replace_op(op, new_op)
        else:
            Rewriter.insert_op(new_op, InsertPoint.before(op))
            self.packs[op.results[0]] = (tuple(new_op.results), op)

    @staticmethod
    def _annotation(object_type: Attribute) -> str:
        return cast(py.PyObjectType, object_type).annotation.data

    def _resolve_annotation(
        self, ir: Operation | BlockArgument, object_type: Attribute, kind: str
    ) -> Attribute:
        annotation = self._annotation(object_type)
        ir_type = self.type_converter.type_registry.resolve_attribute(
            annotation, self.type_converter.globals
        )
        if ir_type is None:
            raise self._error(ir, f"Unsupported function {kind} type: '{annotation}'")
        return ir_type

    @staticmethod
    def _position(ir: Operation | BlockArgument) -> tuple[int, int]:
        """The line and one-based column of an operation or argument, or zeros."""
        if isinstance(location := ir.location, FileLineColLoc):
            return location.line.data, location.column.data
        return 0, 0

    def _location(self, op: Operation) -> Location:
        return Location(self.file or "<unknown>", *self._position(op))

    def _error_at(self, op: Operation) -> ErrorConstructor:
        """Construct errors located at `op`."""
        return lambda msg: self._error(op, msg)

    def _error(
        self, ir: Operation | BlockArgument, msg: str
    ) -> CodeGenerationException:
        line, column = self._position(ir)
        # Errors report the zero-based column of the Python AST.
        return CodeGenerationException(self.file, line, max(column - 1, 0), msg)