From 3e706b3eea262c27cc23c2d13da32235f2dd7efa Mon Sep 17 00:00:00 2001
From: Ben <thatbenbierens@gmail.com>
Date: Tue, 17 Dec 2024 12:52:03 +0100
Subject: [PATCH] Restore cidOrTreeCid

---
 codex/blockexchange/engine/discovery.nim      | 19 ++++++-------------
 codex/blockexchange/peers/peerctxstore.nim    |  7 +------
 codex/blocktype.nim                           |  6 ++++++
 .../codex/blockexchange/testpendingblocks.nim |  4 ++--
 4 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/codex/blockexchange/engine/discovery.nim b/codex/blockexchange/engine/discovery.nim
index d97d5707..0b159452 100644
--- a/codex/blockexchange/engine/discovery.nim
+++ b/codex/blockexchange/engine/discovery.nim
@@ -56,19 +56,10 @@ type
     discoveryLoopSleep: Duration                                 # Discovery loop sleep
     inFlightDiscReqs*: Table[Cid, Future[seq[SignedPeerRecord]]] # Inflight discovery requests
 
-proc getCid(address: BlockAddress): Cid =
-  # We advertise and discover only the CID part of a block address.
-  # Indices are ignored. This means that multiple blocks of the same tree will
-  # have a single DHT entry.
-  if address.leaf:
-    address.treeCid
-  else:
-    address.cid
-
 proc discoveryQueueLoop(b: DiscoveryEngine) {.async: (raises: []).} =
   while b.discEngineRunning:
     for address in toSeq(b.pendingBlocks.wantList):
-      let cid = address.getCid()
+      let cid = address.cidOrTreeCid
       try:
         await b.discoveryQueue.put(cid)
       except CancelledError:
@@ -134,15 +125,17 @@ proc discoveryTaskLoop(b: DiscoveryEngine) {.async: (raises: []).} =
 
   info "Exiting discovery task runner"
 
-proc queueFindBlocksReq*(b: DiscoveryEngine, addresses: seq[BlockAddress]) {.inline.} =
-  for address in addresses:
-    let cid = address.getCid()
+proc queueFindBlocksReq*(b: DiscoveryEngine, cids: seq[Cid]) {.inline.} =
+  for cid in cids:
     if cid notin b.discoveryQueue:
       try:
         b.discoveryQueue.putNoWait(cid)
       except CatchableError as exc:
         warn "Exception queueing discovery request", exc = exc.msg
 
+proc queueFindBlocksReq*(b: DiscoveryEngine, addresses: seq[BlockAddress]) {.inline.} =
+  b.queueFindBlocksReq(addresses.mapIt(it.cidOrTreeCid))
+
 proc start*(b: DiscoveryEngine) {.async.} =
   ## Start the discengine task
   ##
diff --git a/codex/blockexchange/peers/peerctxstore.nim b/codex/blockexchange/peers/peerctxstore.nim
index b145dd31..2d838094 100644
--- a/codex/blockexchange/peers/peerctxstore.nim
+++ b/codex/blockexchange/peers/peerctxstore.nim
@@ -65,12 +65,7 @@ func peersHave*(self: PeerCtxStore, address: BlockAddress): seq[BlockExcPeerCtx]
   toSeq(self.peers.values).filterIt( it.peerHave.anyIt( it == address ) )
 
 func countPeersWhoHave*(self: PeerCtxStore, cid: Cid): int =
-  proc getCid(address: BlockAddress): Cid =
-    if address.leaf:
-      address.treeCid
-    else:
-      address.cid
-  self.peers.values.countIt(it.peerHave.anyIt( it.getCid() == cid ) )
+  self.peers.values.countIt(it.peerHave.anyIt( it.cidOrTreeCid == cid ) )
 
 func peersWant*(self: PeerCtxStore, address: BlockAddress): seq[BlockExcPeerCtx] =
   toSeq(self.peers.values).filterIt( it.peerWants.anyIt( it == address ) )
diff --git a/codex/blocktype.nim b/codex/blocktype.nim
index a149068d..c44e4fd8 100644
--- a/codex/blocktype.nim
+++ b/codex/blocktype.nim
@@ -66,6 +66,12 @@ proc `$`*(a: BlockAddress): string =
   else:
     "cid: " & $a.cid
 
+proc cidOrTreeCid*(a: BlockAddress): Cid =
+  if a.leaf:
+    a.treeCid
+  else:
+    a.cid
+
 proc address*(b: Block): BlockAddress =
   BlockAddress(leaf: false, cid: b.cid)
 
diff --git a/tests/codex/blockexchange/testpendingblocks.nim b/tests/codex/blockexchange/testpendingblocks.nim
index dd94c4da..d3337cd3 100644
--- a/tests/codex/blockexchange/testpendingblocks.nim
+++ b/tests/codex/blockexchange/testpendingblocks.nim
@@ -63,8 +63,8 @@ checksuite "Pending Blocks":
     discard blks.mapIt( pendingBlocks.getWantHandle( it.cid ) )
 
     check:
-      blks.mapIt( $it.cid ).sorted(cmp[string]) ==
-      toSeq(pendingBlocks.wantListBlockCids).mapIt( $it ).sorted(cmp[string])
+      blks.mapIt( $BlockAddress.init(it.cid) ).sorted(cmp[string]) ==
+      toSeq(pendingBlocks.wantList).mapIt( $it ).sorted(cmp[string])
 
   test "Should get want handles list":
     let