"""
transpiler/ssa.py — SSA Construction and Dominance Analysis.
"""

from __future__ import annotations
from typing import Dict, List, Set, Optional, Any, Union, cast
import copy
from transpiler.middleend import cfg, ir
from transpiler.optimization import loop_detector
from transpiler.infrastructure import visitor, logger


class DominatorAnalysis:
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.idoms: Dict[int, Optional[int]] = {} # block_id -> idom_id
        self.dominators: Dict[int, Set[int]] = {} # block_id -> set of dominator_ids
        self.frontiers: Dict[int, Set[int]] = {} # block_id -> set of block_ids
        self.dom_tree: Dict[int, List[int]] = {}

    def run(self):
        self._compute_dominators()
        self._compute_idoms()
        self._compute_frontiers()
        self._build_dom_tree()

    def _compute_dominators(self):
        # Entry node
        entry_id = self.cfg.entry_block.id
        all_blocks = set(self.cfg.blocks.keys())
        
        self.dominators = {bid: all_blocks.copy() for bid in all_blocks}
        self.dominators[entry_id] = {entry_id}
        
        changed = True
        while changed:
            changed = False
            for bid in all_blocks:
                if bid == entry_id: continue
                
                block = self.cfg.blocks[bid]
                preds_list = [p for p in block.predecessors if entry_id in self.dominators.get(p.id, set())]
                
                if len(preds_list) == 0:
                    # If it has predecessors but none are reachable yet, 
                    # we don't update yet (stay at all_blocks)
                    if block.predecessors:
                        continue
                    preds_doms = set()
                else:
                    preds_doms = self.dominators[preds_list[0].id].copy()
                    for i in range(1, len(preds_list)):
                        p = preds_list[i]
                        preds_doms.intersection_update(self.dominators[p.id])
                
                new_doms = {bid} | preds_doms
                if new_doms != self.dominators[bid]:
                    self.dominators[bid] = new_doms
                    changed = True
        

    def _compute_idoms(self):
        entry_id = self.cfg.entry_block.id
        for bid, doms in self.dominators.items():
            if bid == entry_id:
                self.idoms[bid] = None
                continue
            
            # If the entry block does not dominate this block, it's unreachable.
            if entry_id not in doms:
                self.idoms[bid] = None
                continue
            
            # A strict dominator d of n is idom(n) if d dominates all other strict dominators of n.
            strict_doms = doms - {bid}
            
            # Use the size-based heuristic which is correct for simple connected graphs
            idom = None
            max_dom_count = -1
            for d in strict_doms:
                if len(self.dominators[d]) > max_dom_count:
                    max_dom_count = len(self.dominators[d])
                    idom = d
            if idom is not None:
                self.idoms[bid] = idom

    def _compute_frontiers(self):
        # DF(n) is the set of all nodes w such that n dominates a predecessor of w, but n does not strictly dominate w.
        for bid in self.cfg.blocks:
            self.frontiers[bid] = set()
            
        for bid, block in self.cfg.blocks.items():
            if len(block.predecessors) >= 2:
                for p in block.predecessors:
                    runner = p.id
                    target_idom = self.idoms.get(bid)
                    seen = set()
                    while target_idom is not None and runner != target_idom and runner is not None:
                        if runner in seen:
                            break
                        seen.add(runner)
                        self.frontiers.get(runner, set()).add(bid)
                        runner = self.idoms.get(runner)

    def _build_dom_tree(self):
        self.dom_tree = {bid: [] for bid in self.cfg.blocks}
        for bid, idom in self.idoms.items():
            if idom is not None:
                self.dom_tree[idom].append(bid)

class SSAConstructor:
    def __init__(self, target_cfg: cfg.CFG, dom_analysis: DominatorAnalysis, initial_vars: Optional[List[str]] = None):
        self.cfg = target_cfg
        self.dom = dom_analysis
        self.var_versions: Dict[str, int] = {} # var_name -> latest_version
        self.var_stacks: Dict[str, List[str]] = {} # var_name -> stack of versioned_names
        
        if initial_vars:
            for v in initial_vars:
                # Seed arguments as version 1
                new_name = f"{v}_1"
                self.var_versions[v] = 1
                self.var_stacks[v] = [new_name]

        self.phi_placed: Dict[str, Set[int]] = {} # var_name -> set of block_ids
        self.dom_children: Dict[int, List[int]] = self.dom.dom_tree # Initialize from DominatorAnalysis

    def run(self):
        self.construct()

    def construct(self):
        self._place_phis()
        # self.dom_children is already initialized in __init__ from self.dom.dom_tree
        self._rename_vars(self.cfg.entry_block.id)

    def _place_phis(self):
        # 1. Collect all variable assignments
        defs: Dict[str, Set[int]] = {}
        for bid, block in self.cfg.blocks.items():
            for stmt in block.statements:
                # logger.debug(f"Block {bid} statement: {type(stmt)} -> {stmt}")
                if isinstance(stmt, ir.IRAssignment):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    target = stmt_assign.target
                    if isinstance(target, ir.IRVariable):
                        defs.setdefault(target.name, set()).add(bid)
                        if not target.original_name:
                            target.original_name = target.name
                elif isinstance(stmt, ir.IRMultiAssign):
                    for target in stmt.targets:
                        if isinstance(target, ir.IRVariable):
                            defs.setdefault(target.name, set()).add(bid)
                            if not target.original_name:
                                target.original_name = target.name

        # 2. Iteratively place Phis
        for var_name, def_blocks in defs.items():
            # if "found_at" in var_name or "inv_idx" in var_name:
            #     logger.debug(f"Placing Phis for {var_name}, defs in {def_blocks}")
            #     for bid in sorted(list(def_blocks)):
            #         logger.debug(f"  DF({bid}) = {self.dom.frontiers.get(bid, set())}")

            w = list(def_blocks)
            while w:
                x = w.pop()
                for y in self.dom.frontiers[x]:
                    if y not in self.phi_placed.get(var_name, set()):
                        # if "found_at" in var_name or "inv_idx" in var_name:
                        #     print(f"DEBUG: Placing Phi for {var_name} in Block {y}")
                        # Place Phi in block Y for var_name
                        phi = ir.IRPhi(options={})
                        phi_stmt = ir.IRAssignment(
                            target=ir.IRVariable(name=var_name, original_name=var_name), 
                            value=phi
                        )
                        # We know blocks exists and has y
                        self.cfg.blocks[y].statements.insert(0, phi_stmt)
                        
                        self.phi_placed.setdefault(var_name, set()).add(y)
                        if y not in def_blocks:
                            w.append(y)

    def _rename_vars(self, bid: int):
        block = self.cfg.blocks[bid]
        pushed = []

        # 1. Process Phis and regular assignments in this block
        for stmt in block.statements:
            self._rename_in_stmt(stmt, pushed)
        
        # 2. Fill in Phi parameters for successors
        for succ in block.successors:
            for stmt in succ.statements:
                if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.value, ir.IRPhi):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    phi_val = cast(ir.IRPhi, stmt_assign.value)
                    target_var = cast(ir.IRVariable, stmt_assign.target)
                    orig_name = target_var.original_name
                    if orig_name in self.var_stacks and self.var_stacks[orig_name]:
                        new_name = self.var_stacks[orig_name][-1]
                        phi_val.options[bid] = ir.IRVariable(name=new_name, original_name=orig_name)
                    else:
                        # Undefined: use nil or a special Uninitialized node
                        phi_val.options[bid] = ir.IRLiteral(value=None)

        # 3. Recursively visit children in Dominator Tree
        for child_id in self.dom_children.get(bid, []):
            self._rename_vars(child_id)
        
        # 4. Pop stacks
        for name in pushed:
            self.var_stacks[name].pop()

    def _gen_new_name(self, name: str) -> str:
        version = self.var_versions.get(name, 0) + 1
        self.var_versions[name] = version
        return f"{name}_{version}"

    def _rename_in_expr(self, expr: ir.IRExpression):
        # We handle this by mutating the original node IF it's an IRVariable,
        # but the linearizer ensures they aren't shared across different statements/uses.
        if isinstance(expr, ir.IRVariable):
            # Ensure it HAS an original name before we lose it
            if not expr.original_name:
                expr.original_name = expr.name
            
            orig = expr.original_name
            stack = self.var_stacks.get(orig)
            if stack:
                expr.name = stack[-1]
                expr.original_name = orig
        elif hasattr(expr, "__dict__"):
            for attr, v in expr.__dict__.items():
                if isinstance(v, ir.IRExpression):
                    self._rename_in_expr(v)
                elif isinstance(v, list):
                    for item in v:
                        if isinstance(item, ir.IRExpression):
                            self._rename_in_expr(item)
                elif isinstance(v, ir.IRIfElseExpression):
                    self._rename_in_expr(v.condition)
                    self._rename_in_expr(v.then_expr)
                    self._rename_in_expr(v.else_expr)

    def _rename_in_stmt(self, stmt: ir.IRStatement, pushed: List[str]):
        """Rename variables in a statement, potentially pushing new versions to stack."""
        if isinstance(stmt, ir.IRAssignment):
            stmt_assign = cast(ir.IRAssignment, stmt)
            # 1. Rename RHS
            if not isinstance(stmt_assign.value, ir.IRPhi):
                self._rename_in_expr(stmt_assign.value)
            
            # 2. Rename LHS (target)
            target = stmt_assign.target
            if isinstance(target, ir.IRVariable):
                target_var = cast(ir.IRVariable, target)
                old_name = target_var.original_name or target_var.name
                new_name = self._gen_new_name(old_name)
                stmt_assign.target = ir.IRVariable(
                    name=new_name,
                    original_name=old_name,
                    inferred_type=target_var.inferred_type,
                    lineno=target_var.lineno,
                    col_offset=target_var.col_offset
                )
                self.var_stacks.setdefault(old_name, []).append(new_name)
                pushed.append(old_name)
        elif isinstance(stmt, ir.IRMultiAssign):
            # 1. Rename RHS
            if stmt.rhs_call:
                self._rename_in_expr(stmt.rhs_call)
            for val in stmt.rhs_values:
                self._rename_in_expr(val)
            
            # 2. Rename LHS targets
            for i, target in enumerate(stmt.targets):
                if isinstance(target, ir.IRVariable):
                    old_name = target.original_name or target.name
                    new_name = self._gen_new_name(old_name)
                    stmt.targets[i] = ir.IRVariable(
                        name=new_name,
                        original_name=old_name,
                        inferred_type=target.inferred_type,
                        lineno=target.lineno,
                        col_offset=target.col_offset
                    )
                    self.var_stacks.setdefault(old_name, []).append(new_name)
                    pushed.append(old_name)
        else:
            if isinstance(stmt, ir.IRFunctionDef):
                # SSA Constructor must NOT recurse into function bodies during 
                # module-level optimization. Function bodies are handled in their own
                # isolated SSA pass.
                return
            
            # Recursively handle any nested statement lists (body, then_block, else_block, etc.)
            new_pushed = [] # Inner scope (legacy/unused for functions now)
            
            if hasattr(stmt, "__dict__"):
                for k, v in stmt.__dict__.items():
                    if isinstance(v, list):
                        for item in v:
                            if isinstance(item, ir.IRStatement):
                                self._rename_in_stmt(item, pushed if not isinstance(stmt, ir.IRFunctionDef) else new_pushed)
                            elif isinstance(item, ir.IRExpression):
                                self._rename_in_expr(item)
                    elif isinstance(v, ir.IRExpression):
                        self._rename_in_expr(v)
            
            # Pop inner scope if it was a function
            for name in new_pushed:
                self.var_stacks[name].pop()

    def _rename_in_stmt_generic(self, stmt: ir.IRStatement):
        # Deprecated: use _rename_in_stmt instead which handles recursion better
        pass

class SSARedundantTempElimination:
    """Eliminates assignments like 'local x = y' where x is a single-use temporary."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg

    def run(self) -> bool:
        changed = False
        # Build use counts and definition map
        use_counts: Dict[str, int] = {}
        def_counts: Dict[str, int] = {}
        defs: Dict[str, ir.IRAssignment] = {}
        
        for bid, block in self.cfg.blocks.items():
            for stmt in block.statements:
                if isinstance(stmt, ir.IRFunctionDef): continue
                if isinstance(stmt, ir.IRAssignment):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    if isinstance(stmt_assign.target, ir.IRVariable):
                        name = stmt_assign.target.name
                        def_counts[name] = def_counts.get(name, 0) + 1
                        defs[name] = stmt_assign
                self._collect_uses(stmt, use_counts)
        
        # Identify variables that are just copies: x = y
        aliases: Dict[str, ir.IRExpression] = {}
        for name, stmt in defs.items():
            # ONLY eliminate if defined exactly once
            if def_counts.get(name, 0) == 1 and isinstance(stmt.value, ir.IRVariable):
                # ONLY eliminate if it's a temporary (starts with _)
                if name.startswith("_"):
                    aliases[name] = stmt.value
        
        if aliases:
            for bid, block in self.cfg.blocks.items():
                new_stmts = []
                for stmt in block.statements:
                    # If this is the definition of a redundant temp, skip it
                    if isinstance(stmt, ir.IRAssignment):
                        stmt_assign = cast(ir.IRAssignment, stmt)
                        if isinstance(stmt_assign.target, ir.IRVariable) and stmt_assign.target.name in aliases:
                            changed = True
                            continue
                    # Replace uses
                    if self._replace_aliases(stmt, aliases):
                        changed = True
                    new_stmts.append(stmt)
                block.statements = new_stmts
                
        return changed

    def _collect_uses(self, node: Any, counts: Dict[str, int]):
        if isinstance(node, ir.IRVariable):
            counts[node.name] = counts.get(node.name, 0) + 1
        elif hasattr(node, "__dict__"):
            # DO NOT recurse into nested IR bodies (function/if/while)
            # We only want to collect uses in the expressions of the current block.
            for k, v in node.__dict__.items():
                if k in ("body", "then_block", "else_block", "elif_blocks"):
                    continue
                if k == "target" and isinstance(node, ir.IRAssignment):
                    if isinstance(v, ir.IRVariable):
                        continue
                if isinstance(v, (ir.IRNode, list, dict)):
                    self._collect_node_uses(v, counts)

    def _collect_node_uses(self, node: Any, counts: Dict[str, int]):
        if isinstance(node, list):
            for i in node: self._collect_node_uses(i, counts)
        elif isinstance(node, dict):
            for v in node.values(): self._collect_node_uses(v, counts)
        else:
            self._collect_uses(node, counts)

    def _replace_aliases(self, node: Any, aliases: Dict[str, ir.IRExpression]) -> bool:
        changed = False
        if hasattr(node, "__dict__"):
            for k, v in node.__dict__.items():
                if k == "target" and isinstance(node, ir.IRAssignment): continue
                if isinstance(v, ir.IRVariable) and v.name in aliases:
                    setattr(node, k, copy.deepcopy(aliases[v.name]))
                    changed = True
                elif isinstance(v, (ir.IRNode, list)):
                    if self._replace_node_aliases(v, aliases):
                        changed = True
        return changed

    def _replace_node_aliases(self, node: Any, aliases: Dict[str, ir.IRExpression]) -> bool:
        changed = False
        if isinstance(node, list):
            for i, item in enumerate(node):
                if isinstance(item, ir.IRVariable) and item.name in aliases:
                    node[i] = copy.deepcopy(aliases[item.name])
                    changed = True
                else:
                    if self._replace_node_aliases(item, aliases):
                        changed = True
        elif isinstance(node, ir.IRNode):
            if self._replace_aliases(node, aliases):
                changed = True
        return changed

class SSADCE:
    """Performs Dead Code Elimination on SSA IR."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.use_counts: Dict[str, int] = {}
        self.captured_names: Set[str] = set()

    def run(self) -> bool:
        changed = True
        overall_changed = False
        while changed:
            changed = False
            self.use_counts = {}
            self.captured_names = set()
            # 1. Count uses recursively
            for bid, block in self.cfg.blocks.items():
                for stmt in block.statements:
                    self._count_uses(stmt)
                if block.terminator:
                    self._count_uses(block.terminator)
            
            for bid, block in self.cfg.blocks.items():
                if self._remove_dead_recursive(block):
                    changed = True
            
            if changed: overall_changed = True
        return overall_changed

    def _remove_dead_recursive(self, container: Any) -> bool:
        changed = False
        stmts: Optional[List[ir.IRStatement]] = None
        if hasattr(container, "statements"): stmts = container.statements
        elif isinstance(container, list): stmts = container
        
        if stmts is not None:
             new_stmts = []
             for stmt in stmts:
                 if isinstance(stmt, ir.IRFunctionDef):
                     new_stmts.append(stmt)
                     continue
                 
                 is_dead = False
                 if isinstance(stmt, ir.IRAssignment):
                      stmt_assign = cast(ir.IRAssignment, stmt)
                      if isinstance(stmt_assign.target, ir.IRVariable):
                           target_var = cast(ir.IRVariable, stmt_assign.target)
                           count = self.use_counts.get(target_var.name, 0)
                           if count == 0:
                               # Module-level optimization: 
                               # Never delete assignments to user-defined variables (non-underscore) 
                               # as they might be used by function closures not visible in this CFG.
                               is_user_global = not target_var.name.startswith("_")
                               orig_name = getattr(target_var, "original_name", target_var.name) or target_var.name
                               is_captured = orig_name in self.captured_names
                               
                               # Eliminate if no side effects, AND (it's a temp OR it's a local user variable that isn't captured)
                               # Note: In module-level CFG, all user variables are considered module globals
                               # but right now optimize_ir processes each function's CFG separately, 
                               # so we can safely delete unused locals if they aren't captured.
                               
                               if not self._has_side_effects(stmt_assign.value) and (not is_user_global or not is_captured):
                                   is_dead = True
                                   changed = True
                                   # logger.debug(f"DCE REMOVING assignment to {target_var.name} (use count 0)")
                                   logger.explain("DCE", f"Dead assignment removed: {target_var.name}")
                               else:
                                   pass # logger.debug(f"DCE KEEPING {target_var.name} | is_user={is_user_global} is_captured={is_captured} side_eff={self._has_side_effects(stmt_assign.value)}")
                               # logger.debug(f"DCE KEEPING {target_var.name} (USED: {count})")
                 elif isinstance(stmt, ir.IRMultiAssign):
                       # Dead if ALL targets are unused AND no side effects
                       all_dead = True
                       for target in stmt.targets:
                           if isinstance(target, ir.IRVariable):
                               if self.use_counts.get(target.name, 0) > 0:
                                   all_dead = False
                                   break
                       
                       if all_dead:
                           has_side_effect = False
                           if stmt.rhs_call and self._has_side_effects(stmt.rhs_call):
                               has_side_effect = True
                           else:
                               for val in stmt.rhs_values:
                                   if self._has_side_effects(val):
                                       has_side_effect = True
                                       break
                           
                           if not has_side_effect:
                               is_dead = True
                               changed = True
                               logger.explain("DCE", "Dead multi-assignment removed")

                 
                 if not is_dead:
                      # DO NOT recurse into nested IR bodies here.
                      # When running on a CFG, structural statements (if/while) 
                      # should be treated as opaque. Their bodies were handled
                      # either by CFG edges or by an isolated optimization pass.
                      new_stmts.append(stmt)
                 else:
                      if isinstance(stmt, ir.IRAssignment):
                          stmt_assign = cast(ir.IRAssignment, stmt)
                          if isinstance(stmt_assign.target, ir.IRVariable):
                              target_var = cast(ir.IRVariable, stmt_assign.target)
                              logger.explain("DCE", f"Removing DEAD assignment to {target_var.name}")
                      elif isinstance(stmt, ir.IRMultiAssign):
                           logger.explain("DCE", "Removing DEAD multi-assignment")
             
             if isinstance(container, list): 
                 container.clear()
                 container.extend(new_stmts)
             else: container.statements = new_stmts
        return changed
        
    def _count_uses(self, stmt: ir.IRStatement):
        if isinstance(stmt, ir.IRFunctionDef):
            # To safely allow DCE to eliminate unused variables in the outer function,
            # we MUST track variables captured by closures.
            self._find_captured_uses(stmt)
            return
        
        # Recursively find all variable uses in the statement
        if isinstance(stmt, ir.IRMultiAssign):
            if stmt.rhs_call:
                self._find_uses(stmt.rhs_call)
            for val in stmt.rhs_values:
                self._find_uses(val)
            return

        if hasattr(stmt, "__dict__"):
            for k, v in stmt.__dict__.items():
                if isinstance(stmt, ir.IRAssignment) and k == "target":
                    if isinstance(v, ir.IRVariable):
                        continue
                if isinstance(v, ir.IRExpression):
                    self._find_uses(v)
                elif isinstance(v, list):
                    for item in v:
                        if isinstance(item, ir.IRExpression):
                            self._find_uses(item)
                        elif isinstance(item, ir.IRStatement):
                            self._count_uses(item) # Nested statements

    def _find_uses(self, expr: ir.IRExpression):
        if isinstance(expr, ir.IRVariable):
            self.use_counts[expr.name] = self.use_counts.get(expr.name, 0) + 1
        elif isinstance(expr, ir.IRPhi):
            phi_expr = cast(ir.IRPhi, expr)
            for opt in phi_expr.options.values():
                self._find_uses(opt)
        elif hasattr(expr, "__dict__"):
            for v in expr.__dict__.values():
                if isinstance(v, ir.IRExpression):
                    self._find_uses(v)
                elif isinstance(v, list):
                    for item in v:
                        if isinstance(item, ir.IRExpression):
                            self._find_uses(item)
                elif isinstance(v, ir.IRIfElseExpression):
                    self._find_uses(v.condition)
                    self._find_uses(v.then_expr)
                    self._find_uses(v.else_expr)

    def _find_captured_uses(self, node: Any):
        """Scans a nested function body to find uses of variables by their *original* names."""
        if isinstance(node, ir.IRVariable):
            name = getattr(node, "original_name", node.name) or node.name
            self.captured_names.add(name)
            # We artificially increment use_counts for ALL current SSA versions of this original_name
            # so they are not deleted if captured. But we might not know their exact SSA version here.
            # That's why we rely on `self.captured_names` instead.
        elif hasattr(node, "__dict__"):
            for k, v in node.__dict__.items():
                if k == "target" and isinstance(node, ir.IRAssignment):
                    continue # definitions inside the closure don't capture outer variables
                if isinstance(v, (ir.IRExpression, ir.IRStatement)):
                    self._find_captured_uses(v)
                elif isinstance(v, list):
                    for item in v:
                        if isinstance(item, (ir.IRExpression, ir.IRStatement)):
                            self._find_captured_uses(item)

    def _has_side_effects(self, expr: ir.IRExpression) -> bool:
        # If purity analysis has been run, use the lattice
        if expr.effect_level != ir.EffectLevel.UNKNOWN:
            return expr.effect_level in (ir.EffectLevel.WRITES_MEMORY, ir.EffectLevel.HAS_SIDE_EFFECT)
        
        # Fallback to conservative guess
        if isinstance(expr, (ir.IRVariable, ir.IRLiteral, ir.IRTableLiteral)):
            return False
            
        if isinstance(expr, (ir.IRFunctionCall, ir.IRMethodCall, ir.IRAwait)):
            return True
        if hasattr(expr, "__dict__"):
            for v in expr.__dict__.values():
                if isinstance(v, ir.IRExpression):
                    if self._has_side_effects(v): return True
                elif isinstance(v, list):
                    for item in v:
                        if isinstance(item, ir.IRExpression):
                            if self._has_side_effects(item): return True
                elif isinstance(v, ir.IRIfElseExpression):
                    if self._has_side_effects(v.condition): return True
                    if self._has_side_effects(v.then_expr): return True
                    if self._has_side_effects(v.else_expr): return True
        return False

class SSAMarkAndSweepDCE:
    """
    Performs context-aware Dead Code Elimination on SSA IR using a mark-and-sweep algorithm.
    Correctly handles dead cycles of Phi nodes.
    """
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.live_stmts: Set[int] = set() # Set of statement ids (id(stmt))
        self.worklist: List[ir.IRStatement] = []
        self.var_to_def: Dict[str, ir.IRStatement] = {}
        self.captured_names: Set[str] = set()

    def run(self, initial_vars: Optional[List[str]] = None) -> bool:
        self.live_stmts = set()
        self.worklist = []
        self.var_to_def = {}
        self.captured_names = set()
        self.initial_vars = set(initial_vars) if initial_vars else set()

        # 1. Map all versioned variables to their definitions and identify loop variables
        self.loop_vars = set()
        for bid, block in self.cfg.blocks.items():
            for stmt in block.statements:
                if isinstance(stmt, ir.IRAssignment):
                    if isinstance(stmt.target, ir.IRVariable):
                        self.var_to_def[stmt.target.name] = stmt
                elif isinstance(stmt, ir.IRMultiAssign):
                    for target in stmt.targets:
                        if isinstance(target, ir.IRVariable):
                            self.var_to_def[target.name] = stmt
                elif isinstance(stmt, ir.IRFunctionDef):
                    # Track captured names
                    self._find_captured_uses(stmt)
                    
            if block.terminator:
                if isinstance(block.terminator, ir.IRGenericForTerminator):
                    for v in block.terminator.vars:
                        self.loop_vars.add(v)
                elif isinstance(block.terminator, ir.IRBranch):
                    # IRForStatement is compiled to IRBranch and IRGenericForTerminator,
                    # but maybe we should just catch any loop var. The terminator handles it.
                    pass

        # 2. Initial Mark: Side-effecting statements and control flow are live
        for bid, block in self.cfg.blocks.items():
            for stmt in block.statements:
                if self._is_initially_live(stmt):
                    self._mark_live(stmt)
            
            if block.terminator and self._is_initially_live(block.terminator):
                self._mark_live(block.terminator)

        # 2.5. Mark assignments to non-local initial_vars as live
        # (These are upvalues/nonlocals whose effects are visible to the parent function)
        for name, stmt in self.var_to_def.items():
            orig_name = None
            if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.target, ir.IRVariable):
                orig_name = stmt.target.original_name
            elif isinstance(stmt, ir.IRMultiAssign):
                for t in stmt.targets:
                    if isinstance(t, ir.IRVariable) and t.name == name:
                        orig_name = t.original_name
                        break
            
            if orig_name in self.initial_vars:
                # IMPORTANT: Only treat as side-effecting if it's NOT a parameter.
                # In RPy, parameters are passed in initial_vars too. 
                # But SSA pass usually identifies upvalues separately?
                # Actually, in _optimize_block_list, initial_vars contains BOTH.
                # However, our Python source has 'nonlocal' now.
                # For now, let's keep all assignments to initial_vars to be safe.
                self._mark_live(stmt)

        # 3. Mark any assignment to a captured variable as live
        for var_name, stmt in self.var_to_def.items():
            orig_name = None
            if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.target, ir.IRVariable):
                orig_name = stmt.target.original_name
            elif isinstance(stmt, ir.IRMultiAssign):
                for t in stmt.targets:
                    if isinstance(t, ir.IRVariable) and t.name == var_name:
                        orig_name = t.original_name
                        break
            
            if orig_name in self.captured_names:
                # Loop variables captured by lambdas should NOT be forced to keep their
                # injected 'None' initializations, because the loop natively provides
                # the scope and Luau correctly captures the per-iteration variable.
                # If we keep the None init, generator.py will emit a shadowing `local item = nil`.
                if orig_name in self.loop_vars:
                    if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.value, ir.IRLiteral) and stmt.value.value is None:
                        continue
                self._mark_live(stmt)

        # 4. Propagation
        while self.worklist:
            stmt = self.worklist.pop()
            self._mark_uses_live(stmt)

        # 5. Sweep
        changed = False
        for bid, block in self.cfg.blocks.items():
            new_stmts = []
            for stmt in block.statements:
                if id(stmt) in self.live_stmts:
                    new_stmts.append(stmt)
                else:
                    changed = True
                    if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.target, ir.IRVariable):
                        logger.explain("DCE", f"Removing DEAD assignment to {stmt.target.name}")
            block.statements = new_stmts
            
        return changed

    def _mark_live(self, stmt: ir.IRStatement):
        if id(stmt) not in self.live_stmts:
            self.live_stmts.add(id(stmt))
            self.worklist.append(stmt)

    def _is_initially_live(self, stmt: Union[ir.IRStatement, ir.IRExpression]) -> bool:
        if isinstance(stmt, (ir.IRReturnStatement, ir.IRMultiReturn, ir.IRJump, ir.IRBranch, ir.IRGenericForTerminator)):
            return True
        if isinstance(stmt, (ir.IRFunctionDef, ir.IRClassDef, ir.IRImport, ir.IRImportFrom, ir.IRPersistentAssign)):
            return True
        
        if isinstance(stmt, ir.IRAssignment):
            if not isinstance(stmt.target, ir.IRVariable):
                return True
            if self._has_side_effects(stmt.value):
                return True
            return False
            
        if isinstance(stmt, ir.IRMultiAssign):
            if stmt.rhs_call and self._has_side_effects(stmt.rhs_call):
                return True
            for val in stmt.rhs_values:
                if self._has_side_effects(val): return True
            return False
        
        # Standalone expressions in the statement list (like function calls for their side-effects)
        if isinstance(stmt, ir.IRExpression):
            return self._has_side_effects(stmt)
            
        return False

    def _mark_uses_live(self, node: Any):
        if isinstance(node, ir.IRVariable):
            def_stmt = self.var_to_def.get(node.name)
            if def_stmt:
                self._mark_live(def_stmt)
        elif isinstance(node, ir.IRPhi):
            for opt in node.options.values():
                self._mark_uses_live(opt)
        elif hasattr(node, "__dict__"):
            for k, v in node.__dict__.items():
                if isinstance(node, ir.IRAssignment) and k == "target":
                    if isinstance(v, ir.IRVariable): continue
                if isinstance(v, (ir.IRExpression, ir.IRStatement, list, dict)):
                    self._mark_node_uses_live(v)

    def _mark_node_uses_live(self, node: Any):
        if isinstance(node, list):
            for i in node: self._mark_node_uses_live(i)
        elif isinstance(node, dict):
            for v in node.values(): self._mark_node_uses_live(v)
        elif isinstance(node, (ir.IRExpression, ir.IRStatement)):
            self._mark_uses_live(node)

    def _find_captured_uses(self, node: Any):
        temp_dce = SSADCE(self.cfg)
        temp_dce._find_captured_uses(node)
        self.captured_names.update(temp_dce.captured_names)

    def _has_side_effects(self, expr: ir.IRExpression) -> bool:
        temp_dce = SSADCE(self.cfg)
        return temp_dce._has_side_effects(expr)

class SSACSE(visitor.IRTransformer):
    """Performs Global Value Numbering (CSE) on SSA IR."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.expr_to_var: Dict[str, tuple[str, str]] = {} # canonical_expr_str -> (var_name, original_name)
        self.changed = False

    def run(self) -> bool:
        self.changed = False
        # Use a scoped approach based on dominator tree
        dom = DominatorAnalysis(self.cfg)
        dom.run()
        
        # Stack for scoped expr_to_var maps
        self._process_scoped(self.cfg.entry_block.id, {}, dom)
        
        return self.changed

    def _process_scoped(self, bid: int, scoped_exprs: Dict[str, tuple[str, str]], dom: DominatorAnalysis):
        # 1. Capture current state for this block
        block_exprs = scoped_exprs.copy()
        block = self.cfg.blocks[bid]
        
        # 2. Process statements in this block
        self.expr_to_var = block_exprs # Temporarily point to current scope
        block.statements = self.visit(block.statements)
        if block.terminator:
            block.terminator = self.visit(block.terminator)
            
        # 3. Recurse to dominator tree children
        for child_id in dom.dom_tree.get(bid, []):
            self._process_scoped(child_id, block_exprs, dom)

    def visit_IRFunctionDef(self, node: ir.IRFunctionDef) -> ir.IRFunctionDef:
        # Skip recursing into function bodies during CSE
        return node

    def visit_IRAssignment(self, stmt: ir.IRAssignment) -> ir.IRAssignment:
        if isinstance(stmt.target, ir.IRVariable):
            target_var = cast(ir.IRVariable, stmt.target)
            # Canonicalize the value
            val = stmt.value
            if self._is_cse_candidate(val):
                canon = self._get_canonical(val)
                if canon in self.expr_to_var:
                    # Only replace if it's a temporary variable.
                    # If we replace a user variable assignment with another version of the same variable
                    # (e.g. i_5 = i_3), it will become a no-op (i = i) during SSA destruction, 
                    # effectively stripping the assignment!
                    if target_var.name.startswith("_"):
                        logger.explain("CSE", f"Duplicate expression eliminated: {target_var.name} -> {replacement_name}")
                        stmt.value = ir.IRVariable(name=replacement_name, original_name=replacement_orig)
                        self.changed = True
                    else:
                        # For user variables, we can still use this computation as a source for FUTURE eliminations
                        # but we cannot eliminate THIS assignment.
                        pass
                else:
                    # Store name AND original_name of the first computation
                    orig = target_var.original_name or target_var.name
                    self.expr_to_var[canon] = (target_var.name, orig)
        return stmt
        
    def visit_IRIfElseExpression(self, node: ir.IRIfElseExpression) -> ir.IRIfElseExpression:
        node.condition = self.visit(node.condition)
        node.then_expr = self.visit(node.then_expr)
        node.else_expr = self.visit(node.else_expr)
        return node

    def _is_cse_candidate(self, expr: ir.IRExpression) -> bool:
        # Only pure/reads-only expressions are candidates
        if expr.effect_level != ir.EffectLevel.UNKNOWN:
            return expr.effect_level in (ir.EffectLevel.PURE, ir.EffectLevel.READS_MEMORY)
        # Fallback: literal or binary/unary ops without calls
        if isinstance(expr, (ir.IRLiteral, ir.IRVariable)): return False # redundant
        if isinstance(expr, (ir.IRBinaryOperation, ir.IRUnaryOperation)):
            return not self._has_calls(expr)
        return False

    def _has_calls(self, expr: ir.IRExpression) -> bool:
        if isinstance(expr, (ir.IRFunctionCall, ir.IRMethodCall)): return True
        # Use IRVisitor to find calls recursively
        class CallFinder(visitor.IRVisitor):
            def __init__(self): self.has_call = False
            def visit_IRFunctionCall(self, node): self.has_call = True
            def visit_IRMethodCall(self, node): self.has_call = True
        
        finder = CallFinder()
        finder.visit(expr)
        return finder.has_call

    def _get_canonical(self, expr: ir.IRExpression) -> str:
        if isinstance(expr, ir.IRLiteral): return f"L({repr(expr.value)})"
        if isinstance(expr, ir.IRVariable): return f"V({expr.name})"
        if isinstance(expr, ir.IRIfElseExpression):
            return f"IF({self._get_canonical(expr.condition)}, {self._get_canonical(expr.then_expr)}, {self._get_canonical(expr.else_expr)})"
        if isinstance(expr, ir.IRPropertyAccess):
            return f"PROP({self._get_canonical(expr.receiver)}, {expr.property})"
        if isinstance(expr, ir.IRIndexOperation):
            return f"INDEX({self._get_canonical(expr.receiver)}, {self._get_canonical(expr.index)})"
        if isinstance(expr, ir.IRUnaryOperation):
            return f"UN({expr.operator}, {self._get_canonical(expr.operand)})"
        if isinstance(expr, ir.IRBinaryOperation):
            bin_op = cast(ir.IRBinaryOperation, expr)
            l = self._get_canonical(bin_op.left)
            r = self._get_canonical(bin_op.right)
            op = bin_op.operator
            # Commutativity
            if op in ("Add", "Mult", "Eq", "NotEq", "BitAnd", "BitOr", "BitXor"):
                if l > r: l, r = r, l
            return f"B({l},{op},{r})"
        if isinstance(expr, ir.IRUnaryOperation):
            un_op = cast(ir.IRUnaryOperation, expr)
            return f"U({un_op.operator},{self._get_canonical(un_op.operand)})"
        return str(id(expr))

class _SSAUseChecker(visitor.IRVisitor):
    def __init__(self, current_bid: int, defs: Dict[str, int], dominators: Dict[int, Set[int]]):
        self.bid = current_bid
        self.defs = defs
        self.dominators = dominators

    def visit_IRVariable(self, var: ir.IRVariable):
        if var.name not in self.defs: return
        def_bid = self.defs[var.name]
        if def_bid not in self.dominators[self.bid]:
            # SSA dominance check: def must dominate use
            pass

    def visit_IRPhi(self, phi: ir.IRPhi):
        # Phi uses are special: they are defined in predecessor blocks
        for pred_id, opt in phi.options.items():
            if isinstance(opt, ir.IRVariable):
                var = cast(ir.IRVariable, opt)
                if var.name in self.defs:
                    def_bid = self.defs[var.name]
                    if def_bid not in self.dominators[pred_id]:
                        raise ValueError(f"SSA violation: Def of {var.name} in B{def_bid} does not dominate use in Phi edge from B{pred_id}")

class SSAVerifier:
    """Verifies SSA invariants in the CFG."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.defs: Dict[str, int] = {} # var -> block_id
        self.current_block_id = -1

    def run(self):
        self.defs = {}
        # 1. First pass: collect all definitions and check for multiple defs
        for bid, block in self.cfg.blocks.items():
            self.current_block_id = bid
            for stmt in block.statements:
                self._collect_defs(stmt)

        # 2. Second pass: check dominance of all uses
        ana = DominatorAnalysis(self.cfg)
        ana.run()
        for bid, block in self.cfg.blocks.items():
            for stmt in block.statements:
                checker = _SSAUseChecker(bid, self.defs, ana.dominators)
                checker.visit(stmt)

    def _collect_defs(self, node: Any):
        if isinstance(node, ir.IRAssignment):
            if isinstance(node.target, ir.IRVariable):
                target_var = cast(ir.IRVariable, node.target)
                name = target_var.name
                if name in self.defs:
                    raise ValueError(f"SSA violation: Multiple definitions of {name}")
                self.defs[name] = self.current_block_id
        elif isinstance(node, ir.IRFunctionDef):
            func_def = cast(ir.IRFunctionDef, node)
            for s in func_def.body: self._collect_defs(s)

class SSAConstantPropagation:
    """Performs Constant Propagation and Folding on SSA IR."""
    UNKNOWN = object()

    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.constants: Dict[str, Any] = {} # name -> literal value

    def run(self) -> bool:
        overall_changed = False
        limit = 20
        iterations = 0
        changed = True
        while changed and iterations < limit:
            changed = False
            iterations += 1
            
            # 1. Collect and fold recursively
            # print(f"--- ConstProp Iteration {iterations} ---")
            for bid, block in self.cfg.blocks.items():
                if self._collect_and_fold_recursive(block):
                    changed = True
            
            # 2. Propagate recursively
            for bid, block in self.cfg.blocks.items():
                for stmt in block.statements:
                    if self._propagate_in_stmt(stmt):
                        changed = True
            
            if changed: overall_changed = True
            
        if iterations >= limit:
             logger.explain("ConstProp", "WARNING: Constant propagation reached iteration limit. Possible infinite loop detected.")
        return overall_changed

    def _collect_and_fold_recursive(self, container: Any) -> bool:
        changed = False
        stmts: Optional[List[ir.IRStatement]] = None
        if hasattr(container, "statements"): stmts = container.statements
        elif isinstance(container, list): stmts = container
        
        if stmts:
            for stmt in stmts:
                if isinstance(stmt, ir.IRFunctionDef):
                    continue # Don't recurse into function bodies during block-level CP
                
                if isinstance(stmt, ir.IRAssignment):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    if isinstance(stmt_assign.target, ir.IRVariable):
                        target_var = cast(ir.IRVariable, stmt_assign.target)
                        old_val = self.constants.get(target_var.name, self.UNKNOWN)
                        
                        # Fold expressions using known constants
                        if isinstance(stmt_assign.value, ir.IRVariable):
                            val_var = cast(ir.IRVariable, stmt_assign.value)
                            c = self._get_const(val_var)
                            if c is not self.UNKNOWN:
                                stmt_assign.value = ir.IRLiteral(value=c)
                                changed = True
                        
                        if isinstance(stmt_assign.value, ir.IRLiteral):
                            literal_val = cast(ir.IRLiteral, stmt_assign.value)
                            self.constants[target_var.name] = literal_val.value
                            # print(f"  Def: {target_var.name} = {literal_val.value}")
                        elif isinstance(stmt_assign.value, ir.IRPhi):
                            phi_val = cast(ir.IRPhi, stmt_assign.value)
                            # print(f"  Evaluating Phi for {target_var.name} with options: {list(phi_val.options.keys())}")
                            values = []
                            for v in phi_val.options.values():
                                c = self._get_const(v)
                                if c is not self.UNKNOWN: values.append(c)
                                else: 
                                    # if isinstance(v, ir.IRVariable): 
                                        # v_var = cast(ir.IRVariable, v)
                                        # print(f"    Option {v_var.name} is UNKNOWN")
                                    values = None; break
                            
                            if values and len(values) > 0:
                                # print(f"  Phi values for {target_var.name}: {values}")
                                if all(v == values[0] for v in values):
                                    self.constants[target_var.name] = values[0]
                                    stmt_assign.value = ir.IRLiteral(value=values[0], lineno=stmt_assign.lineno)
                                    changed = True
                                    # print(f"  PhiFold: {target_var.name} = {values[0]}")
                                        
                        elif isinstance(stmt_assign.value, ir.IRBinaryOperation):
                            bin_op = cast(ir.IRBinaryOperation, stmt_assign.value)
                            l = self._get_const(bin_op.left)
                            r = self._get_const(bin_op.right)
                            if l is not self.UNKNOWN and r is not self.UNKNOWN:
                                res = self._fold_bin(l, bin_op.operator, r)
                                if res is not None:
                                    logger.explain("ConstFold", f"Folding {target_var.name}: {l} {bin_op.operator} {r} = {res}")
                                    stmt_assign.value = ir.IRLiteral(value=res, lineno=stmt_assign.lineno)
                                    self.constants[target_var.name] = res
                                    changed = True
                        elif isinstance(stmt_assign.value, ir.IRUnaryOperation):
                            un_op = cast(ir.IRUnaryOperation, stmt_assign.value)
                            val = self._get_const(un_op.operand)
                            if val is not self.UNKNOWN:
                                res = self._fold_un(un_op.operator, val)
                                if res is not None:
                                    stmt_assign.value = ir.IRLiteral(value=res, lineno=stmt_assign.lineno)
                                    self.constants[target_var.name] = res
                                    changed = True
                        
                        if self.constants.get(target_var.name, self.UNKNOWN) != old_val:
                            changed = True
                
                # Recurse into nested blocks/functions
                if hasattr(stmt, "__dict__"):
                    for v in stmt.__dict__.values():
                        if isinstance(v, list):
                            if self._collect_and_fold_recursive(v): changed = True
        return changed

    def _get_const(self, expr: ir.IRExpression) -> Any:
        if isinstance(expr, ir.IRLiteral): return expr.value
        if isinstance(expr, ir.IRVariable): return self.constants.get(expr.name, self.UNKNOWN)
        return self.UNKNOWN

    def _propagate_in_stmt(self, stmt: ir.IRStatement) -> bool:
        changed = False
        if hasattr(stmt, "__dict__"):
            for k, v in stmt.__dict__.items():
                if isinstance(stmt, ir.IRAssignment) and k == "target":
                    if isinstance(v, ir.IRVariable):
                        continue
                if isinstance(v, ir.IRExpression):
                    new_v = self._propagate_in_expr(v)
                    if new_v is not v:
                        setattr(stmt, k, new_v)
                        changed = True
                elif isinstance(v, list):
                    # CRITICAL: Skip recursing into structured bodies during CFG-level ConstProp
                    if k in ("body", "then_block", "else_block", "elif_blocks"):
                        continue
                    for i, item in enumerate(v):
                        if isinstance(item, ir.IRExpression):
                            new_item = self._propagate_in_expr(item)
                            if new_item is not item:
                                v[i] = new_item
                                changed = True
                        elif isinstance(item, ir.IRStatement):
                            # Function bodies are handled separately
                            if isinstance(item, ir.IRFunctionDef):
                                continue
                            if self._propagate_in_stmt(item):
                                changed = True
        
        # Specific folding logic for branches
        if isinstance(stmt, ir.IRBranch):
            branch_stmt = cast(ir.IRBranch, stmt)
            c_val = self._get_const(branch_stmt.condition)
            if c_val is not self.UNKNOWN:
                is_true = bool(c_val)
                if not isinstance(branch_stmt.condition, ir.IRLiteral) or branch_stmt.condition.value != is_true:
                    branch_stmt.condition = ir.IRLiteral(value=is_true, lineno=branch_stmt.condition.lineno)
                    changed = True
        return changed

    def _propagate_in_expr(self, expr: ir.IRExpression) -> ir.IRExpression:
        if isinstance(expr, ir.IRVariable) and expr.name in self.constants:
            val = self.constants[expr.name]
            # print(f"  Propagating {expr.name} -> {val}")
            return ir.IRLiteral(value=val, lineno=expr.lineno)
            
        if isinstance(expr, ir.IRBinaryOperation):
            bin_op = cast(ir.IRBinaryOperation, expr)
            bin_op.left = self._propagate_in_expr(bin_op.left)
            bin_op.right = self._propagate_in_expr(bin_op.right)
            l = self._get_const(bin_op.left)
            r = self._get_const(bin_op.right)
            if l is not self.UNKNOWN and r is not self.UNKNOWN:
                res = self._fold_bin(l, bin_op.operator, r)
                if res is not None:
                    return ir.IRLiteral(value=res, lineno=expr.lineno)
                    
        elif isinstance(expr, ir.IRUnaryOperation):
            un_op = cast(ir.IRUnaryOperation, expr)
            un_op.operand = self._propagate_in_expr(un_op.operand)
            val = self._get_const(un_op.operand)
            if val is not self.UNKNOWN:
                res = self._fold_un(un_op.operator, val)
                if res is not None:
                    return ir.IRLiteral(value=res, lineno=expr.lineno)
        
        elif isinstance(expr, ir.IRIfElseExpression):
            expr.condition = self._propagate_in_expr(expr.condition)
            expr.then_expr = self._propagate_in_expr(expr.then_expr)
            expr.else_expr = self._propagate_in_expr(expr.else_expr)
            # Constant folding for ternary itself
            c_cond = self._get_const(expr.condition)
            if c_cond is not self.UNKNOWN:
                return expr.then_expr if bool(c_cond) else expr.else_expr

        elif hasattr(expr, "__dict__"):
            for k, v in expr.__dict__.items():
                if k in ("left", "right", "operand"): continue # Already handled
                if isinstance(v, ir.IRExpression):
                    new_v = self._propagate_in_expr(v)
                    if new_v is not v:
                        setattr(expr, k, new_v)
                elif isinstance(v, list):
                    for i, item in enumerate(v):
                        if isinstance(item, ir.IRExpression):
                            new_i = self._propagate_in_expr(item)
                            if new_i is not item:
                                v[i] = new_i
        return expr

    def _fold_bin(self, l: Any, op: str, r: Any) -> Optional[Any]:
        # print(f"DEBUG FOLD: {l} {repr(op)} {r}")
        try:
            if op == "Add": return l + r
            if op == "Sub": return l - r
            if op == "Mult":
                res = l * r
                # print(f"  FOLDED MULT: {res}")
                return res
            if op == "Div": return l / r if r != 0 else None
            if op == "FloorDiv": return l // r if r != 0 else None
            if op == "Mod": return l % r if r != 0 else None
            if op == "Pow": return l ** r
            if op == "LShift": return l << r
            if op == "RShift": return l >> r
            if op == "BitOr": return l | r
            if op == "BitXor": return l ^ r
            if op == "BitAnd": return l & r
            if op == "Gt": return l > r
            if op == "Lt": return l < r
            if op == "GtE": return l >= r
            if op == "LtE": return l <= r
            if op == "Eq": return l == r
            if op == "NotEq": return l != r
        except: pass
        return None

    def _fold_un(self, op: str, val: Any) -> Optional[Any]:
        try:
            if op == "Invert": return ~val
            if op == "Not": return not val
            if op == "UAdd": return +val
            if op == "USub": return -val
        except: pass
        return None

class DeSSA:
    """Converts SSA IR back to normal form (eliminates Phi nodes)."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg

    def run(self):
        # 1. Eliminate Phi nodes by inserting assignments in predecessors
        for bid, block in self.cfg.blocks.items():
            new_statements = []
            for stmt in block.statements:
                if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.value, ir.IRPhi):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    phi = cast(ir.IRPhi, stmt_assign.value)
                    target_var = cast(ir.IRVariable, stmt_assign.target)
                    for pred_id, source_expr in phi.options.items():
                        pred_block = self.cfg.blocks[pred_id]
                        # Propagate original_name if possible
                        orig = getattr(target_var, "original_name", target_var.name)
                        assignment = ir.IRAssignment(
                            target=ir.IRVariable(name=target_var.name, original_name=orig, lineno=target_var.lineno), 
                            value=source_expr,
                            lineno=target_var.lineno
                        )
                        if pred_block.statements and isinstance(pred_block.statements[-1], (ir.IRJump, ir.IRBranch)):
                            pred_block.statements.insert(-1, assignment)
                        else:
                            pred_block.statements.append(assignment)
                    continue # Do not add the Phi-assignment to the current block
                new_statements.append(stmt)
            block.statements = new_statements

        # 2. Rename variables back to original names and remove self-assignments
        for bid, block in self.cfg.blocks.items():
            new_stmts = []
            for stmt in block.statements:
                self._rename_back(stmt)
                if isinstance(stmt, ir.IRAssignment):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    if isinstance(stmt_assign.target, ir.IRVariable) and isinstance(stmt_assign.value, ir.IRVariable):
                        if stmt_assign.target.name == stmt_assign.value.name:
                            continue
                new_stmts.append(stmt)
            block.statements = new_stmts

    def _rename_back(self, stmt_or_expr: Union[ir.IRStatement, ir.IRExpression]):
        if isinstance(stmt_or_expr, ir.IRFunctionDef):
            import re
            # Also rename formal parameters back to original names
            func_def = cast(ir.IRFunctionDef, stmt_or_expr)
            new_args = []
            for arg in func_def.args:
                # Optimized regex-free check for _N suffix
                if "_" in arg:
                    parts = arg.rsplit("_", 1)
                    if parts[1].isdigit():
                        new_args.append(parts[0])
                        continue
                new_args.append(arg)
            func_def.args = new_args
        
        if isinstance(stmt_or_expr, ir.IRVariable):
            var_expr = cast(ir.IRVariable, stmt_or_expr)
            orig_name = getattr(var_expr, "original_name", None)
            if orig_name:
                var_expr.name = orig_name

        elif hasattr(stmt_or_expr, "__dict__"):
            for v in stmt_or_expr.__dict__.values():
                if isinstance(v, (ir.IRStatement, ir.IRExpression)):
                    self._rename_back(v)
                elif isinstance(v, list):
                    for item in v:
                        if isinstance(item, (ir.IRStatement, ir.IRExpression)):
                            self._rename_back(item)

class PhiSimplification:
    """Simplifies redundant Phi nodes."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg

    def run(self) -> bool:
        changed = False
        for bid, block in self.cfg.blocks.items():
            new_stmts = []
            for stmt in block.statements:
                if isinstance(stmt, ir.IRAssignment) and isinstance(stmt.value, ir.IRPhi):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    phi = cast(ir.IRPhi, stmt_assign.value)
                    phi_target = stmt_assign.target.name
                    options = [v for v in phi.options.values() if not (isinstance(v, ir.IRVariable) and v.name == phi_target)]
                    
                    # Case 1: All non-self options are the same literal or variable
                    if options:
                        first = options[0]
                        if all(self._is_same(first, o) for o in options):
                            stmt_assign.value = first
                            changed = True
                    elif not phi.options:
                        # Empty phi (should not happen in valid SSA)
                        pass
                new_stmts.append(stmt)
            block.statements = new_stmts
        return changed

    def _is_same(self, e1: ir.IRExpression, e2: ir.IRExpression) -> bool:
        if isinstance(e1, ir.IRLiteral) and isinstance(e2, ir.IRLiteral):
            return e1.value == e2.value
        if isinstance(e1, ir.IRVariable) and isinstance(e2, ir.IRVariable):
            return e1.name == e2.name
        return False

class SSALICM:
    """Loop Invariant Code Motion on SSA IR."""
    def __init__(self, target_cfg: cfg.CFG, loops: Dict[int, loop_detector.LoopInfo]):
        self.cfg = target_cfg
        self.loops = loops

    def run(self) -> bool:
        changed = False
        # Process loops inner-to-outer (optional but good)
        for _, loop in sorted(self.loops.items(), key=lambda x: len(x[1].body_ids)):
            if self._process_loop(loop):
                changed = True
        return changed

    def _process_loop(self, loop: loop_detector.LoopInfo) -> bool:
        changed = False
        # 1. Identify all definitions in the loop
        defs_in_loop = set()
        for bid in loop.body_ids:
            block = self.cfg.blocks[bid]
            for stmt in block.statements:
                if isinstance(stmt, ir.IRAssignment):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    if isinstance(stmt_assign.target, ir.IRVariable):
                        defs_in_loop.add(stmt_assign.target.name)

        # 2. Iterate until we can't hoist anymore
        hoisted_anything = True
        while hoisted_anything:
            hoisted_anything = False
            for bid in sorted(list(loop.body_ids)):
                block = self.cfg.blocks[bid]
                new_stmts = []
                for stmt in block.statements:
                    if self._can_hoist(stmt, def_in_loop=defs_in_loop):
                        # Hoist it!
                        self._hoist(stmt, loop.header_id)
                        hoisted_anything = True
                        changed = True
                        if isinstance(stmt, ir.IRAssignment):
                            stmt_assign = cast(ir.IRAssignment, stmt)
                            if isinstance(stmt_assign.target, ir.IRVariable):
                                target_var = cast(ir.IRVariable, stmt_assign.target)
                                logger.explain("LICM", f"Hoisted loop-invariant: {target_var.name}")
                                # It's no longer defined in the loop body
                                defs_in_loop.discard(target_var.name)
                        continue
                    new_stmts.append(stmt)
                block.statements = new_stmts
        return changed

    def _can_hoist(self, stmt: ir.IRStatement, def_in_loop: Set[str]) -> bool:
        if not isinstance(stmt, ir.IRAssignment):
            return False
        stmt_assign = cast(ir.IRAssignment, stmt)
        if not isinstance(stmt_assign.target, ir.IRVariable):
            return False
            
        # NEVER hoist a Phi node
        if isinstance(stmt_assign.value, ir.IRPhi):
            return False

        # 1. Expression must be loop invariant (operands not defined in loop)
        if not self._is_invariant(stmt_assign.value, def_in_loop):
            return False
            
        # 2. Expression must be pure or read-only (no side effects)
        if stmt_assign.value.effect_level not in (ir.EffectLevel.PURE, ir.EffectLevel.READS_MEMORY):
            return False
            
        return True

    def _is_invariant(self, expr: ir.IRExpression, def_in_loop: Set[str]) -> bool:
        class InvariantVisitor(visitor.IRVisitor):
            def __init__(self, defs: Set[str]):
                self.defs = defs
                self.is_invariant = True
            def visit_IRVariable(self, node: ir.IRVariable):
                if node.name in self.defs:
                    self.is_invariant = False
            def visit_IRBinaryOperation(self, node: ir.IRBinaryOperation):
                self.visit(node.left)
                self.visit(node.right)
            def visit_IRLiteral(self, node: ir.IRLiteral):
                pass
            def visit_IRIfElseExpression(self, node: ir.IRIfElseExpression):
                self.visit(node.condition)
                self.visit(node.then_expr)
                self.visit(node.else_expr)
            def visit_IRPropertyAccess(self, node: ir.IRPropertyAccess):
                self.visit(node.receiver)
            def visit_IRUnaryOperation(self, node: ir.IRUnaryOperation):
                self.visit(node.operand)
            def visit_IRFunctionCall(self, node): self.is_invariant = False
            def visit_IRMethodCall(self, node): self.is_invariant = False
            def visit_IRAwait(self, node): self.is_invariant = False
            def visit_IRList(self, node: ir.IRList):
                for item in node.elements: self.visit(item)
            def visit_IRDict(self, node: ir.IRDict):
                for k, v in node.pairs:
                    self.visit(k)
                    self.visit(v)
            def visit_IRPhi(self, node): self.is_invariant = False
            def visit_IRSuperCall(self, node): self.is_invariant = False
            def visit_IRIndexOperation(self, node): 
                self.visit(node.receiver)
                self.visit(node.index)

        checker = InvariantVisitor(def_in_loop)
        checker.visit(expr)
        return checker.is_invariant

    def _hoist(self, stmt: ir.IRStatement, header_id: int):
        # Find pre-header of the loop header
        header = self.cfg.blocks[header_id]
        # In simple RPy loops, the pre-header is the unique predecessor that doesn't reach back?
        # A more robust way: create a dedicated pre-header block.
        # For now, we search for a predecessor that is NOT in the loop.
        pre_header = None
        for pred in header.predecessors:
            if pred.id not in self.loops[header_id].body_ids:
                pre_header = pred
                break
        
        if pre_header:
            # Place at the end of pre-header, before branch/jump
            if pre_header.statements and isinstance(pre_header.statements[-1], (ir.IRJump, ir.IRBranch)):
                pre_header.statements.insert(-1, stmt)
            else:
                pre_header.statements.append(stmt)
        else:
            # Fallback: keep in loop (should not happen for natural loops)
            pass

class SSACopyPropagation:
    """Performs Copy Propagation on SSA IR."""
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.aliases: Dict[str, tuple[str, str]] = {} # name -> (replacement_name, replacement_original_name)

    def run(self) -> bool:
        self.aliases = {}
        overall_changed = False
        changed = True
        
        while changed:
            changed = False
            # 1. Count definitions and collect candidates
            def_counts: Dict[str, int] = {}
            for bid, block in self.cfg.blocks.items():
                for stmt in block.statements:
                    if isinstance(stmt, ir.IRFunctionDef): continue
                    if isinstance(stmt, ir.IRAssignment):
                        stmt_assign = cast(ir.IRAssignment, stmt)
                        if isinstance(stmt_assign.target, ir.IRVariable):
                            target_name = stmt_assign.target.name
                            def_counts[target_name] = def_counts.get(target_name, 0) + 1

            for bid, block in self.cfg.blocks.items():
                for stmt in block.statements:
                    if isinstance(stmt, ir.IRFunctionDef): continue
                    if isinstance(stmt, ir.IRAssignment):
                        stmt_assign = cast(ir.IRAssignment, stmt)
                        if isinstance(stmt_assign.target, ir.IRVariable):
                            target_name = stmt_assign.target.name
                            if def_counts.get(target_name, 0) == 1 and isinstance(stmt_assign.value, ir.IRVariable):
                                val_var = cast(ir.IRVariable, stmt_assign.value)
                                src_name = val_var.name
                                src_orig = val_var.original_name or val_var.name
                                # Follow alias chain to get the ultimate source
                                while src_name in self.aliases:
                                    src_name, src_orig = self.aliases[src_name]
                                
                                if src_name != target_name:
                                    current = self.aliases.get(target_name)
                                    if current != (src_name, src_orig):
                                        self.aliases[target_name] = (src_name, src_orig)
                                        changed = True
            
            # 2. Propagate aliases
            for bid, block in self.cfg.blocks.items():
                for stmt in block.statements:
                    if self._propagate_in_stmt(stmt):
                        changed = True
            
            if changed: overall_changed = True
            
        return overall_changed

    def _propagate_in_stmt(self, stmt: ir.IRStatement) -> bool:
        class Propagator(visitor.IRTransformer):
            def __init__(self, aliases: Dict[str, tuple[str, str]]):
                self.aliases = aliases
                self.changed = False
            
            def visit_IRAssignment(self, node: ir.IRAssignment) -> ir.IRAssignment:
                # Don't propagate into target of assignment
                node.value = self.visit(node.value)
                return node
            
            def visit_IRFunctionDef(self, node: ir.IRFunctionDef) -> ir.IRFunctionDef:
                return node
            
            def visit_IRWhileStatement(self, node: ir.IRWhileStatement) -> ir.IRWhileStatement:
                # Only visit condition, skip body
                node.condition = self.visit(node.condition)
                return node
                
            def visit_IRIfStatement(self, node: ir.IRIfStatement) -> ir.IRIfStatement:
                # Only visit condition, skip bodies
                node.condition = self.visit(node.condition)
                return node
                node.then_expr = self.visit(node.then_expr)
                node.else_expr = self.visit(node.else_expr)
                return node

        prop = Propagator(self.aliases)
        prop.visit(stmt)
        return prop.changed

class SSAExpressionInliner:
    """
    Inlines SSA variables that are used exactly once into their use-site.
    This reconstructs nested expressions from the 3-address IR.
    """
    def __init__(self, target_cfg: cfg.CFG):
        self.cfg = target_cfg
        self.use_counts: Dict[str, int] = {}
        self.definitions: Dict[str, ir.IRExpression] = {}
        self.def_counts: Dict[str, int] = {}
        self.inlined_names: Set[str] = set()

    def run(self) -> bool:
        changed = False
        iteration_changed = True
        
        while iteration_changed:
            iteration_changed = False
            self.use_counts = {}
            self.definitions = {}

            # 1. Count uses and collect definitions recursively
            self.def_counts = {}
            self.use_counts = {}
            for bid, block in self.cfg.blocks.items():
                for stmt in block.statements:
                    if isinstance(stmt, ir.IRFunctionDef): continue
                    self._collect_recursive(stmt)
                if block.terminator:
                    self._collect_recursive(block.terminator)
            
            # 2. Identify inlining candidates (single-use, single-def, pure)
            candidates: Dict[str, ir.IRExpression] = {}
            for name, count in self.use_counts.items():
                if count == 1 and self.def_counts.get(name, 0) == 1 and name in self.definitions:
                    expr = self.definitions[name]
                    eff = getattr(expr, "effect_level", ir.EffectLevel.UNKNOWN)
                    # ONLY inline if it's a temporary (starts with _)
                    if name.startswith("_") and eff == ir.EffectLevel.PURE and not isinstance(expr, ir.IRPhi):
                        candidates[name] = expr

            if not candidates:
                break

            # 3. Apply inlining recursively
            for bid, block in self.cfg.blocks.items():
                if self._apply_recursive(block, candidates):
                    changed = True
                    iteration_changed = True
        
        return changed

    def _collect_recursive(self, node: Any):
        if isinstance(node, ir.IRFunctionDef): return
        if isinstance(node, ir.IRAssignment):
            stmt_assign = cast(ir.IRAssignment, node)
            if isinstance(stmt_assign.target, ir.IRVariable):
                self.definitions[stmt_assign.target.name] = stmt_assign.value
                self.def_counts[stmt_assign.target.name] = self.def_counts.get(stmt_assign.target.name, 0) + 1
        elif isinstance(node, ir.IRMultiAssign):
            # IRMultiAssign targets are NOT inlined as expressions easily
            # But we must track their definitions so they aren't considered dead mistakenly
            for target in node.targets:
                if isinstance(target, ir.IRVariable):
                    # For now, we don't inline multi-returns because Luau syntax 
                    # doesn't allow 'f(g())' to mean something other than f's 1st arg.
                    # We keep them separate.
                    self.definitions[target.name] = ir.IRLiteral(value=None) # dummy
                    self.def_counts[target.name] = self.def_counts.get(target.name, 0) + 1
        
        if isinstance(node, ir.IRVariable):
            self.use_counts[node.name] = self.use_counts.get(node.name, 0) + 1
        elif hasattr(node, "__dict__"):
            for k, v in node.__dict__.items():
                if isinstance(node, ir.IRAssignment) and k == "target":
                    continue
                if isinstance(v, (ir.IRNode, list, dict)):
                    if isinstance(v, list):
                        for item in v: self._collect_recursive(item)
                    elif isinstance(v, dict):
                        # DO NOT recurse into IRPhi.options during collection
                        if isinstance(node, ir.IRPhi):
                             continue
                        for item in v.values(): self._collect_recursive(item)
                    elif isinstance(v, ir.IRIfElseExpression):
                        self._collect_recursive(v.condition)
                        self._collect_recursive(v.then_expr)
                        self._collect_recursive(v.else_expr)
                    else:
                        self._collect_recursive(v)

    def _apply_recursive(self, container: Any, candidates: Dict[str, ir.IRExpression]) -> bool:
        if isinstance(container, ir.IRFunctionDef): return False
        changed = False
        stmts: Optional[List[ir.IRStatement]] = None
        if hasattr(container, "statements"): stmts = container.statements
        elif isinstance(container, list): stmts = container
        
        if stmts is not None:
             new_stmts = []
             for stmt in stmts:
                 skip = False
                 if isinstance(stmt, ir.IRAssignment):
                    stmt_assign = cast(ir.IRAssignment, stmt)
                    if isinstance(stmt_assign.target, ir.IRVariable):
                        if stmt_assign.target.name in candidates:
                            skip = True
                            changed = True
                            continue
                 elif isinstance(stmt, ir.IRMultiAssign):
                     # Never skip MultiAssign (hard to inline)
                     pass
                 
                 if not skip:
                     if self._inline_in_stmt(stmt, candidates):
                         changed = True
                     new_stmts.append(stmt)
             
             if isinstance(container, list):
                 container.clear()
                 container.extend(new_stmts)
             else: container.statements = new_stmts
        return changed

    def _inline_in_stmt(self, stmt: ir.IRStatement, candidates: Dict[str, ir.IRExpression]) -> bool:
        changed = False
        if hasattr(stmt, "__dict__"):
            for k, v in stmt.__dict__.items():
                if isinstance(stmt, ir.IRAssignment) and k == "target":
                    if isinstance(v, ir.IRVariable):
                        continue
                if isinstance(v, ir.IRExpression):
                    new_v = self._inline_in_expr(v, candidates)
                    if new_v is not v:
                        setattr(stmt, k, new_v)
                        changed = True
                elif isinstance(v, list):
                    for i, item in enumerate(v):
                        if isinstance(item, ir.IRExpression):
                            new_item = self._inline_in_expr(item, candidates)
                            if new_item is not item:
                                v[i] = new_item
                                changed = True
                        elif isinstance(item, ir.IRStatement):
                            if self._inline_in_stmt(item, candidates):
                                changed = True
                elif isinstance(v, ir.IRIfElseExpression):
                    v.condition = self._inline_in_expr(v.condition, candidates)
                    v.then_expr = self._inline_in_expr(v.then_expr, candidates)
                    v.else_expr = self._inline_in_expr(v.else_expr, candidates)
                    changed = True
        
        # Also recurse into nested blocks/bodies
        if hasattr(stmt, "body") and isinstance(stmt.body, list):
             if self._apply_recursive(stmt.body, candidates): changed = True
        return changed

    def _inline_in_expr(self, expr: ir.IRExpression, candidates: Dict[str, ir.IRExpression]) -> ir.IRExpression:
        changed = False
        if isinstance(expr, ir.IRVariable) and expr.name in candidates:
            import copy
            return copy.deepcopy(candidates[expr.name])
            
        if isinstance(expr, ir.IRMultiReturn):
            for i, val in enumerate(expr.values):
                expr.values[i] = self._inline_in_expr(val, candidates)
            return expr

        if hasattr(expr, "__dict__"):
            for k, v in expr.__dict__.items():
                if isinstance(v, ir.IRExpression):
                    new_v = self._inline_in_expr(v, candidates)
                    if new_v is not v:
                        setattr(expr, k, new_v)
                elif isinstance(v, list):
                    for i, item in enumerate(v):
                        if isinstance(item, ir.IRExpression):
                            new_item = self._inline_in_expr(item, candidates)
                            if new_item is not item:
                                v[i] = new_item
                                changed = True
                elif isinstance(v, dict):
                    # Handle dictionaries like IRPhi.options
                    # CRITICAL: DO NOT inline into Phi operands!
                    # Phi nodes MUST have variables as inputs for DeSSA.
                    if isinstance(expr, ir.IRPhi):
                        continue
                        
                    for key, val in v.items():
                        if isinstance(val, ir.IRExpression):
                            new_val = self._inline_in_expr(val, candidates)
                            if new_val is not val:
                                v[key] = new_val
                                changed = True
                elif isinstance(v, ir.IRIfElseExpression):
                    v.condition = self._inline_in_expr(v.condition, candidates)
                    v.then_expr = self._inline_in_expr(v.then_expr, candidates)
                    v.else_expr = self._inline_in_expr(v.else_expr, candidates)
                    changed = True
        return expr
