from ScEpTIC.AST.elements.instructions.memory_operations import AllocaOperation
from ScEpTIC.AST.elements.instructions.other_operations import CallOperation


class RegisterSavingOptimization:
    """
    Identifies the minimum set of registes that each checkpoint need to save
    """

    def __init__(self, checkpoint_placer):
        self.checkpoint_placer = checkpoint_placer


    def set_checkpoints_registers(self):
        """
        Identifies and sets the registers that each checkpoints need to save
        """

        self.checkpoint_placer._split_function_bodies()

        for name in self.checkpoint_placer.function_body_slices.keys():
            self._reset_registers_info(name)
            self._set_checkpoints_general_purpose_regs(name)
            self._set_checkpoints_esp_reg(name)

        self.checkpoint_placer._synchronize_slices_to_body()


    def _reset_registers_info(self, name):
        """
        Resets the register-saving information associated to each checkpoint
        """
        slices = self.checkpoint_placer.function_body_slices[name]

        for slice in slices:
            checkpoints = self._get_checkpoints(slice)

            for checkpoint in checkpoints:
                checkpoint.checkpoint_save_pc = True
                checkpoint.checkpoint_save_esp = False
                checkpoint.checkpoint_save_regs = set()


    def _get_checkpoints(self, slice):
        """
        Returns a list of all the checkpoints included in the slices
        """
        instructions = []

        for instruction in slice:
            if isinstance(instruction, CallOperation) and instruction.name == self.checkpoint_placer.checkpoint_function_name:
                instructions.append(instruction)

        return instructions


    def _set_checkpoints_general_purpose_regs(self, name):
        """
        Sets the general purpose registers that each checkpoint in the function need to save
        """

        # account for function's argument registers
        n_args = len(self.checkpoint_placer.functions[name].arguments)

        slices = self.checkpoint_placer.function_body_slices[name]
        alloca_regs = self._get_alloca_regs(slices, n_args)

        alive_regs = set()

        # analyze slices in reverse order to identify general purpose registers to save
        for slice in reversed(slices):
            to_save = set()

            # retrieve regs infos
            regs_external_uses, regs_defs = self._get_regs_first_use_def(slice, alloca_regs)

            # add live set to registers to save
            for reg in alive_regs:
                to_save.add(reg)

            # if a defined reg is in alive_regs -> remove it from alive_regs
            for reg in regs_defs:
                if reg in alive_regs:
                    alive_regs.remove(reg)

            # add to alive_regs the used regs
            for reg in regs_external_uses:
                alive_regs.add(reg)

            self._set_checkpoints_save_regs(slice, to_save)


    def _get_alloca_regs(self, slices, n_args):
        """
        Returns the registers that are target of alloca instructions
        """
        alloca_regs = set()

        # accounts for registers that contain arguments
        for i in range(n_args):
            alloca_regs.add(f"%{i}")

        # order does not matter
        for slice in slices:
            for instruction in slice:
                if isinstance(instruction, AllocaOperation):
                    alloca_regs.add(instruction.target.value)

        return alloca_regs


    def _get_regs_first_use_def(self, slice, alloca_regs):
        """
        Returns two lists:
            - registers whose use happens before their defs
            - registers defined
        """
        reg_first_uses = set()
        reg_defs = set()

        for instruction in slice:
            for reg in instruction.get_uses():
                # reg not defined + ignore alloca regs
                if reg not in reg_defs and reg not in alloca_regs:
                    reg_first_uses.add(reg)

            for reg in instruction.get_defs():
                # ignore alloca regs
                if reg not in alloca_regs:
                    reg_defs.add(reg)

        return reg_first_uses, reg_defs


    def _set_checkpoints_save_regs(self, slice, to_save):
        """
        Sets the checkpoints contained in the slice to save the given registers

        Assumption (safe) -> no reg is used after a function call
        This is a safe assumption for the benchmarks we are considering, as we consider LLVM IR virtual regs
        """
        checkpoints = self._get_checkpoints(slice)

        for checkpoint in checkpoints:
            checkpoint.checkpoint_save_pc = True

            for reg in to_save:
                checkpoint.checkpoint_save_regs.add(reg)

            if len(to_save) > 0:
                checkpoint.checkpoint_save_esp = True


    def _set_checkpoints_esp_reg(self, name):
        """
        Sets which checkpoint in the function need to save the stack pointer
        """
        slices = self.checkpoint_placer.function_body_slices[name]

        # all checkpoints need to save stack pointer
        #for slice in slices:
        #    self._set_checkpoint_save_esp(slice)

        # first slice always need to save stack pointer
        self._set_checkpoint_save_esp(slices[0])

        # a checkpoint at the end of a slice need to save the stack pointer only if the slice calls a function with a checkpoint
        for slice in slices[1:]:
            for instruction in slice:
                # call to a function that is not a checkpoint
                if isinstance(instruction, CallOperation) and instruction.name != self.checkpoint_placer.checkpoint_function_name:
                    # the function contains a checkpoint, that is, saves the stack pointer
                    if self._function_has_checkpoint(instruction.name):
                        self._set_checkpoint_save_esp(slice)
                        # go to next slice
                        break


    def _set_checkpoint_save_esp(self, slice):
        """
        Sets the checkpoints contained in the given slice to save the stack pointer
        """
        checkpoints = self._get_checkpoints(slice)

        for checkpoint in checkpoints:
            checkpoint.checkpoint_save_esp = True


    def _function_has_checkpoint(self, name):
        """
        Returns if a given function contains a checkpoint
        """

        # builtin or input function
        if name not in self.checkpoint_placer.function_body_slices:
            return False

        slices = self.checkpoint_placer.function_body_slices[name]

        for slice in slices:
            for instruction in slice:
                if isinstance(instruction, CallOperation) and instruction.name == self.checkpoint_placer.checkpoint_function_name:
                    return True

        return False