cryptachia/ghost: rewrite on_block to remove dependency on Chain

cryptarchia/cryptarchia.py

@@ -398,7 +398,7 @@ class Follower:
     def __init__(self, genesis_state: LedgerState, config: Config):
         self.config = config
         self.forks = []
-        self.local_chain = Chain([], genesis=genesis_state.block.id())
+        self.local_chain = genesis_state.block.id()
         self.genesis_state = genesis_state
         self.ledger_state = {genesis_state.block.id(): genesis_state.copy()}
         self.epoch_state = {}
@@ -482,66 +482,31 @@ def verify_slot_leader(
 
         return True
 
-    # Try appending this block to an existing chain and return whether
-    # the operation was successful
-    def try_extend_chains(self, block: BlockHeader) -> Optional[Chain]:
-        if self.tip_id() == block.parent:
-            return self.local_chain
-
-        for chain in self.forks:
-            if chain.tip_id() == block.parent:
-                return chain
-
-        return None
-
-    def try_create_fork(self, block: BlockHeader) -> Optional[Chain]:
-        if self.genesis_state.block.id() == block.parent:
-            # this block is forking off the genesis state
-            return Chain(blocks=[], genesis=self.genesis_state.block.id())
-
-        chains = self.forks + [self.local_chain]
-        for chain in chains:
-            block_position = chain.block_position(block.parent)
-            if block_position is not None:
-                return Chain(
-                    blocks=chain.blocks[: block_position + 1],
-                    genesis=self.genesis_state.block.id(),
-                )
-
-        return None
-
     def on_block(self, block: BlockHeader):
         if not self.validate_header(block):
             logger.warning("invalid header")
             return
 
-        # check if the new block extends an existing chain
-        new_chain = self.try_extend_chains(block)
-        if new_chain is None:
-            # 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:
-                logger.warning("missing parent block")
-                # otherwise, we're missing the parent block
-                # in that case, just ignore the block
-                return
-
         new_state = self.ledger_state[block.parent].copy()
         new_state.apply(block)
         self.ledger_state[block.id()] = new_state
 
-        new_chain.blocks.append(block)
+        if block.parent == self.local_chain:
+            # simply extending the local chain
+            self.local_chain = block.id()
+        else:
+            # otherwise, this block creates a fork
+            self.forks.append(block.id())
+
+            # remove any existing fork that is superceded by this block
+            if block.parent in self.forks:
+                self.forks.remove(block.parent)
 
-        # We may need to switch forks, lets run the fork choice rule to check.
-        new_chain_head = self.fork_choice()
-        if new_chain_head != self.local_chain.tip_id():
-            assert new_chain_head == new_chain.tip_id()
-            self.forks.remove(new_chain)
+            # We may need to switch forks, lets run the fork choice rule to check.
+            new_tip = self.fork_choice()
             self.forks.append(self.local_chain)
-            self.local_chain = new_chain
+            self.forks.remove(new_tip)
+            self.local_chain = new_tip
 
     def unimported_orphans(self) -> list[BlockHeader]:
         """
@@ -554,9 +519,9 @@ def unimported_orphans(self) -> list[BlockHeader]:
 
         for fork in self.forks:
             _, fork_depth = common_prefix_depth(
-                tip_state.block.id(), fork.tip_id(), self.ledger_state
+                tip_state.block.id(), fork, self.ledger_state
             )
-            fork_chain = chain_suffix(fork.tip_id(), fork_depth, self.ledger_state)
+            fork_chain = chain_suffix(fork, fork_depth, self.ledger_state)
             for block_state in fork_chain:
                 b = block_state.block
                 if b.leader_proof.nullifier not in tip_state.nullifiers:
@@ -566,20 +531,20 @@ def unimported_orphans(self) -> list[BlockHeader]:
         return orphans
 
     # Evaluate the fork choice rule and return the chain we should be following
-    def fork_choice(self) -> Chain:
+    def fork_choice(self) -> Id:
         return maxvalid_bg(
-            self.local_chain.tip_id(),
-            [f.tip_id() for f in self.forks],
+            self.local_chain,
+            self.forks,
             self.ledger_state,
             k=self.config.k,
             s=self.config.s,
         )
 
     def tip(self) -> BlockHeader:
-        return self.local_chain.tip()
+        return self.tip_state().block
 
     def tip_id(self) -> Id:
-        return self.local_chain.tip_id()
+        return self.local_chain
 
     def tip_state(self) -> LedgerState:
         return self.ledger_state[self.tip_id()]

cryptarchia/test_fork_choice.py

@@ -156,7 +156,7 @@ def test_fork_choice_integration(self):
         follower.on_block(b1)
 
         assert follower.tip_id() == b1.id()
-        assert follower.forks == []
+        assert follower.forks == [], follower.forks
 
         # -- then we fork --
         #
@@ -174,7 +174,7 @@ def test_fork_choice_integration(self):
         follower.on_block(b3)
 
         assert follower.tip_id() == b2.id()
-        assert len(follower.forks) == 1 and follower.forks[0].tip_id() == b3.id()
+        assert len(follower.forks) == 1 and follower.forks[0] == b3.id()
 
         # -- extend the fork causing a re-org --
         #
@@ -189,4 +189,4 @@ def test_fork_choice_integration(self):
         follower.on_block(b4)
 
         assert follower.tip_id() == b4.id()
-        assert len(follower.forks) == 1 and follower.forks[0].tip_id() == b2.id()
+        assert len(follower.forks) == 1 and follower.forks[0] == b2.id(), follower.forks

cryptarchia/test_ledger_state_update.py

@@ -89,7 +89,7 @@ def test_fork_creation(self):
         follower.on_block(block_2)
         assert follower.tip() == block_1
         assert len(follower.forks) == 1, f"{len(follower.forks)}"
-        assert follower.forks[0].tip() == block_2
+        assert follower.forks[0] == block_2.id()
 
         # coin_2 wins slot 1 and chooses to extend from block_1
         # coin_3 also wins slot 1 and but chooses to extend from block_2
@@ -100,16 +100,16 @@ def test_fork_creation(self):
         follower.on_block(block_4)
         assert follower.tip() == block_3
         assert len(follower.forks) == 1, f"{len(follower.forks)}"
-        assert follower.forks[0].tip() == block_4
+        assert follower.forks[0] == block_4.id()
 
         # coin_4 wins slot 1 and but chooses to extend from block_2 as well
         # The block is accepted. A new fork is created "from the block_2".
         block_5 = mk_block(parent=block_2, slot=1, coin=coins[4])
         follower.on_block(block_5)
         assert follower.tip() == block_3
         assert len(follower.forks) == 2, f"{len(follower.forks)}"
-        assert follower.forks[0].tip() == block_4
-        assert follower.forks[1].tip() == block_5
+        assert follower.forks[0] == block_4.id()
+        assert follower.forks[1] == block_5.id()
 
         # A block based on an unknown parent is not accepted.
         # Nothing changes from the local chain and forks.
@@ -118,8 +118,8 @@ def test_fork_creation(self):
         follower.on_block(block_6)
         assert follower.tip() == block_3
         assert len(follower.forks) == 2, f"{len(follower.forks)}"
-        assert follower.forks[0].tip() == block_4
-        assert follower.forks[1].tip() == block_5
+        assert follower.forks[0] == block_4.id()
+        assert follower.forks[1] == block_5.id()
 
     def test_epoch_transition(self):
         leader_coins = [Coin(sk=i, value=100) for i in range(4)]

cryptarchia/test_orphaned_proofs.py

@@ -43,7 +43,7 @@ def test_simple_orphan_import(self):
             follower.on_block(b)
 
         assert follower.tip() == b2
-        assert [f.tip() for f in follower.forks] == [b3]
+        assert [f for f in follower.forks] == [b3.id()]
         assert follower.unimported_orphans() == [b3]
 
         # -- extend with import --
@@ -58,7 +58,7 @@ def test_simple_orphan_import(self):
         follower.on_block(b4)
 
         assert follower.tip() == b4
-        assert [f.tip() for f in follower.forks] == [b3]
+        assert [f for f in follower.forks] == [b3.id()]
         assert follower.unimported_orphans() == []
 
     def test_orphan_proof_import_from_long_running_fork(self):
@@ -89,7 +89,7 @@ def test_orphan_proof_import_from_long_running_fork(self):
             follower.on_block(b)
 
         assert follower.tip() == b3
-        assert [f.tip() for f in follower.forks] == [b5]
+        assert [f for f in follower.forks] == [b5.id()]
         assert follower.unimported_orphans() == [b4, b5]
 
         # -- extend b3, importing the fork --
@@ -104,7 +104,7 @@ def test_orphan_proof_import_from_long_running_fork(self):
         follower.on_block(b6)
 
         assert follower.tip() == b6
-        assert [f.tip() for f in follower.forks] == [b5]
+        assert [f for f in follower.forks] == [b5.id()]
 
     def test_orphan_proof_import_from_fork_without_direct_shared_parent(self):
         coins = [Coin(sk=i, value=10) for i in range(2)]
@@ -135,7 +135,7 @@ def test_orphan_proof_import_from_fork_without_direct_shared_parent(self):
             follower.on_block(b)
 
         assert follower.tip() == b4
-        assert [f.tip() for f in follower.forks] == [b7]
+        assert [f for f in follower.forks] == [b7.id()]
         assert follower.unimported_orphans() == [b5, b6, b7]
 
         # -- extend b4, importing the forks --
@@ -153,7 +153,7 @@ def test_orphan_proof_import_from_fork_without_direct_shared_parent(self):
         follower.on_block(b8)
 
         assert follower.tip() == b8
-        assert [f.tip() for f in follower.forks] == [b7]
+        assert [f for f in follower.forks] == [b7.id()]
         assert follower.unimported_orphans() == []
 
     def test_unimported_orphans(self):
@@ -199,7 +199,7 @@ def test_unimported_orphans(self):
             follower.on_block(b)
 
         assert follower.tip() == b3
-        assert [f.tip() for f in follower.forks] == [b5, b6]
+        assert [f for f in follower.forks] == [b5.id(), b6.id()]
         assert follower.unimported_orphans() == [b4, b5, b6]
 
         b7, c_a = mk_block(b3, 4, c_a, orphaned_proofs=[b4, b5, b6]), c_a.evolve()