feat: Implement StateJournal and in-memory snapshots for state rollback

Author: SIDDHANTCOOKIE

minichain/chain.py

@@ -49,6 +49,9 @@ def __init__(self, genesis_path="genesis.json"):
         self.state = State()
         self.chain_id = "minichain-default"
         self._lock = threading.RLock()
+        import collections
+        self._state_snapshots = collections.OrderedDict()
+        self._max_snapshots = 10
         self._create_genesis_block(genesis_path)
 
     def _create_genesis_block(self, genesis_path):
@@ -121,6 +124,7 @@ def _create_genesis_block(self, genesis_path):
 
         # Snapshot the state exactly after genesis allocation for clean reorg rebuilds
         self._genesis_state_snapshot = self.state.snapshot()
+        self._state_snapshots[genesis_block.hash] = self.state.snapshot()
 
     @property
     def last_block(self):
@@ -214,10 +218,18 @@ def add_block(self, block):
                 return status
 
             # All transactions valid → commit state and append block
+            if hasattr(temp_state.accounts, 'commit'):
+                temp_state.accounts.commit()
+                temp_state.accounts = temp_state.accounts.backing
             self.state = temp_state
             self.current_difficulty = new_difficulty
             self.avg_block_time = new_avg
             self.chain.append(block)
+            
+            self._state_snapshots[block.hash] = self.state.snapshot()
+            while len(self._state_snapshots) > self._max_snapshots:
+                self._state_snapshots.popitem(last=False)
+            
             return ValidationStatus.VALID
 
     def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
@@ -262,15 +274,34 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
 
             temp_state = State()
             temp_state.chain_id = self.chain_id
-            temp_state.restore(self._genesis_state_snapshot)
-
+            
+            fork_base_hash = self.chain[fork_idx - 1].hash if fork_idx > 0 else None
+            
             temp_difficulty = proposed_chain[0].difficulty
             temp_avg_block_time = self.target_block_time
+            
+            if fork_base_hash and fork_base_hash in self._state_snapshots:
+                logger.info("Reorg optimization: Restoring state from in-memory snapshot at block %s", fork_idx - 1)
+                temp_state.restore(self._state_snapshots[fork_base_hash])
+                
+                # Fast forward difficulty and avg_block_time without executing state
+                for i in range(1, fork_idx):
+                    block_time = proposed_chain[i].timestamp - proposed_chain[i-1].timestamp
+                    temp_avg_block_time = self.alpha * block_time + (1 - self.alpha) * temp_avg_block_time
+                    temp_difficulty = self._next_difficulty(temp_difficulty, temp_avg_block_time)
+                
+                start_idx = fork_idx
+            else:
+                temp_state.restore(self._genesis_state_snapshot)
+                start_idx = 1
 
-            for i in range(1, len(proposed_chain)):
+            for i in range(start_idx, len(proposed_chain)):
                 status, temp_difficulty, temp_avg_block_time = self._apply_block(
                     proposed_chain[i - 1], proposed_chain[i], temp_state, temp_difficulty, temp_avg_block_time
                 )
+                if hasattr(temp_state.accounts, 'commit'):
+                    temp_state.accounts.commit()
+                    temp_state.accounts = temp_state.accounts.backing
                 if status != ValidationStatus.VALID:
                     logger.warning("Reorg failed at block %s", proposed_chain[i].index)
                     return False, []
@@ -283,5 +314,11 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
             self.state = temp_state
             self.current_difficulty = temp_difficulty
             self.avg_block_time = temp_avg_block_time
+            
+            # Repopulate snapshots for the new chain tip
+            self._state_snapshots[self.last_block.hash] = self.state.snapshot()
+            while len(self._state_snapshots) > self._max_snapshots:
+                self._state_snapshots.popitem(last=False)
+                
             logger.info("Reorg successful! Switched to new chain tip: Block %s", self.last_block.index)
             return True, orphans

minichain/state.py

@@ -10,6 +10,66 @@
 
 logger = logging.getLogger(__name__)
 
+class StateJournal:
+    """
+    An in-memory proxy dictionary that caches reads and writes to avoid
+    expensive deep copies of the entire state dictionary during transactions.
+    """
+    def __init__(self, backing_dict):
+        self.backing = backing_dict
+        self.cache = {}
+
+    def __getitem__(self, key):
+        if key not in self.cache:
+            if key in self.backing:
+                import copy
+                self.cache[key] = copy.deepcopy(self.backing[key])
+            else:
+                raise KeyError(key)
+        return self.cache[key]
+
+    def __setitem__(self, key, value):
+        self.cache[key] = value
+
+    def __delitem__(self, key):
+        raise NotImplementedError("Account deletion not supported in StateJournal")
+
+    def __contains__(self, key):
+        return key in self.cache or key in self.backing
+
+    def get(self, key, default=None):
+        try:
+            return self.__getitem__(key)
+        except KeyError:
+            return default
+
+    def items(self):
+        res = self.backing.copy()
+        res.update(self.cache)
+        return res.items()
+
+    def update(self, other_dict):
+        if hasattr(other_dict, 'items'):
+            for k, v in other_dict.items():
+                self[k] = v
+        else:
+            for k, v in other_dict:
+                self[k] = v
+
+    def copy(self):
+        res = self.backing.copy()
+        res.update(self.cache)
+        return res
+
+    def commit(self):
+        """Flushes cached modifications to the backing dictionary."""
+        self.backing.update(self.cache)
+        self.cache.clear()
+
+    def rollback(self):
+        """Discards modifications."""
+        self.cache.clear()
+
 
 class State:
     def __init__(self):
@@ -69,9 +129,11 @@ def verify_transaction_logic(self, tx):
     def copy(self):
         """
         Return an independent copy of state for transactional validation.
+        Uses StateJournal for O(1) cloning instead of deepcopy.
         """
-        new_state = copy.deepcopy(self)
-        new_state.contract_machine = ContractMachine(new_state) # Reinitialize contract_machine
+        new_state = State()
+        new_state.accounts = StateJournal(self.accounts)
+        new_state.contract_machine = ContractMachine(new_state)
         new_state.chain_id = self.chain_id
         return new_state
 
@@ -124,22 +186,22 @@ def apply_transaction(self, tx):
 
 
     def _apply_validated_tx(self, tx):
+        original_accounts = self.accounts
+        journal = StateJournal(original_accounts)
+        self.accounts = journal
+
         sender = self.accounts[tx.sender]
         total_cost = tx.amount + (getattr(tx, 'gas_limit', 0) * getattr(tx, 'fee_per_gas', 0))
 
         sender['balance'] -= total_cost
         sender['nonce'] += 1
 
-        import copy
-        state_snapshot = copy.deepcopy(self.accounts)
-
         def rollback_and_refund(error_message, gas_used):
-            self.accounts = copy.deepcopy(state_snapshot)
+            journal.rollback()
+            self.accounts = original_accounts
             refund_acc = self.accounts[tx.sender]
-            refund_acc['balance'] += tx.amount
-            gas_refund = getattr(tx, 'gas_limit', 0) - gas_used
-            if gas_refund > 0:
-                refund_acc['balance'] += (gas_refund * getattr(tx, 'fee_per_gas', 0))
+            refund_acc['balance'] -= (gas_used * getattr(tx, 'fee_per_gas', 0))
+            refund_acc['nonce'] += 1
             return Receipt(tx.tx_id, status=0, error_message=error_message, gas_used=gas_used)
 
         # LOGIC BRANCH 1: Contract Deployment
@@ -162,6 +224,9 @@ def rollback_and_refund(error_message, gas_used):
             gas_refund = gas_used - code_gas
             if gas_refund > 0:
                 self.accounts[tx.sender]['balance'] += (gas_refund * getattr(tx, 'fee_per_gas', 0))
+            
+            journal.commit()
+            self.accounts = original_accounts
             return Receipt(tx.tx_id, status=1, contract_address=contract_address, gas_used=code_gas)
 
         # LOGIC BRANCH 2: Contract Call
@@ -187,12 +252,17 @@ def rollback_and_refund(error_message, gas_used):
             if gas_refund > 0:
                 self.accounts[tx.sender]['balance'] += (gas_refund * getattr(tx, 'fee_per_gas', 0))
 
+            journal.commit()
+            self.accounts = original_accounts
             return Receipt(tx.tx_id, status=1, gas_used=gas_used)
 
         # LOGIC BRANCH 3: Regular Transfer
         receiver = self.get_account(tx.receiver)
         receiver['balance'] += tx.amount
         gas_used = getattr(tx, 'gas_limit', 0)
+        
+        journal.commit()
+        self.accounts = original_accounts
         return Receipt(tx.tx_id, status=1, gas_used=gas_used)
 
     def execute_internal_call(self, sender, receiver_address, amount, payload, gas_limit, depth, is_top_level=False):