sceptic-biomem/ScEpTIC/emulator/memory/physical/physical_state_retention.py

import copy
import json

from ScEpTIC.emulator.memory import utils
from ScEpTIC.emulator.energy.memory.physical_memory_energy_model import PhysicalMemoryEnergyModel
from ScEpTIC.emulator.memory.physical.physical_cell_state_enum import PhysicalCellStateEnum
from ScEpTIC.emulator.memory.physical.physical_memory_array import PhysicalMemoryArray
from ScEpTIC.emulator.memory.physical.physical_memory_management_unit import PhysicalMemoryManagementUnit
from ScEpTIC.emulator.register_file.program_counter import ProgramCounter


class PhysicalStateRetention:
    """
    State retention on physical memory
    """

    def __init__(self, n_registers=12, register_size=4, gst_size=4096, stack_size=4096, version_bytes=1):
        """
        :param n_registers: number of registers
        :param register_size: size of registers in bytes
        :param gst_size: size of global symbol table in bytes
        :param stack_size: size of stack in bytes
        """
        self.version_bytes = version_bytes
        self.n_registers = n_registers
        self.register_size = register_size
        ProgramCounter._word_size = register_size
        self.gst_size = gst_size
        self.stack_size = stack_size
        self.state_size = self.n_registers * self.register_size + self.gst_size + self.stack_size

        self.write_counter = 0

        self.mmu = None
        self.energy_model = None

        state_map = {
            # 0 -> register_size-1
            'pc': 0,
            # register_size -> 2*register_size - 1
            'esp': self.register_size,
            # 2*register_size -> n_registers*register_size - 1
            'registers': [x for x in range(self.register_size*2, self.register_size*(2+(self.n_registers-2)), self.register_size)],
            # n_registers*register_size -> n_registers*register_size + gst_size - 1
            'gst': self.register_size * self.n_registers,
            # n_registers*register_size + gst_size
            'stack': self.register_size * self.n_registers + self.gst_size,
        }

        self.state_map = {
            'version': 0,
            'data': {
                0: {},
                1: {},
            }
        }

        self.dumps = {0: None, 1: None}

        for element, addr in state_map.items():
            for version in [0, 1]:
                offset = self.version_bytes + version * self.state_size

                if isinstance(addr, list):
                    self.state_map['data'][version][element] = [x+offset for x in addr]
                else:
                    self.state_map['data'][version][element] = addr + offset

    def init_memory(self, model_name, regulator_efficiency=0.8, read_before_write=True, wear_only_if_change=False, refresh_on_read=True, die_on_last_cycle=False):
        """
        Initializes the physical memory
        :param model_name: the memory model name
        :param regulator_efficiency: the voltage regulator efficiency for the MMU
        :param read_before_write: read data before write; the MMU writes only cells whose value will change
        """
        size = self.version_bytes + 2 * self.state_size
        name = 'STATE_RETENTION_MEMORY'

        self.energy_model = PhysicalMemoryEnergyModel(model_name, regulator_efficiency)
        self.energy_model.set_name(name)
        self.mmu = PhysicalMemoryManagementUnit(0, size, name, read_before_write)
        cell_array = PhysicalMemoryArray.from_energy_model(self.energy_model, size, wear_only_if_change, refresh_on_read, die_on_last_cycle)

        self.mmu.attach_cell_array(cell_array)
        self.mmu.attach_state_map(self.state_map)

    def __get_version(self):
        return int(self.mmu.memory_read(self.state_map['version'], self.version_bytes*8, sync_memory=False)[-1])

    def __update_version(self, version):
        self.mmu.memory_write(version, self.state_map['version'], self.version_bytes*8, sync_memory=False)
        if version != self.__get_version():
            self.mmu.log_data("VERSION_UPDATE", self.state_map['version'], self.version_bytes*8, "CHECKPOINT_VERSION", self.__get_version(), version)
        elif PhysicalMemoryArray.log_all_operations:
            self.mmu.log_data("VERSION_UPDATE", self.state_map['version'], self.version_bytes*8, "CHECKPOINT_VERSION", self.__get_version(), version)

    def __get_virtual_addr(self, address):
        # remove version offset
        address -= self.version_bytes

        if address >= self.state_size:
            address -= self.state_size

        return address

    def get_mmu(self):
        """
        Returns the mmu
        """
        return self.mmu

    def get_energy_model(self):
        """
        Returns the energy model
        """
        return self.energy_model

    def write(self, dump):
        """
        Execute the memory write of the saved state
        :param dump: state to write
        """
        self.mmu.log_data('STATE_SAVE_START', None, None, None, None, None)
        version = (self.__get_version() + 1) % 2
        self.dumps[version] = copy.deepcopy(dump)
        #print(f"Writing dump into version #{version}")
        self.__save_registers(dump, version)
        self.__save_stack(dump, version)
        self.__save_gst(dump, version)
        self.__update_version(version)
        self.__save_mem_json()
        self.mmu.log_data('STATE_SAVE_END', None, None, None, None, None)
        self.write_counter = self.write_counter + 1

    def read(self):
        """
        Execute the memory read of the saved state and synchronizes the dump state accordingly to the underlying physical memory state
        """
        self.mmu.log_data('STATE_RESTORE_START', None, None, None, None, None)
        version = self.__get_version()
        dump = copy.deepcopy(self.dumps[version])
        print(f"Reading dump from version #{version}")
        self.__restore_registers(dump, version)
        self.__restore_stack(dump, version)
        self.__restore_gst(dump, version)
        self.__trigger_restore_log()
        self.mmu.log_data('STATE_RESTORE_END', None, None, None, None, None)
        return dump

    def __save_registers(self, dump, version):
        """
        Save registers
        :param dump: state to write
        :param version: memory version
        """
        if 'register_file' not in dump or dump['register_file'] is None:
            return

        # get state map of current data version
        state_map = self.state_map['data'][version]

        # Save PC
        pc_address = state_map['pc']
        #pc_val = {'function_name': dump['register_file']['pc'].function_name, 'instruction_number': dump['register_file']['pc'].instruction_number}
        pc_val = dump['register_file']['pc'].get_address()
        self.mmu.memory_write(pc_val, pc_address, self.register_size*8, sync_memory=True, sync_addr=self.__get_virtual_addr(pc_address))

        # Save ESP
        esp_address = state_map['esp']
        self.mmu.memory_write(dump['stack']['top_address'], esp_address, self.register_size*8, sync_memory=True, sync_addr=self.__get_virtual_addr(esp_address))

        reg_id = 0
        for reg_stack_id in reversed(sorted(dump['register_file']['reg_tracker'].keys())):
            regs = dump['register_file']['reg_tracker'][reg_stack_id]['reg_list']
            for reg_name in regs:
                reg_val = dump['register_file']['reg_tracker'][reg_stack_id][reg_name].value
                reg_address = state_map['registers'][reg_id]
                self.mmu.memory_write(reg_val, reg_address, self.register_size*8, sync_memory=True, sync_addr=self.__get_virtual_addr(reg_address))
                reg_id += 1

                if reg_id >= (self.n_registers - 2):
                    return

    def __restore_registers(self, dump, version):
        """
        Restores registers
        :param dump: state to read
        :param version: memory version
        """
        if 'register_file' not in dump or dump['register_file'] is None:
            return

        # get state map of current data version
        state_map = self.state_map['data'][version]

        version_status = self.mmu.check_status(self.state_map['version'], self.version_bytes * 8)

        # Read PC and check if it is consistent
        pc_address = state_map['pc']
        pc = self.mmu.memory_read(pc_address, self.register_size*8, sync_memory=True, sync_addr=self.__get_virtual_addr(pc_address))
        #pc_val = {'function_name': dump['register_file']['pc'].function_name, 'instruction_number': dump['register_file']['pc'].instruction_number}
        pc_val = dump['register_file']['pc'].get_address()
        match, _, _ = utils.check_data(pc_val, pc)

        # Invalidate PC
        if not match:
            pc = utils.from_binary(pc, int)
            new_pc = dump['register_file']['pc'].get_pc_str_from_address(pc)

            expected_pc = {'function_name': dump['register_file']['pc'].function_name, 'instruction_number': dump['register_file']['pc'].instruction_number}

            status = self.mmu.check_status(pc_address, self.register_size*8)
            if status == PhysicalCellStateEnum.OK:
                status = f"Version status: {version_status.value}; Cell status: {status.value}"
            else:
                status = status.value

            print(f"State saved for PC degraded ({status}): {hex(pc)} ({new_pc}) instead of {hex(pc_val)} ({expected_pc})")
            self.mmu.log_data("RESTORE", pc_address, self.register_size*8, "CHECKPOINT_PC", pc, pc_val)

            dump['register_file']['pc'].function_name = new_pc['function_name']
            dump['register_file']['pc'].instruction_number = new_pc['instruction_number']
        elif PhysicalMemoryArray.log_all_operations:
            pc = utils.from_binary(pc, int)
            self.mmu.log_data("RESTORE", pc_address, self.register_size*8, "CHECKPOINT_PC", pc, pc)

        # Read ESP and check if it is consistent
        esp_address = state_map['esp']
        esp = self.mmu.memory_read(esp_address, self.register_size*8, sync_memory=True, sync_addr=self.__get_virtual_addr(esp_address))
        match, _, bits_val = utils.check_data(dump['stack']['top_address'], esp)

        # Update ESP
        if not match:
            status = self.mmu.check_status(esp_address, self.register_size*8)
            if status == PhysicalCellStateEnum.OK:
                status = f"Version status: {version_status.value}; Cell status: {status.value}"
            else:
                status = status.value

            print(f"State saved for EBP degraded ({status}): {bits_val} instead of {dump['stack']['top_address']}")
            self.mmu.log_data("RESTORE", esp_address, self.register_size*8, "CHECKPOINT_ESP", bits_val, dump['stack']['top_address'])
            dump['stack']['top_address'] = bits_val
        elif PhysicalMemoryArray.log_all_operations:
            self.mmu.log_data("RESTORE", esp_address, self.register_size*8, "CHECKPOINT_ESP", esp, esp)

        # Read registers and check if they are consistent
        reg_id = 0
        for reg_stack_id in reversed(sorted(dump['register_file']['reg_tracker'].keys())):
            regs = dump['register_file']['reg_tracker'][reg_stack_id]['reg_list']
            for reg_name in regs:
                reg_val = dump['register_file']['reg_tracker'][reg_stack_id][reg_name].value
                reg_address = state_map['registers'][reg_id]
                physical_reg_val = self.mmu.memory_read(reg_address, self.register_size*8, sync_memory=True, sync_addr=self.__get_virtual_addr(reg_address))
                match, _, bits_val = utils.check_data(reg_val, physical_reg_val)

                if not match:
                    status = self.mmu.check_status(reg_address, self.register_size*8)
                    if status == PhysicalCellStateEnum.OK:
                        status = f"Version status: {version_status.value}; Cell status: {status.value}"
                    else:
                        status = status.value

                    print(f"Saved state for register {reg_name} degraded ({status}): {bits_val} instead of {reg_val}")
                    self.mmu.log_data("RESTORE", reg_address, self.register_size*8, f"CHECKPOINT_{reg_name}", bits_val, reg_val)

                    if reg_stack_id == len(dump['register_file']['reg_tracker']):
                        dump['register_file']['registers'][reg_name].value = bits_val
                    else:
                        dump['register_file']['reg_stack'][reg_stack_id][reg_name].value = bits_val

                elif PhysicalMemoryArray.log_all_operations:
                    self.mmu.log_data("RESTORE", reg_address, self.register_size*8, f"CHECKPOINT_{reg_name}", reg_val, reg_val)

                reg_id += 1

                if reg_id >= (self.n_registers - 2):
                    return

    def __save_memory(self, address, memory_cells):
        """
        Save memory
        :param address: memory start address
        :param memory_cells: memory cells list
        """
        for memory_cell in memory_cells.values():
            val = memory_cell.get_content()
            size = memory_cell.get_bit_size()
            self.mmu.memory_write(val, address, size, sync_memory=True, sync_addr=self.__get_virtual_addr(address))

            # Update next address
            address += int(size / 8)

    def __restore_memory(self, address, memory_cells):
        """
        Restore memory and synchronizes the dump state accordingly to the underlying physical memory state
        :param address: memory start address
        :param memory_cells: memory cells list
        """
        version_status = self.mmu.check_status(self.state_map['version'], self.version_bytes * 8)

        for memory_cell in memory_cells.values():
            val = memory_cell.get_content()
            size = memory_cell.get_bit_size()
            physical_val = self.mmu.memory_read(address, size, sync_memory=True, sync_addr=self.__get_virtual_addr(address))

            match, _, bits_val = utils.check_data(val, physical_val)

            if not match:
                status = self.mmu.check_status(address, size)
                if status == PhysicalCellStateEnum.OK:
                    status = f"Version status: {version_status.value}; Cell status: {status.value}"
                else:
                    status = status.value

                print(f"State saved for memory cell {memory_cell.absolute_address} degraded ({status}): {bits_val} instead of {val}")
                self.mmu.log_data("RESTORE", address, size, f"{memory_cell.absolute_address}", bits_val, val)
                memory_cell.set_content(bits_val)
            elif PhysicalMemoryArray.log_all_operations:
                self.mmu.log_data("RESTORE", address, size, f"{memory_cell.absolute_address}", val, val)

            # Update next address
            address += int(size / 8)

    def __save_stack(self, dump, version):
        """
        Saves the stack
        :param dump: state to write
        :param version: memory version
        """
        if 'stack' not in dump or dump['stack'] is None:
            return

        # get state map of current data version
        state_map = self.state_map['data'][version]

        address = state_map['stack']
        self.__save_memory(address, dump['stack']['memory'])

    def __restore_stack(self, dump, version):
        """
        Restores the stack and synchronizes the dump state accordingly to the underlying physical memory state
        :param dump: state to read
        :param version: memory version
        """
        if 'stack' not in dump or dump['stack'] is None:
            return

        # get state map of current data version
        state_map = self.state_map['data'][version]

        address = state_map['stack']
        self.__restore_memory(address, dump['stack']['memory'])

    def __save_gst(self, dump, version):
        """
        Saves the gst
        :param dump: state to write
        :param version: memory version
        """
        if 'volatile_gst' not in dump or dump['volatile_gst'] is None:
            return

        # get state map of current data version
        state_map = self.state_map['data'][version]

        address = state_map['gst']
        self.__save_memory(address, dump['volatile_gst']['memory'])

    def __restore_gst(self, dump, version):
        """
        Restores the gst and synchronizes the dump state accordingly to the underlying physical memory state
        :param dump: state to read
        :param version: memory version
        """
        if 'volatile_gst' not in dump or dump['volatile_gst'] is None:
            return

        # get state map of current data version
        state_map = self.state_map['data'][version]

        address = state_map['gst']
        self.__restore_memory(address, dump['volatile_gst']['memory'])

    def __save_mem_json(self):
        self.mmu.write_save_trace(self.write_counter)

    def __trigger_restore_log(self):
        self.mmu.write_restore_trace(self.write_counter - 1)