Refactor SlotsBuilder and DataSampler types to use generic parameters SomeTree and SomeHash

2026-01-15 20:03:09 +00:00 · 2025-05-28 19:04:22 -06:00 · 2025-05-28 19:04:22 -06:00 · 6ec73f396d
commit 6ec73f396d
parent 086e6f10e0
2 changed files with 104 additions and 105 deletions
--- a/codex/slots/builder/builder.nim
+++ b/codex/slots/builder/builder.nim
@ -34,107 +34,125 @@ export converters, asynciter
 logScope:
  topics = "codex slotsbuilder"

-type SlotsBuilder*[T, H] = ref object of RootObj
+type SlotsBuilder*[SomeTree, SomeHash] = ref object of RootObj
  store: BlockStore
  manifest: Manifest # current manifest
  strategy: IndexingStrategy # indexing strategy
  cellSize: NBytes # cell size
  numSlotBlocks: Natural
    # number of blocks per slot (should yield a power of two number of cells)
-  slotRoots: seq[H] # roots of the slots
+  slotRoots: seq[SomeHash] # roots of the slots
  emptyBlock: seq[byte] # empty block
-  verifiableTree: ?T # verification tree (dataset tree)
-  emptyDigestTree: T # empty digest tree for empty blocks
+  verifiableTree: ?SomeTree # verification tree (dataset tree)
+  emptyDigestTree: SomeTree # empty digest tree for empty blocks

-func verifiable*[T, H](self: SlotsBuilder[T, H]): bool {.inline.} =
+func verifiable*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): bool {.inline.} =
  ## Returns true if the slots are verifiable.
  ##

  self.manifest.verifiable

-func slotRoots*[T, H](self: SlotsBuilder[T, H]): seq[H] {.inline.} =
+func slotRoots*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): seq[SomeHash] {.inline.} =
  ## Returns the slot roots.
  ##

  self.slotRoots

-func verifyTree*[T, H](self: SlotsBuilder[T, H]): ?T {.inline.} =
+func verifyTree*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): ?SomeTree {.inline.} =
  ## Returns the slots tree (verification tree).
  ##

  self.verifiableTree

-func verifyRoot*[T, H](self: SlotsBuilder[T, H]): ?H {.inline.} =
+func verifyRoot*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): ?SomeHash {.inline.} =
  ## Returns the slots root (verification root).
  ##

  if tree =? self.verifyTree and root =? tree.root:
    return some root

-func numSlots*[T, H](self: SlotsBuilder[T, H]): Natural =
+func numSlots*[SomeTree, SomeHash](self: SlotsBuilder[SomeTree, SomeHash]): Natural =
  ## Number of slots.
  ##

  self.manifest.numSlots

-func numSlotBlocks*[T, H](self: SlotsBuilder[T, H]): Natural =
+func numSlotBlocks*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): Natural =
  ## Number of blocks per slot.
  ##

  self.numSlotBlocks

-func numBlocks*[T, H](self: SlotsBuilder[T, H]): Natural =
+func numBlocks*[SomeTree, SomeHash](self: SlotsBuilder[SomeTree, SomeHash]): Natural =
  ## Number of blocks.
  ##

  self.numSlotBlocks * self.manifest.numSlots

-func slotBytes*[T, H](self: SlotsBuilder[T, H]): NBytes =
+func slotBytes*[SomeTree, SomeHash](self: SlotsBuilder[SomeTree, SomeHash]): NBytes =
  ## Number of bytes per slot.
  ##

  (self.manifest.blockSize.int * self.numSlotBlocks).NBytes

-func numBlockCells*[T, H](self: SlotsBuilder[T, H]): Natural =
+func numBlockCells*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): Natural =
  ## Number of cells per block.
  ##

  (self.manifest.blockSize div self.cellSize).Natural

-func cellSize*[T, H](self: SlotsBuilder[T, H]): NBytes =
+func cellSize*[SomeTree, SomeHash](self: SlotsBuilder[SomeTree, SomeHash]): NBytes =
  ## Cell size.
  ##

  self.cellSize

-func numSlotCells*[T, H](self: SlotsBuilder[T, H]): Natural =
+func numSlotCells*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
+): Natural =
  ## Number of cells per slot.
  ##

  self.numBlockCells * self.numSlotBlocks

-func slotIndiciesIter*[T, H](self: SlotsBuilder[T, H], slot: Natural): ?!Iter[int] =
+func slotIndiciesIter*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], slot: Natural
+): ?!Iter[int] =
  ## Returns the slot indices.
  ##

  self.strategy.getIndicies(slot).catch

-func slotIndicies*[T, H](self: SlotsBuilder[T, H], slot: Natural): seq[int] =
+func slotIndicies*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], slot: Natural
+): seq[int] =
  ## Returns the slot indices.
  ##

  if iter =? self.strategy.getIndicies(slot).catch:
    return toSeq(iter)

-func manifest*[T, H](self: SlotsBuilder[T, H]): Manifest =
+func manifest*[SomeTree, SomeHash](self: SlotsBuilder[SomeTree, SomeHash]): Manifest =
  ## Returns the manifest.
  ##

  self.manifest

-proc buildBlockTree*[T, H](
-    self: SlotsBuilder[T, H], blkIdx: Natural, slotPos: Natural
-): Future[?!(seq[byte], T)] {.async: (raises: [CancelledError]).} =
+proc buildBlockTree*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], blkIdx: Natural, slotPos: Natural
+): Future[?!(seq[byte], SomeTree)] {.async: (raises: [CancelledError]).} =
  ## Build the block digest tree and return a tuple with the
  ## block data and the tree.
  ##
@ -152,22 +170,17 @@ proc buildBlockTree*[T, H](
    trace "Returning empty digest tree for pad block"
    return success (self.emptyBlock, self.emptyDigestTree)

-  without blk =? await self.store.getBlock(self.manifest.treeCid, blkIdx), err:
-    error "Failed to get block CID for tree at index", err = err.msg
-    return failure(err)
+  let blk = ?await self.store.getBlock(self.manifest.treeCid, blkIdx)

  if blk.isEmpty:
    success (self.emptyBlock, self.emptyDigestTree)
  else:
-    without tree =? T.digestTree(blk.data, self.cellSize.int), err:
-      error "Failed to create digest for block", err = err.msg
-      return failure(err)
-
+    let tree = ?SomeTree.digestTree(blk.data, self.cellSize.int)
    success (blk.data, tree)

-proc getCellHashes*[T, H](
-    self: SlotsBuilder[T, H], slotIndex: Natural
-): Future[?!seq[H]] {.async: (raises: [CancelledError, IndexingError]).} =
+proc getCellHashes*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], slotIndex: Natural
+): Future[?!seq[SomeHash]] {.async: (raises: [CancelledError]).} =
  ## Collect all the cells from a block and return
  ## their hashes.
  ##
@ -184,7 +197,7 @@ proc getCellHashes*[T, H](
    slotIndex = slotIndex

  let hashes = collect(newSeq):
-    for i, blkIdx in self.strategy.getIndicies(slotIndex):
+    for i, blkIdx in ?self.strategy.getIndicies(slotIndex).catch:
      logScope:
        blkIdx = blkIdx
        pos = i
@ -200,25 +213,18 @@ proc getCellHashes*[T, H](

  success hashes

-proc buildSlotTree*[T, H](
-    self: SlotsBuilder[T, H], slotIndex: Natural
-): Future[?!T] {.async: (raises: [CancelledError]).} =
+proc buildSlotTree*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], slotIndex: Natural
+): Future[?!SomeTree] {.async: (raises: [CancelledError]).} =
  ## Build the slot tree from the block digest hashes
  ## and return the tree.

-  try:
-    without cellHashes =? (await self.getCellHashes(slotIndex)), err:
-      error "Failed to select slot blocks", err = err.msg
-      return failure(err)
+  let cellHashes = ?await self.getCellHashes(slotIndex)
+  SomeTree.init(cellHashes)

-    T.init(cellHashes)
-  except IndexingError as err:
-    error "Failed to build slot tree", err = err.msg
-    return failure(err)
-
-proc buildSlot*[T, H](
-    self: SlotsBuilder[T, H], slotIndex: Natural
-): Future[?!H] {.async: (raises: [CancelledError]).} =
+proc buildSlot*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], slotIndex: Natural
+): Future[?!SomeHash] {.async: (raises: [CancelledError]).} =
  ## Build a slot tree and store the proofs in
  ## the block store.
  ##
@ -244,18 +250,17 @@ proc buildSlot*[T, H](
      error "Failed to get proof for slot tree", err = err.msg
      return failure(err)

-    if err =?
-        (await self.store.putCidAndProof(treeCid, i, cellCid, encodableProof)).errorOption:
-      error "Failed to store slot tree", err = err.msg
-      return failure(err)
+    ?(await self.store.putCidAndProof(treeCid, i, cellCid, encodableProof))

  tree.root()

-func buildVerifyTree*[T, H](self: SlotsBuilder[T, H], slotRoots: openArray[H]): ?!T =
-  T.init(@slotRoots)
+func buildVerifyTree*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash], slotRoots: openArray[SomeHash]
+): ?!SomeTree =
+  SomeTree.init(@slotRoots)

-proc buildSlots*[T, H](
-    self: SlotsBuilder[T, H]
+proc buildSlots*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
 ): Future[?!void] {.async: (raises: [CancelledError]).} =
  ## Build all slot trees and store them in the block store.
  ##
@ -269,10 +274,7 @@ proc buildSlots*[T, H](
  if self.slotRoots.len == 0:
    self.slotRoots = collect(newSeq):
      for i in 0 ..< self.manifest.numSlots:
-        without slotRoot =? (await self.buildSlot(i)), err:
-          error "Failed to build slot", err = err.msg, index = i
-          return failure(err)
-        slotRoot
+        ?(await self.buildSlot(i))

  without tree =? self.buildVerifyTree(self.slotRoots) and root =? tree.root, err:
    error "Failed to build slot roots tree", err = err.msg
@ -286,17 +288,15 @@ proc buildSlots*[T, H](

  success()

-proc buildManifest*[T, H](
-    self: SlotsBuilder[T, H]
+proc buildManifest*[SomeTree, SomeHash](
+    self: SlotsBuilder[SomeTree, SomeHash]
 ): Future[?!Manifest] {.async: (raises: [CancelledError]).} =
-  if err =? (await self.buildSlots()).errorOption:
-    error "Failed to build slot roots", err = err.msg
-    return failure(err)
+  ## Build the manifest from the slots and return it.
+  ##

-  without rootCids =? self.slotRoots.toSlotCids(), err:
-    error "Failed to map slot roots to CIDs", err = err.msg
-    return failure(err)
+  ?(await self.buildSlots()) # build all slots first

+  let rootCids = ?self.slotRoots.toSlotCids()
  without rootProvingCidRes =? self.verifyRoot .? toVerifyCid() and
    rootProvingCid =? rootProvingCidRes, err:
    error "Failed to map slot roots to CIDs", err = err.msg
@ -306,13 +306,13 @@ proc buildManifest*[T, H](
    self.manifest, rootProvingCid, rootCids, self.cellSize, self.strategy.strategyType
  )

-proc new*[T, H](
-    _: type SlotsBuilder[T, H],
+proc new*[SomeTree, SomeHash](
+    _: type SlotsBuilder[SomeTree, SomeHash],
    store: BlockStore,
    manifest: Manifest,
    strategy = SteppedStrategy,
    cellSize = DefaultCellSize,
-): ?!SlotsBuilder[T, H] =
+): ?!SlotsBuilder[SomeTree, SomeHash] =
  if not manifest.protected:
    trace "Manifest is not protected."
    return failure("Manifest is not protected.")
@ -352,7 +352,7 @@ proc new*[T, H](
    numBlocksTotal = numSlotBlocksTotal * manifest.numSlots # number of blocks per slot

    emptyBlock = newSeq[byte](manifest.blockSize.int)
-    emptyDigestTree = ?T.digestTree(emptyBlock, cellSize.int)
+    emptyDigestTree = ?SomeTree.digestTree(emptyBlock, cellSize.int)

    strategy = ?strategy.init(0, numBlocksTotal - 1, manifest.numSlots).catch

@ -368,7 +368,7 @@ proc new*[T, H](

  trace "Creating slots builder"

-  var self = SlotsBuilder[T, H](
+  var self = SlotsBuilder[SomeTree, SomeHash](
    store: store,
    manifest: manifest,
    strategy: strategy,
--- a/codex/slots/sampler/sampler.nim
+++ b/codex/slots/sampler/sampler.nim
@ -29,14 +29,14 @@ import ./utils
 logScope:
  topics = "codex datasampler"

-type DataSampler*[T, H] = ref object of RootObj
+type DataSampler*[SomeTree, SomeHash] = ref object of RootObj
  index: Natural
  blockStore: BlockStore
-  builder: SlotsBuilder[T, H]
+  builder: SlotsBuilder[SomeTree, SomeHash]

-func getCell*[T, H](
-    self: DataSampler[T, H], blkBytes: seq[byte], blkCellIdx: Natural
-): seq[H] =
+func getCell*[SomeTree, SomeHash](
+    self: DataSampler[SomeTree, SomeHash], blkBytes: seq[byte], blkCellIdx: Natural
+): seq[SomeHash] =
  let
    cellSize = self.builder.cellSize.uint64
    dataStart = cellSize * blkCellIdx.uint64
@ -44,11 +44,14 @@ func getCell*[T, H](

  doAssert (dataEnd - dataStart) == cellSize, "Invalid cell size"

-  blkBytes[dataStart ..< dataEnd].elements(H).toSeq()
+  blkBytes[dataStart ..< dataEnd].elements(SomeHash).toSeq()

-proc getSample*[T, H](
-    self: DataSampler[T, H], cellIdx: int, slotTreeCid: Cid, slotRoot: H
-): Future[?!Sample[H]] {.async: (raises: [CancelledError]).} =
+proc getSample*[SomeTree, SomeHash](
+    self: DataSampler[SomeTree, SomeHash],
+    cellIdx: int,
+    slotTreeCid: Cid,
+    slotRoot: SomeHash,
+): Future[?!Sample[SomeHash]] {.async: (raises: [CancelledError]).} =
  let
    cellsPerBlock = self.builder.numBlockCells
    blkCellIdx = cellIdx.toCellInBlk(cellsPerBlock) # block cell index
@ -77,27 +80,22 @@ proc getSample*[T, H](
    cellProof = blkTree.getProof(blkCellIdx).valueOr:
      return failure("Failed to get proof from block tree")

-  success Sample[H](cellData: cellData, merklePaths: (cellProof.path & slotProof.path))
+  success Sample[SomeHash](
+    cellData: cellData, merklePaths: (cellProof.path & slotProof.path)
+  )

-proc getProofInput*[T, H](
-    self: DataSampler[T, H], entropy: ProofChallenge, nSamples: Natural
-): Future[?!ProofInputs[H]] {.async: (raises: [CancelledError]).} =
+proc getProofInput*[SomeTree, SomeHash](
+    self: DataSampler[SomeTree, SomeHash], entropy: ProofChallenge, nSamples: Natural
+): Future[?!ProofInputs[SomeHash]] {.async: (raises: [CancelledError]).} =
  ## Generate proofs as input to the proving circuit.
  ##

  let
-    entropy = H.fromBytes(array[31, byte].initCopyFrom(entropy[0 .. 30]))
-      # truncate to 31 bytes, otherwise it _might_ be greater than mod
-
-    verifyTree = self.builder.verifyTree.toFailure.valueOr:
-      return failure("Failed to get verify tree")
-
-    slotProof = verifyTree.getProof(self.index).valueOr:
-      return failure("Failed to get slot proof")
-
-    datasetRoot = verifyTree.root().valueOr:
-      return failure("Failed to get dataset root")
-
+    # truncate to 31 bytes, otherwise it _might_ be greater than mod
+    entropy = SomeHash.fromBytes(array[31, byte].initCopyFrom(entropy[0 .. 30]))
+    verifyTree = ?self.builder.verifyTree.toFailure
+    slotProof = ?verifyTree.getProof(self.index)
+    datasetRoot = ?verifyTree.root()
    slotTreeCid = self.builder.manifest.slotRoots[self.index]
    slotRoot = self.builder.slotRoots[self.index]
    cellIdxs = entropy.cellIndices(slotRoot, self.builder.numSlotCells, nSamples)
@ -108,10 +106,9 @@ proc getProofInput*[T, H](
  trace "Collecting input for proof"
  let samples = collect(newSeq):
    for cellIdx in cellIdxs:
-      (await self.getSample(cellIdx, slotTreeCid, slotRoot)).valueOr:
-        return failure("Failed to get sample")
+      ?(await self.getSample(cellIdx, slotTreeCid, slotRoot))

-  success ProofInputs[H](
+  success ProofInputs[SomeHash](
    entropy: entropy,
    datasetRoot: datasetRoot,
    slotProof: slotProof.path,
@ -122,12 +119,12 @@ proc getProofInput*[T, H](
    samples: samples,
  )

-proc new*[T, H](
-    _: type DataSampler[T, H],
+proc new*[SomeTree, SomeHash](
+    _: type DataSampler[SomeTree, SomeHash],
    index: Natural,
    blockStore: BlockStore,
-    builder: SlotsBuilder[T, H],
-): ?!DataSampler[T, H] =
+    builder: SlotsBuilder[SomeTree, SomeHash],
+): ?!DataSampler[SomeTree, SomeHash] =
  if index > builder.slotRoots.high:
    error "Slot index is out of range"
    return failure("Slot index is out of range")
@ -135,4 +132,6 @@ proc new*[T, H](
  if not builder.verifiable:
    return failure("Cannot instantiate DataSampler for non-verifiable builder")

-  success DataSampler[T, H](index: index, blockStore: blockStore, builder: builder)
+  success DataSampler[SomeTree, SomeHash](
+    index: index, blockStore: blockStore, builder: builder
+  )