From 4e3387aa97194b58374eb191d422a16cc8dcb5e9 Mon Sep 17 00:00:00 2001 From: siddhant Date: Tue, 4 Aug 2026 00:14:43 +0530 Subject: [PATCH 1/2] feat: Implement StateJournal and in-memory snapshots for state rollback --- minichain/chain.py | 43 ++++++++++++++++++++-- minichain/state.py | 90 ++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 120 insertions(+), 13 deletions(-) diff --git a/minichain/chain.py b/minichain/chain.py index cf46ce2..3a3d3b2 100644 --- a/minichain/chain.py +++ b/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 diff --git a/minichain/state.py b/minichain/state.py index 0d80764..9cc0e2d 100644 --- a/minichain/state.py +++ b/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): From 91a0e8727fcdaba5656d40d25b16779b84165ad4 Mon Sep 17 00:00:00 2001 From: siddhant Date: Fri, 7 Aug 2026 01:24:15 +0530 Subject: [PATCH 2/2] refactor: use deque for state snapshots and configurable max snapshots size --- minichain/chain.py | 25 ++++++++++++++----------- minichain/node_config.py | 3 +++ 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/minichain/chain.py b/minichain/chain.py index 3a3d3b2..c865435 100644 --- a/minichain/chain.py +++ b/minichain/chain.py @@ -50,8 +50,8 @@ def __init__(self, genesis_path="genesis.json"): self.chain_id = "minichain-default" self._lock = threading.RLock() import collections - self._state_snapshots = collections.OrderedDict() - self._max_snapshots = 10 + from .node_config import MAX_STATE_SNAPSHOTS + self._state_snapshots = collections.deque(maxlen=MAX_STATE_SNAPSHOTS) self._create_genesis_block(genesis_path) def _create_genesis_block(self, genesis_path): @@ -124,7 +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() + self._state_snapshots.append((genesis_block.hash, self.state.snapshot())) @property def last_block(self): @@ -226,9 +226,7 @@ def add_block(self, block): 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) + self._state_snapshots.append((block.hash, self.state.snapshot())) return ValidationStatus.VALID @@ -280,9 +278,16 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]: 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: + snapshot_found = None + if fork_base_hash: + for h, snap in self._state_snapshots: + if h == fork_base_hash: + snapshot_found = snap + break + + if snapshot_found is not None: 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]) + temp_state.restore(snapshot_found) # Fast forward difficulty and avg_block_time without executing state for i in range(1, fork_idx): @@ -316,9 +321,7 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]: 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) + self._state_snapshots.append((self.last_block.hash, self.state.snapshot())) logger.info("Reorg successful! Switched to new chain tip: Block %s", self.last_block.index) return True, orphans diff --git a/minichain/node_config.py b/minichain/node_config.py index 68154a3..4effeab 100644 --- a/minichain/node_config.py +++ b/minichain/node_config.py @@ -14,3 +14,6 @@ # Mining Config MINING_MAX_NONCE = 10_000_000 # Number of hashes to attempt before yielding the mining thread + +# State Config +MAX_STATE_SNAPSHOTS = 10 # Number of recent block states to keep in memory for reorg optimization