"""Simple copy propagation optimization.

Example input:

    x = f()
    y = x

The register x is redundant and we can directly assign its value to y:

    y = f()

This can optimize away registers that are assigned to once.
"""

from __future__ import annotations

from mypyc.ir.func_ir import FuncIR
from mypyc.ir.ops import Assign, AssignMulti, LoadAddress, LoadErrorValue, Register, Value
from mypyc.irbuild.ll_builder import LowLevelIRBuilder
from mypyc.options import CompilerOptions
from mypyc.sametype import is_same_type
from mypyc.transform.ir_transform import IRTransform


def do_copy_propagation(fn: FuncIR, options: CompilerOptions) -> None:
    """Perform copy propagation optimization for fn."""

    # Anything with an assignment count >1 will not be optimized
    # here, as it would be require data flow analysis and we want to
    # keep this simple and fast, at least until we've made data flow
    # analysis much faster.
    counts: dict[Value, int] = {}
    replacements: dict[Value, Value] = {}
    for arg in fn.arg_regs:
        # Arguments are always assigned to initially
        counts[arg] = 1

    for block in fn.blocks:
        for op in block.ops:
            if isinstance(op, Assign):
                c = counts.get(op.dest, 0)
                counts[op.dest] = c + 1
                # Does this look like a supported assignment?
                # TODO: Something needs LoadErrorValue assignments to be preserved?
                if (
                    c == 0
                    and is_same_type(op.dest.type, op.src.type)
                    and not isinstance(op.src, LoadErrorValue)
                ):
                    replacements[op.dest] = op.src
                elif c == 1:
                    # Too many assignments -- don't replace this one
                    replacements.pop(op.dest, 0)
            elif isinstance(op, AssignMulti):
                # Copy propagation not supported for AssignMulti destinations
                counts[op.dest] = 2
                replacements.pop(op.dest, 0)
            elif isinstance(op, LoadAddress):
                # We don't support taking the address of an arbitrary Value,
                # so we'll need to preserve the operands of LoadAddress.
                if isinstance(op.src, Register):
                    counts[op.src] = 2
                    replacements.pop(op.src, 0)

    # Follow chains of propagation with more than one assignment.
    for src, dst in list(replacements.items()):
        if counts.get(dst, 0) > 1:
            # Not supported
            del replacements[src]
        else:
            while dst in replacements:
                dst = replacements[dst]
                if counts.get(dst, 0) > 1:
                    # Not supported
                    del replacements[src]
        if src in replacements:
            replacements[src] = dst

    builder = LowLevelIRBuilder(None, options)
    transform = CopyPropagationTransform(builder, replacements)
    transform.transform_blocks(fn.blocks)
    fn.blocks = builder.blocks


class CopyPropagationTransform(IRTransform):
    def __init__(self, builder: LowLevelIRBuilder, map: dict[Value, Value]) -> None:
        super().__init__(builder)
        self.op_map.update(map)
        self.removed = set(map)

    def visit_assign(self, op: Assign) -> Value | None:
        if op.dest in self.removed:
            return None
        return self.add(op)