from Lib import RiscV
from Lib.Operands import Temporary, Operand, S
from Lib.Statement import Instruction
from Lib.Allocator import Allocator
from typing import List


class AllInMemAllocator(Allocator):

    def replace(self, old_instr: Instruction) -> List[Instruction]:
        """Replace Temporary operands with the corresponding allocated
        memory location."""
        numreg = 1
        before: List[Instruction] = []
        after: List[Instruction] = []
        new_args: List[Operand] = []

        for arg in old_instr.defined():
            if isinstance(arg, Temporary):
                after.append(RiscV.sd(
                    S[3], arg.get_alloced_loc()
                ))
                new_args.append(S[3])
            else:
                new_args.append(arg)

        for i, arg in enumerate(old_instr.used(), 1):
            if isinstance(arg, Temporary):
                before.append(RiscV.ld(
                    S[i], arg.get_alloced_loc()
                ))
                new_args.append(S[i])
            else:
                new_args.append(arg)

        new_instr = old_instr.with_args(new_args)
        return before + [new_instr] + after

    def prepare(self) -> None:
        """Allocate all temporaries to memory.
        Invariants:
        - Expanded instructions can use s2 and s3
          (to store the values of temporaries before the actual instruction).
        """
        self._fdata._pool.set_temp_allocation(
            {temp: self._fdata.fresh_offset()
             for temp in self._fdata._pool.get_all_temps()})