CAP/MiniC/TP05/ExitSSA.py

"""
CAP, SSA Intro, Elimination and Optimisations
Functions to convert a CFG out of SSA Form.
"""

from typing import cast, List
from Lib import RiscV
from Lib.CFG import Block, BlockInstr, CFG
from Lib.Operands import Temporary, DataLocation
from Lib.Statement import AbsoluteJump, Label
from Lib.Terminator import BranchingTerminator, Return
from Lib.PhiNode import PhiNode
from TP05.SequentializeMoves import sequentialize_moves, generate_smart_move


def generate_moves_from_phis(phis: List[PhiNode], parent: Block, is_smart: bool=False) -> List[BlockInstr]:
    """
    `generate_moves_from_phis(phis, parent)` builds a list of move instructions
    to be inserted in a new block between `parent` and the block with phi nodes
    `phis`.

    This is an helper function called during SSA exit.
    """
    # (Lab 5a, Exercise 6)
    lbl = parent.get_label()
    if is_smart:
        parallel_moves: set[tuple[DataLocation, DataLocation]] = set()
        moves: List[BlockInstr] = []
        for phi in phis:
            if lbl not in phi.get_srcs():
                continue
            src = phi.get_srcs()[lbl]
            if isinstance(src, Temporary) or isinstance(phi.var, Temporary):
                loc1, loc2 = phi.var, src
                if isinstance(loc1, Temporary):
                    loc1 = loc1.get_alloced_loc()
                if isinstance(loc2, Temporary):
                    loc2 = loc2.get_alloced_loc()
                moves += generate_smart_move(loc1, cast(DataLocation, loc2))
            else:
                src = phi.get_srcs()[lbl]
                parallel_moves.add((phi.var, cast(DataLocation, src)))
        return sequentialize_moves(parallel_moves)+moves
    return [
        RiscV.mv(
            phi.var,
            phi.get_srcs()[lbl]
        ) for phi in phis if lbl in phi.get_srcs()
    ]


def exit_ssa(cfg: CFG, is_smart: bool) -> None:
    """
    `exit_ssa(cfg)` replaces phi nodes with move instructions to exit SSA form.

    `is_smart` is set to true when smart register allocation is enabled (Lab 5b).
    """
    for b in cfg.get_blocks():
        phis = cast(List[PhiNode], b.get_phis())  # Use cast for Pyright
        b.remove_all_phis()  # Remove all phi nodes in the block
        parents: List[Block] = b.get_in().copy()  # Copy as we modify it by adding blocks
        for parent in parents:
            # Add the block containing 'moves' to 'cfg'
            moves_label = cfg.fdata.fresh_label("merge_phi")
            block_moves = Block(moves_label, generate_moves_from_phis(phis, parent, is_smart=is_smart), AbsoluteJump(b.get_label()))
            cfg.add_block(block_moves)

            # Update parent terminator
            j = parent.get_terminator()
            match j:
                case AbsoluteJump():
                    new_terminator = AbsoluteJump(moves_label)
                case BranchingTerminator():
                    lbl_else : Label = moves_label if (j.label_else == b.get_label()) else j.label_else
                    lbl_then : Label = moves_label if (j.label_then == b.get_label()) else j.label_then
                    new_terminator = BranchingTerminator(j.cond, j.op1, j.op2, lbl_else, lbl_then)
                case Return():
                    new_terminator = AbsoluteJump(moves_label)
                    block_moves.set_terminator(Return())

            parent.set_terminator(new_terminator)