from __future__ import annotations
from typing import TYPE_CHECKING
from cfa.BackwardAnalysis import BackwardAnalysis, Var

if TYPE_CHECKING:
    from cfg.CFG import CFG

class LiveVariables(BackwardAnalysis):
    def __init__(self, cfg: "CFG") -> None:
        # Base populates uses, defs, _func_scope, etc.
        super().__init__(cfg)

        self.gen: dict[int, set[Var]] = {}
        self.kill: dict[int, set[Var]] = {}
        self.incoming:  dict[int, set[Var]] = {}
        self.outgoing: dict[int, set[Var]] = {}

        self.__init_sets()
        self.solve()

    # Initialize gen, kill, in, and out sets for all CFG nodes.
    def __init_sets(self) -> None:
        for node in self.cfg.nodes():
            nid = node.id

            # GEN(n) = USE(n); KILL(n) = DEF(n)
            self.gen[nid] = set(self.uses[nid])
            self.kill[nid] = set(self.defs[nid])

            # IN(n) = GEN(n) = USE(n); OUT(n) = empty
            self.incoming[nid] = set(self.gen[nid])
            self.outgoing[nid] = set()

    # Update the lists until the fixpoint.
    def solve(self) -> None:
        nodes = list(self.cfg.nodes())
        known: set[int] = set(n.id for n in nodes)

        # while there are changes do
        changes = True
        while changes:
            changes = False

            # for all v IN V do
            for node in nodes:
                nid = node.id

                # OUT(n) = UNION IN(s) for all successors s
                new_out: set[Var] = set()
                for child in node.children:
                    if child.id in known:
                        new_out |= self.incoming[child.id]

                # IN(n) = (OUT(n) MINUS KILL(n)) UNION GEN(n)
                new_in: set[Var] = (new_out - self.kill[nid]) | self.gen[nid]

                if new_out != self.outgoing[nid] or new_in != self.incoming[nid]:
                    self.outgoing[nid] = new_out
                    self.incoming[nid] = new_in
                    changes = True  # there are changes -> loop again

    # Return the living variables within each node
    def live_vars_by_node(self) -> dict[int, set[Var]]:
        return {nid: set(vs) for nid, vs in self.incoming.items() if vs}