perf: compare fork work from the divergence point and compute the state root incrementally

minichain/block.py

@@ -63,6 +63,9 @@ def __init__(
         self.hash: Optional[str] = None
         self.state_root: Optional[str] = state_root
         self.receipt_root: Optional[str] = receipt_root
+        # Difficulty EMA after this block. Local bookkeeping set when the block is
+        # accepted: not part of the header, the hash, or the wire format.
+        self.avg_block_time: Optional[float] = None
 
         # Compute merkle roots once
         self.merkle_root: Optional[str] = _calculate_merkle_root(self.transactions)

minichain/chain.py

@@ -168,6 +168,12 @@ def _next_target(self, target, avg_block_time):
 
         return max(MIN_TARGET, min(MAX_TARGET, new_target))
 
+    def _difficulty_after(self, block):
+        """The (target, avg_block_time) in effect for the block that follows `block`."""
+        if block.index == 0:
+            return block.target, self.target_block_time
+        return self._next_target(block.target, block.avg_block_time), block.avg_block_time
+
     def _apply_block(self, prev_block, block, state, target, avg_block_time):
         """
         Canonical block-application pipeline shared by add_block and resolve_conflicts.
@@ -240,6 +246,7 @@ def add_block(self, block):
             self.state = temp_state
             self.current_target = new_target
             self.avg_block_time = new_avg
+            block.avg_block_time = new_avg
             self.chain.append(block)
 
             self._state_snapshots.append((block.hash, self.state.snapshot()))
@@ -249,7 +256,8 @@ def add_block(self, block):
     def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
         """
         Evaluates a competing partial or full chain. If it has strictly greater cumulative work,
-        attempts a reorg. Rebuilds state from genesis to guarantee validity.
+        attempts a reorg. Re-executes the new blocks from the state at the fork point (or from
+        genesis when no snapshot of the fork point is kept) to guarantee validity.
         Returns: (success_bool, list_of_orphaned_transactions)
         """
         from .validators import ValidationStatus
@@ -258,91 +266,96 @@ def resolve_conflicts(self, new_chain_list) -> tuple[bool, list]:
             return False, []
 
         with self._lock:
-            first_block = new_chain_list[0]
+            # Forks almost always share a long prefix and differ only near the tip, so find
+            # the split by walking backward from the incoming tip to the last block we
+            # already have. Blocks up to that one are taken from our own chain, never from
+            # the peer, so the cost here depends on the fork length, not the chain height.
+            diverge = len(new_chain_list)
+            while diverge > 0:
+                b = new_chain_list[diverge - 1]
+                if 0 <= b.index < len(self.chain) and self.chain[b.index].hash == b.hash:
+                    break
+                diverge -= 1
+            if diverge == len(new_chain_list):
+                logger.debug("Incoming chain adds nothing beyond our own. Rejecting.")
+                return False, []
+
+            new_blocks = new_chain_list[diverge:]
+            first_block = new_blocks[0]
             fork_idx = first_block.index
-
-            if fork_idx == 0:
-                if first_block.hash != self.chain[0].hash:
-                    logger.warning("Reorg failed: Genesis hash mismatch.")
-                    return False, []
-            else:
-                if fork_idx > len(self.chain):
-                    logger.warning("Reorg failed: Partial chain does not connect to our history.")
-                    return False, []
-                if first_block.previous_hash != self.chain[fork_idx - 1].hash:
-                    logger.warning("Reorg failed: Partial chain hash mismatch at fork point.")
-                    return False, []
 
-            proposed_chain = self.chain[:fork_idx] + new_chain_list
+            if fork_idx <= 0:
+                logger.warning("Reorg failed: Genesis hash mismatch.")
+                return False, []
+            if fork_idx > len(self.chain):
+                logger.warning("Reorg failed: Partial chain does not connect to our history.")
+                return False, []
+            if first_block.previous_hash != self.chain[fork_idx - 1].hash:
+                logger.warning("Reorg failed: Partial chain hash mismatch at fork point.")
+                return False, []
 
             # Fast PoW Check: Ensure incoming blocks actually satisfy their declared target
             # before we trust their target to calculate total_work and rebuild state.
             from .network_config import MAX_TARGET
-            for b in new_chain_list:
+            for b in new_blocks:
                 if not isinstance(b.target, int) or b.target <= 0 or b.target > MAX_TARGET:
                     logger.warning("Reorg failed: Fast PoW check failed for block %s (invalid target)", b.index)
                     return False, []
                 if int(b.hash, 16) >= b.target:
                     logger.warning("Reorg failed: Fast PoW check failed for block %s (hash >= target)", b.index)
                     return False, []
 
-            current_work = self.get_total_work()
-            new_work = self.get_total_work(proposed_chain)
+            # The shared prefix adds the same work to both sides, so compare only what follows it.
+            old_blocks = self.chain[fork_idx:]
+            current_work = self.get_total_work(old_blocks)
+            new_work = self.get_total_work(new_blocks)
 
             if new_work <= current_work:
                 logger.debug("Incoming chain (work: %s) is not heavier than local chain (work: %s). Rejecting.", new_work, current_work)
                 return False, []
 
             logger.info("Incoming chain is heavier (%s > %s). Attempting reorg...", new_work, current_work)
 
-            original_chain = list(self.chain)
-
             temp_state = State()
             temp_state.chain_id = self.chain_id
 
-            fork_base_hash = self.chain[fork_idx - 1].hash if fork_idx > 0 else None
+            fork_base_hash = self.chain[fork_idx - 1].hash
 
-            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
-
+            snapshot_found = next((snap for h, snap in self._state_snapshots if h == fork_base_hash), None)
+
+            # Resume from the fork base if we kept its state, otherwise from genesis. The
+            # difficulty after the base is read from the base block itself.
             if snapshot_found is not None:
                 logger.info("Reorg optimization: Restoring state from in-memory snapshot at block %s", fork_idx - 1)
+                base_idx = 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:
+                base_idx = 0
                 temp_state.restore(self._genesis_state_snapshot)
-                start_idx = 1
+            temp_target, temp_avg_block_time = self._difficulty_after(self.chain[base_idx])
+            replay = self.chain[base_idx:fork_idx] + new_blocks
 
-            for i in range(start_idx, len(proposed_chain)):
+            for i in range(1, len(replay)):
                 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
+                    replay[i - 1], replay[i], temp_state, temp_target, 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)
+                    logger.warning("Reorg failed at block %s", replay[i].index)
                     return False, []
+                replay[i].avg_block_time = temp_avg_block_time
 
-            old_txs = {tx.tx_id: tx for b in original_chain[1:] for tx in b.transactions}
-            new_tx_ids = {tx.tx_id for b in proposed_chain[1:] for tx in b.transactions}
-            orphans = [tx for tx_id, tx in old_txs.items() if tx_id not in new_tx_ids]
+            # Only blocks we drop can orphan a transaction; the shared prefix is in both chains.
+            new_tx_ids = {tx.tx_id for b in new_blocks for tx in b.transactions}
+            orphans = list({
+                tx.tx_id: tx for b in old_blocks for tx in b.transactions if tx.tx_id not in new_tx_ids
+            }.values())
 
-            self.chain = proposed_chain
+            # A single slice assignment: it touches only the replaced tail, and readers never
+            # see the chain half-updated.
+            self.chain[fork_idx:] = new_blocks
             self.state = temp_state
             self.current_target = temp_target
             self.avg_block_time = temp_avg_block_time

minichain/mpt.py

@@ -2,13 +2,24 @@
 from trie import HexaryTrie
 
 
+class _Overlay(dict):
+    """Copy-on-write node store: writes stay local, reads fall through to `base`."""
+
+    def __init__(self, base):
+        super().__init__()
+        self.base = base
+
+    def __missing__(self, key):
+        return self.base[key]
+
+
 class Trie:
     """
     A Merkle Patricia Trie (MPT) for MiniChain backed by the `trie` library.
     Provides O(log N) state verification via cryptographic state roots.
     """
-    def __init__(self):
-        self._trie = HexaryTrie({})
+    def __init__(self, _trie=None):
+        self._trie = _trie if _trie is not None else HexaryTrie({})
 
     def root_hash(self) -> str:
         """Returns the 32-byte hex hash of the trie's root."""
@@ -22,3 +33,17 @@ def get(self, key_hex: str) -> Optional[str]:
     def put(self, key_hex: str, value: str):
         key = bytes.fromhex(key_hex)
         self._trie.set(key, value.encode())
+
+    def delete(self, key_hex: str):
+        try:
+            key = bytes.fromhex(key_hex)
+        except ValueError:
+            return  # not valid hex, so it was never stored
+        self._trie.delete(key)
+
+    def fork(self) -> "Trie":
+        """A trie at the same root whose later changes never touch this one (O(1))."""
+        return Trie(HexaryTrie(_Overlay(self._trie.db), self._trie.root_hash))
+
+    def node_count(self) -> int:
+        return len(self._trie.db)

minichain/state.py

@@ -22,8 +22,9 @@ def __init__(self, backing_dict):
     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])
+                # Read without marking the backing entry changed: we only keep a copy of it.
+                raw = dict.__getitem__ if isinstance(self.backing, dict) else type(self.backing).__getitem__
+                self.cache[key] = copy.deepcopy(raw(self.backing, key))
             else:
                 raise KeyError(key)
         return self.cache[key]
@@ -71,24 +72,126 @@ def rollback(self):
         self.cache.clear()
 
 
+def _is_empty_account(acc):
+    return acc.get('balance', 0) == 0 and acc.get('nonce', 0) == 0 and not acc.get('code') and not acc.get('storage')
+
+
+def _sync_leaf(trie, address, acc):
+    """Make `trie` reflect `acc` for `address` (empty accounts are not part of the state)."""
+    if acc is None or _is_empty_account(acc):
+        trie.delete(address)
+    else:
+        trie.put(address, json.dumps(acc, sort_keys=True))
+
+
+class AccountStore(dict):
+    """
+    The live accounts dict. It remembers which accounts may have changed since its
+    state trie was last brought up to date, so the state root only costs work
+    proportional to the accounts a block touched, not to all accounts.
+
+    Any access that hands out an account marks it as possibly changed, so callers
+    must fetch an account from the store before mutating it (not keep a stale reference).
+    """
+    # The trie never forgets replaced nodes; rebuild it once they outnumber live ones.
+    _MAX_NODES_PER_ACCOUNT = 6
+    _NODE_SLACK = 1000
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._trie = None
+        self._dirty = set()
+
+    def __getitem__(self, key):
+        value = super().__getitem__(key)
+        self._dirty.add(key)
+        return value
+
+    def get(self, key, default=None):
+        return self[key] if key in self else default
+
+    def __setitem__(self, key, value):
+        self._dirty.add(key)
+        super().__setitem__(key, value)
+
+    def __delitem__(self, key):
+        self._dirty.add(key)
+        super().__delitem__(key)
+
+    def setdefault(self, key, default=None):
+        if key not in self:
+            self[key] = default
+        return self[key]
+
+    def pop(self, key, *default):
+        self._dirty.add(key)
+        return super().pop(key, *default)
+
+    def update(self, *args, **kwargs):
+        for key, value in dict(*args, **kwargs).items():
+            self[key] = value
+
+    def values(self):
+        self._trie = None  # hands out every account at once; cannot track what the caller does
+        return super().values()
+
+    def items(self):
+        self._trie = None
+        return super().items()
+
+    def __deepcopy__(self, memo):
+        clone = AccountStore()
+        for key, value in dict.items(self):
+            dict.__setitem__(clone, key, copy.deepcopy(value, memo))
+        return clone
+
+    def synced_trie(self):
+        """The state trie, brought up to date with every account changed since the last call."""
+        if self._trie is None or self._trie.node_count() > self._MAX_NODES_PER_ACCOUNT * len(self) + self._NODE_SLACK:
+            self._trie = Trie()
+            self._dirty = set(dict.keys(self))
+        for address in self._dirty:
+            _sync_leaf(self._trie, address, dict.get(self, address))
+        self._dirty.clear()
+        return self._trie
+
+
 class State:
     def __init__(self):
         # { address: {'balance': int, 'nonce': int, 'code': str|None, 'storage': dict} }
-        self.accounts = {}
+        self.accounts = AccountStore()
         self.contract_machine = ContractMachine(self)
         self.chain_id = "minichain-default"
 
+    @property
+    def accounts(self):
+        return self._accounts
+
+    @accounts.setter
+    def accounts(self, value):
+        self._accounts = value if isinstance(value, (AccountStore, StateJournal)) else AccountStore(value)
+
     def state_root(self) -> str:
         """
-        Dynamically builds the Merkle Patricia Trie from the current state dictionary
-        and returns the cryptographic state root hash.
+        Returns the Merkle Patricia Trie root of the current state. The trie is kept
+        between calls and only the accounts that changed are updated; a state that is
+        a journal over the live store forks that trie and applies just its own changes.
         """
+        accounts = self.accounts
+        if isinstance(accounts, StateJournal) and isinstance(accounts.backing, AccountStore):
+            trie = accounts.backing.synced_trie().fork()
+            for address, acc in accounts.cache.items():
+                _sync_leaf(trie, address, acc)
+            return trie.root_hash()
+        if isinstance(accounts, AccountStore):
+            return accounts.synced_trie().root_hash()
+        return self._full_state_root()
+
+    def _full_state_root(self) -> str:
+        """Rebuild the trie from every account. Reference implementation for state_root()."""
         trie = Trie()
-        # Sort items to ensure deterministic insertion order if necessary (though MPT is order-independent)
-        for addr, acc in sorted(self.accounts.items()):
-            if acc.get('balance', 0) == 0 and acc.get('nonce', 0) == 0 and not acc.get('code') and not acc.get('storage'):
-                continue
-            trie.put(addr, json.dumps(acc, sort_keys=True))
+        for addr, acc in sorted(dict(self.accounts.items()).items()):
+            _sync_leaf(trie, addr, acc)
         return trie.root_hash()
 
     DEFAULT_MINING_REWARD = DEFAULT_MINING_REWARD