Follower maintains ledger state as it follows the chain

cryptarchia/cryptarchia.py

@@ -23,6 +23,12 @@ class TimeConfig:
     chain_start_time: int
 
 
+@dataclass
+class Config:
+    k: int
+    time: TimeConfig
+
+
 # An absolute unique indentifier of a slot, counting incrementally from 0
 @dataclass
 class Slot:
@@ -37,9 +43,30 @@ class Slot:
 
 
 @dataclass
-class Config:
-    k: int
-    time: TimeConfig
+class Coin:
+    pk: int
+    value: int
+
+    def commitment(self) -> Id:
+        # TODO: mocked until CL is understood
+        pk_bytes = int.to_bytes(self.pk, length=32, byteorder="little")
+        value_bytes = int.to_bytes(self.value, length=32, byteorder="little")
+
+        h = sha256()
+        h.update(pk_bytes)
+        h.update(value_bytes)
+        return h.digest()
+
+    def nullifier(self) -> Id:
+        # TODO: mocked until CL is understood
+        pk_bytes = int.to_bytes(self.pk, length=32, byteorder="little")
+        value_bytes = int.to_bytes(self.value, length=32, byteorder="little")
+
+        h = sha256()
+        h.update(pk_bytes)
+        h.update(value_bytes)
+        h.update(b"\x00")  # extra 0 byte to differentiate from commitment
+        return h.digest()
 
 
 @dataclass
@@ -47,15 +74,17 @@ class MockLeaderProof:
     commitment: Id
     nullifier: Id
 
-    def verify(self):
+    @staticmethod
+    def from_coin(coin: Coin):
+        return MockLeaderProof(commitment=coin.commitment(), nullifier=coin.nullifier())
+
+    def verify(self, slot):
         # TODO: verification not implemented
         return True
 
     def _id_update(self, hasher):
-        commitment_bytes = int.to_bytes(self.commitment, length=32, byteorder="little")
-        nullifier_bytes = int.to_bytes(self.nullifier, length=32, byteorder="little")
-        hasher.update(commitment_bytes)
-        hasher.update(nullifier_bytes)
+        hasher.update(self.commitment)
+        hasher.update(self.nullifier)
 
 
 @dataclass
@@ -121,13 +150,21 @@ class LedgerState:
     """
 
     block: Id = None
-    nonce: bytes = None
+    nonce: Id = None
     total_stake: int = None
     commitments: set[Id] = field(default_factory=set)  # set of commitments
     nullifiers: set[Id] = field(default_factory=set)  # set of nullified
 
-    def is_coin_nullified(self, nullifier: Id) -> bool:
-        return nullifier in self.nullifiers
+    def verify_committed(self, commitment: Id) -> bool:
+        return commitment in self.commitments
+
+    def verify_unspent(self, nullifier: Id) -> bool:
+        return nullifier not in self.nullifiers
+
+    def apply(self, block: BlockHeader):
+        assert block.parent == self.block
+        self.block = block.id()
+        self.nullifiers.add(block.leader_proof.nullifier)
 
 
 class Follower:
@@ -139,23 +176,24 @@ class Follower:
             stake_distribution_snapshot=genesis_state,
             nonce_snapshot=genesis_state,
         )
+        self.ledger_state_snapshot = genesis_state
         self.ledger_state = genesis_state
 
-    def validate_header(block: BlockHeader) -> bool:
+    def validate_header(self, block: BlockHeader) -> bool:
         # TODO: this is not the full block validation spec, only slot leader is verified
         return self.verify_slot_leader(block.slot, block.leader_proof)
 
     def verify_slot_leader(self, slot: Slot, proof: MockLeaderProof) -> bool:
         return (
-            proof.verify(slot) # verify slot leader proof
-            and self.epoch.is_coin_old_enough_to_lead(proof.coin) # verify coin was not recently created
-            and not self.ledger_state.is_coin_nullified(proof.nullifier) # verify the coin has not already been spent
+            proof.verify(slot)  # verify slot leader proof
+            and self.epoch.verify_commitment_is_old_enough_to_lead(proof.commitment)
+            and self.ledger_state.verify_unspent(proof.nullifier)
         )
 
     # Try appending this block to an existing chain and return whether
     # the operation was successful
     def try_extend_chains(self, block: BlockHeader) -> bool:
-        if self.local_chain.tip().id() == block.parent():
+        if self.tip_id() == block.parent:
             self.local_chain.blocks.append(block)
             return True
 
@@ -180,50 +218,47 @@ class Follower:
             return
 
         # check if the new block extends an existing chain
-        if self.try_extend_chains(block):
-            return
+        succeeded_in_extending_a_chain = self.try_extend_chains(block)
+        if not succeeded_in_extending_a_chain:
+            # we failed to extend one of the existing chains,
+            # therefore we might need to create a new fork
+            new_chain = self.try_create_fork(block)
+            if new_chain is not None:
+                self.forks.append(new_chain)
+            else:
+                # otherwise, we're missing the parent block
+                # in that case, just ignore the block
+                return
 
-        # if we get here, we might need to create a fork
-        new_chain = self.try_create_fork(block)
-        if new_chain is not None:
-            self.forks.append(new_chain)
-        # otherwise, we're missing the parent block
-        # in that case, just ignore the block
+        # We may need to switch forks, lets run the fork choice rule to check.
+        new_chain = Follower.fork_choice(self.local_chain, self.forks)
+
+        if new_chain == self.local_chain:
+            # we have not re-org'd therefore we can simply update our ledger state
+            # if this block extend our local chain
+            if self.local_chain.tip() == block:
+                self.ledger_state.apply(block)
+        else:
+            # we have re-org'd, therefore we must roll back out ledger state and
+            # re-apply blocks from the new chain
+            ledger_state = self.ledger_state_snapshot.copy()
+            for block in new_chain.blocks:
+                ledger_state.apply(block)
+
+            self.ledger_state = ledger_state
+            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(local_chain: Chain, forks: List[Chain]) -> Chain:
         # TODO: define k and s
         return maxvalid_bg(local_chain, forks, 0, 0)
 
-    def tip(self) -> BlockHeader:
-        return self.fork_choice()
-
-
-@dataclass
-class Coin:
-    pk: int
-    value: int
-
-    def commitment(self) -> Id:
-        # TODO: mocked until CL is understood
-        pk_bytes = int.to_bytes(self.pk, length=32, byteorder="little")
-        value_bytes = int.to_bytes(self.value, length=32, byteorder="little")
-
-        h = sha256()
-        h.update(pk_bytes)
-        h.update(value_bytes)
-        return h.digest()
-
-    def nullifier(self) -> Id:
-        # TODO: mocked until CL is understood
-        pk_bytes = int.to_bytes(self.pk, length=32, byteorder="little")
-        value_bytes = int.to_bytes(self.value, length=32, byteorder="little")
-
-        h = sha256()
-        h.update(pk_bytes)
-        h.update(value_bytes)
-        h.update(b"\x00")  # extra 0 byte to differentiate from commitment
-        return h.digest()
+    def tip_id(self) -> Id:
+        if self.local_chain.length() > 0:
+            return self.local_chain.tip().id
+        else:
+            return self.ledger_state.block
 
 
 @dataclass
@@ -237,8 +272,8 @@ class EpochState:
     # The nonce snapshot is taken 7k/f slots into the previous epoch
     nonce_snapshot: LedgerState
 
-    def is_coin_old_enough_to_lead(self, coin: Coin) -> bool:
-        return coin in self.stake_distribution.commitments
+    def verify_commitment_is_old_enough_to_lead(self, commitment: Id) -> bool:
+        return self.stake_distribution_snapshot.verify_committed(commitment)
 
     def total_stake(self) -> int:
         """Returns the total stake that will be used to reletivize leadership proofs during this epoch"""
@@ -291,9 +326,7 @@ class Leader:
         self, epoch: EpochState, slot: Slot
     ) -> MockLeaderProof | None:
         if self._is_slot_leader(epoch, slot):
-            return MockLeaderProof(
-                commitment=self.coin.commitment(), nullifier=self.coin.nullifier()
-            )
+            return MockLeaderProof.from_coin(self.coin)
 
     def propose_block(self, slot: Slot, parent: BlockHeader) -> BlockHeader:
         return BlockHeader(parent=parent.id(), slot=slot)

cryptarchia/test_fork_choice.py

@@ -11,6 +11,7 @@ from cryptarchia.cryptarchia import (
     Slot,
     Id,
     MockLeaderProof,
+    Coin,
 )
 
 
@@ -22,7 +23,7 @@ def make_block(parent_id: Id, slot: Slot, content: bytes) -> BlockHeader:
         content_size=1,
         slot=slot,
         content_id=content_id,
-        leader_proof=MockLeaderProof(commitment=0, nullifier=0),
+        leader_proof=MockLeaderProof.from_coin(Coin(pk=0, value=10)),
     )
 
 

cryptarchia/test_ledger_state_update.py

@@ -0,0 +1,67 @@
+from unittest import TestCase
+
+import numpy as np
+
+from .cryptarchia import (
+    Follower,
+    TimeConfig,
+    BlockHeader,
+    Config,
+    Coin,
+    LedgerState,
+    MockLeaderProof,
+    Slot,
+)
+
+
+def config() -> Config:
+    return Config(
+        k=10, 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(
+            block=bytes(32),
+            nonce=bytes(32),
+            total_stake=leader_coin.value,
+            commitments={leader_coin.commitment()},
+            nullifiers=set(),
+        )
+
+        follower = Follower(genesis_state, config())
+
+        block = BlockHeader(
+            slot=Slot(0),
+            parent=genesis_state.block,
+            content_size=1,
+            content_id=bytes(32),
+            leader_proof=MockLeaderProof.from_coin(leader_coin),
+        )
+
+        follower.on_block(block)
+
+        # Follower should have accepted the block
+        assert follower.local_chain.length() == 1
+        assert follower.local_chain.tip() == block
+
+        # 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(
+            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