from ScEpTIC.AST.elements.instructions.memory_operations import StoreOperation
from ScEpTIC.AST.misc.virtual_memory_enum import VirtualMemoryEnum
from ScEpTIC.AST.transformations import virtual_memory

def apply_transformation(functions, vmstate, f_declarations):
    """
    Sets all the memory operations to target NVM
    """

    virtual_memory.apply_transformation(functions, vmstate)

    checkpoint_function_name = vmstate.config.state_retention.state_save_function_name

    for function in functions.values():
        to_remove = []

        for instruction in function.body:
            # Set all to NVM
            if "virtual_memory_target" in dir(instruction):
                instruction.virtual_memory_target = VirtualMemoryEnum.NON_VOLATILE

            # Note: dummy writes resolve uncertainty and we need them to apply our versioning technique
            # Remove volatile copies
            if isinstance(instruction, StoreOperation):
                if instruction.is_virtual_memory_copy:
                    to_remove.append(instruction)

        for instruction in to_remove:
            function.body.remove(instruction)

        function.update_labels()