import copy

from ScEpTIC.AST.elements.instructions.other_operations import CallOperation
from ScEpTIC.AST.misc.virtual_memory_enum import VirtualMemoryEnum


def apply_transformation(functions, vmstate, f_declarations):
    """
    Sets the correct memory map for Mementos
    """
    checkpoint_function_name = vmstate.config.state_retention.state_save_function_name

    registers_to_save = set([x for x in range(0, vmstate.config.register_file.n_physical_registers)])

    for function in functions.values():
        for instruction in function.body:
            if "virtual_memory_target" in dir(instruction):
                instruction.virtual_memory_target = VirtualMemoryEnum.VOLATILE

            if isinstance(instruction, CallOperation) and instruction.name == checkpoint_function_name:
                instruction.checkpoint_save_pc = True
                instruction.checkpoint_save_esp = True
                instruction.checkpoint_save_ram = True
                instruction.checkpoint_save_regs = copy.deepcopy(registers_to_save)