From cc50e2e8970b18c72100c5dd7fa27577af1cc34b Mon Sep 17 00:00:00 2001
From: David Rusu <davidrusu.me@gmail.com>
Date: Fri, 1 Nov 2024 10:17:37 +0400
Subject: [PATCH] cryptarchia/ghost: rewrite unimported_orphan w/
 common_prefix_depths

---
 cryptarchia/cryptarchia.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/cryptarchia/cryptarchia.py b/cryptarchia/cryptarchia.py
index 15dfe41..4caea8f 100644
--- a/cryptarchia/cryptarchia.py
+++ b/cryptarchia/cryptarchia.py
@@ -551,12 +551,17 @@ class Follower:
         tip_state = self.tip_state().copy()
 
         orphans = []
-        for fork in [self.local_chain, *self.forks]:
-            for block in fork.blocks:
-                for b in [*block.orphaned_proofs, block]:
-                    if b.leader_proof.nullifier not in tip_state.nullifiers:
-                        tip_state.nullifiers.add(b.leader_proof.nullifier)
-                        orphans += [b]
+
+        for fork in self.forks:
+            _, fork_depth = common_prefix_depth(
+                tip_state.block.id(), fork.tip_id(), self.ledger_state
+            )
+            fork_chain = chain_suffix(fork.tip_id(), 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:
+                    tip_state.nullifiers.add(b.leader_proof.nullifier)
+                    orphans += [b]
 
         return orphans
 
@@ -716,6 +721,15 @@ class Leader:
         )
 
 
+def chain_suffix(tip: Id, n: int, states: Dict[Id, LedgerState]) -> list[LedgerState]:
+    chain = []
+    for _ in range(n):
+        chain.append(states[tip])
+        tip = states[tip].block.parent
+
+    return reversed(chain)
+
+
 def common_prefix_depth(
     local_chain: Id, fork: Id, states: Dict[Id, LedgerState]
 ) -> (int, int):