from __future__ import annotations

import cfg_build
from cfg.CFG import CFG
from cfg.CFG_Node import CFG_Node
from cfa.ReachedUses import ReachedUses
from syntax import EXPRESSION, ASSIGN
from optimizations.Optimization import Optimization

class DeadAssignmentElimination(Optimization):
    def __init__(self) -> None:
        self._node_by_ast_id: dict[int, CFG_Node] = {}
        self.__ru: ReachedUses | None = None

    def setup(self, ast: EXPRESSION) -> None:
        cfg_build.FUNCTIONS.clear()
        cfg_build.CURRENT_FUNCTION = None
        cfg = CFG(ast)
        self.__ru = ReachedUses(cfg)
        self._node_by_ast_id = {id(n.ast_node): n for n in self.__ru.cfg.nodes() if n.ast_node is not None}

    def apply(self, node: EXPRESSION) -> EXPRESSION:
        # Only ASSIGN nodes can be dead assignments.
        if not isinstance(node, ASSIGN):
            return node

        # Find the CFG node corresponding to this AST node.
        # If none is found, the node was not analyzed — leave it unchanged.
        cfg_node = self._node_by_ast_id.get(id(node))
        if cfg_node is None:
            return node

        ru = self.__ru
        nid = cfg_node.id
        defined_vars = ru.defs.get(nid, set())
        out = ru.out_sets.get(nid, set())

        # Check if any use-fact in OUT(n) belongs to a variable defined at n.
        # If so, the assignment is live — at least one use is reachable from here.
        for (_, var) in out:
            if var in defined_vars:
                return node

        # No reached use found for the assigned variable — the assignment is dead.
        # Replace ASSIGN(var, expr) with expr to preserve side effects on the RHS.
        return node.expr