diff --git a/minichain/chain.py b/minichain/chain.py index 0b9e74f..98e164e 100644 --- a/minichain/chain.py +++ b/minichain/chain.py @@ -51,6 +51,9 @@ def __init__(self, genesis_path="genesis.json"): self.state = State() self.chain_id = "minichain-default" self._lock = threading.RLock() + import collections + 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): @@ -131,6 +134,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.append((genesis_block.hash, self.state.snapshot())) @property def last_block(self): @@ -237,7 +241,9 @@ def add_block(self, block): self.current_target = new_target self.avg_block_time = new_avg self.chain.append(block) - + + self._state_snapshots.append((block.hash, self.state.snapshot())) + return ValidationStatus.VALID def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]: @@ -293,12 +299,35 @@ 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_target = proposed_chain[0].target temp_avg_block_time = self.target_block_time + + 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(snapshot_found) + + # Fast forward target 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_target = self._next_target(temp_target, 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_target, temp_avg_block_time = self._apply_block( proposed_chain[i - 1], proposed_chain[i], temp_state, temp_target, temp_avg_block_time ) @@ -317,6 +346,9 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]: self.state = temp_state self.current_target = temp_target self.avg_block_time = temp_avg_block_time - + + # Repopulate snapshots for the new chain tip + 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 db4b07f..784b54e 100644 --- a/minichain/node_config.py +++ b/minichain/node_config.py @@ -23,3 +23,6 @@ # Keeping the upper limit around 32-bits ensures the nonce string in the JSON block # doesn't become unnecessarily large, and avoids cross-language serialization issues. MINING_INITIAL_NONCE_MAX = 2**32 - 1 + +# State Config +MAX_STATE_SNAPSHOTS = 10 # Number of recent block states to keep in memory for reorg optimization 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):