test ledger state is properly updated on re-org

cryptarchia/cryptarchia.py

@@ -26,8 +26,13 @@ class TimeConfig:
 @dataclass
 class Config:
     k: int
+    active_slot_coeff: float  # 'f', the rate of occupied slots
     time: TimeConfig
 
+    @property
+    def s(self):
+        return int(3 * self.k / self.active_slot_coeff)
+
 
 # An absolute unique indentifier of a slot, counting incrementally from 0
 @dataclass
@@ -158,6 +163,15 @@ class LedgerState:
     commitments: set[Id] = field(default_factory=set)  # set of commitments
     nullifiers: set[Id] = field(default_factory=set)  # set of nullified
 
+    def copy(self):
+        return LedgerState(
+            block=self.block,
+            nonce=self.nonce,
+            total_stake=self.total_stake,
+            commitments=self.commitments.copy(),
+            nullifiers=self.nullifiers.copy(),
+        )
+
     def verify_committed(self, commitment: Id) -> bool:
         return commitment in self.commitments
 
@@ -180,7 +194,7 @@ class Follower:
             nonce_snapshot=genesis_state,
         )
         self.genesis_state = genesis_state
-        self.ledger_state = genesis_state
+        self.ledger_state = genesis_state.copy()
 
     def validate_header(self, block: BlockHeader) -> bool:
         # TODO: this is not the full block validation spec, only slot leader is verified
@@ -201,16 +215,20 @@ class Follower:
             return True
 
         for chain in self.forks:
-            if chain.tip().id() == block.parent():
+            if chain.tip().id() == block.parent:
                 chain.blocks.append(block)
                 return True
 
         return False
 
     def try_create_fork(self, block: BlockHeader) -> Optional[Chain]:
+        if self.genesis_state.block == block.parent:
+            # this block is forking off the genesis state
+            return Chain(blocks=[block])
+
         chains = self.forks + [self.local_chain]
         for chain in chains:
-            if self.chain.contains_block(block):
+            if chain.contains_block(block):
                 block_position = chain.block_position(block)
                 return Chain(blocks=chain.blocks[:block_position] + [block])
 
@@ -234,7 +252,7 @@ class Follower:
                 return
 
         # We may need to switch forks, lets run the fork choice rule to check.
-        new_chain = Follower.fork_choice(self.local_chain, self.forks)
+        new_chain = self.fork_choice()
 
         if new_chain == self.local_chain:
             # we have not re-org'd therefore we can simply update our ledger state
@@ -252,14 +270,14 @@ class Follower:
             self.local_chain = new_chain
 
     # Evaluate the fork choice rule and return the block header of the block that should be the head of the chain
-    @staticmethod
+    def fork_choice(self) -> Chain:
-    def fork_choice(local_chain: Chain, forks: List[Chain]) -> Chain:
+        return maxvalid_bg(
-        # TODO: define k and s
+            self.local_chain, self.forks, k=self.config.k, s=self.config.s
-        return maxvalid_bg(local_chain, forks, 0, 0)
+        )
 
     def tip_id(self) -> Id:
         if self.local_chain.length() > 0:
-            return self.local_chain.tip().id
+            return self.local_chain.tip().id()
         else:
             return self.ledger_state.block
 
@@ -286,11 +304,6 @@ class EpochState:
         return self.nonce_snapshot.nonce
 
 
-@dataclass
-class LeaderConfig:
-    active_slot_coeff: float = 0.05  # 'f', the rate of occupied slots
-
-
 def phi(f: float, alpha: float) -> float:
     """
     params:
@@ -322,7 +335,7 @@ class MOCK_LEADER_VRF:
 
 @dataclass
 class Leader:
-    config: LeaderConfig
+    config: Config
     coin: Coin
 
     def try_prove_slot_leader(

cryptarchia/test_leader.py

@@ -2,7 +2,7 @@ from unittest import TestCase
 
 import numpy as np
 
-from .cryptarchia import Leader, LeaderConfig, EpochState, LedgerState, Coin, phi
+from .cryptarchia import Leader, Config, EpochState, LedgerState, Coin, phi, TimeConfig
 
 
 class TestLeader(TestCase):
@@ -15,8 +15,12 @@ class TestLeader(TestCase):
         )
 
         f = 0.05
-        leader_config = LeaderConfig(active_slot_coeff=f)
+        config = Config(
-        l = Leader(config=leader_config, coin=Coin(pk=0, value=10))
+            k=10,
+            active_slot_coeff=f,
+            time=TimeConfig(slots_per_epoch=1000, slot_duration=1, chain_start_time=0),
+        )
+        l = Leader(config=config, coin=Coin(pk=0, value=10))
 
         # We'll use the Margin of Error equation to decide how many samples we need.
         # https://en.wikipedia.org/wiki/Margin_of_error

cryptarchia/test_ledger_state_update.py

@@ -11,36 +11,48 @@ from .cryptarchia import (
     LedgerState,
     MockLeaderProof,
     Slot,
+    Id,
 )
 
 
+def mk_genesis_state(initial_stake_distribution: list[Coin]) -> LedgerState:
+    return LedgerState(
+        block=bytes(32),
+        nonce=bytes(32),
+        total_stake=sum(c.value for c in initial_stake_distribution),
+        commitments={c.commitment() for c in initial_stake_distribution},
+        nullifiers=set(),
+    )
+
+
+def mk_block(parent: Id, slot: int, coin: Coin, content=bytes(32)) -> BlockHeader:
+    from hashlib import sha256
+
+    return BlockHeader(
+        slot=Slot(slot),
+        parent=parent,
+        content_size=len(content),
+        content_id=sha256(content).digest(),
+        leader_proof=MockLeaderProof.from_coin(coin),
+    )
+
+
 def config() -> Config:
     return Config(
-        k=10, time=TimeConfig(slots_per_epoch=1000, slot_duration=1, chain_start_time=0)
+        k=10,
+        active_slot_coeff=0.05,
+        time=TimeConfig(slots_per_epoch=1000, slot_duration=1, chain_start_time=0),
     )
 
 
 class TestLedgerStateUpdate(TestCase):
     def test_ledger_state_prevents_coin_reuse(self):
         leader_coin = Coin(pk=0, value=100)
-        genesis_state = LedgerState(
+        genesis = mk_genesis_state([leader_coin])
-            block=bytes(32),
-            nonce=bytes(32),
-            total_stake=leader_coin.value,
-            commitments={leader_coin.commitment()},
-            nullifiers=set(),
-        )
 
-        follower = Follower(genesis_state, config())
+        follower = Follower(genesis, config())
-
-        block = BlockHeader(
-            slot=Slot(0),
-            parent=genesis_state.block,
-            content_size=1,
-            content_id=bytes(32),
-            leader_proof=MockLeaderProof.from_coin(leader_coin),
-        )
 
+        block = mk_block(slot=0, parent=genesis.block, coin=leader_coin)
         follower.on_block(block)
 
         # Follower should have accepted the block
@@ -50,18 +62,47 @@ class TestLedgerStateUpdate(TestCase):
         # Follower should have updated their ledger state to mark the leader coin as spent
         assert follower.ledger_state.verify_unspent(leader_coin.nullifier()) == False
 
-        reuse_coin_block = BlockHeader(
+        reuse_coin_block = mk_block(slot=1, parent=block.id, coin=leader_coin)
-            slot=Slot(0),
-            parent=block.id(),
-            content_size=1,
-            content_id=bytes(32),
-            leader_proof=MockLeaderProof(
-                commitment=leader_coin.commitment(),
-                nullifier=leader_coin.nullifier(),
-            ),
-        )
         follower.on_block(block)
 
         # Follower should *not* have accepted the block
         assert follower.local_chain.length() == 1
         assert follower.local_chain.tip() == block
+
+    def test_ledger_state_is_properly_updated_on_reorg(self):
+        coin_1 = Coin(pk=0, value=100)
+        coin_2 = Coin(pk=1, value=100)
+        coin_3 = Coin(pk=1, value=100)
+
+        genesis = mk_genesis_state([coin_1, coin_2, coin_3])
+
+        follower = Follower(genesis, config())
+
+        # 1) coin_1 & coin_2 both concurrently win slot 0
+
+        block_1 = mk_block(parent=genesis.block, slot=0, coin=coin_1)
+        block_2 = mk_block(parent=genesis.block, slot=0, coin=coin_2)
+
+        # 2) follower sees block 1 first
+
+        follower.on_block(block_1)
+        assert follower.tip_id() == block_1.id()
+        assert not follower.ledger_state.verify_unspent(coin_1.nullifier())
+
+        # 3) then sees block 2, but sticks with block_1 as the tip
+
+        follower.on_block(block_2)
+        assert follower.tip_id() == block_1.id()
+        assert len(follower.forks) == 1, f"{len(follower.forks)}"
+
+        # 4) then coin_3 wins slot 1 and chooses to extend from block_2
+
+        block_3 = mk_block(parent=block_2.id(), slot=0, coin=coin_3)
+
+        follower.on_block(block_3)
+
+        # the follower should have switched over to the block_2 fork
+        assert follower.tip_id() == block_3.id()
+
+        # and the original coin_1 should now be removed from the spent pool
+        assert follower.ledger_state.verify_unspent(coin_1.nullifier())