diff --git a/codex/slotbuilder/slotbuilder.nim b/codex/slotbuilder/slotbuilder.nim index 0aa2c476..0737ae8f 100644 --- a/codex/slotbuilder/slotbuilder.nim +++ b/codex/slotbuilder/slotbuilder.nim @@ -6,6 +6,7 @@ import pkg/questionable/results import ../merkletree import ../stores import ../manifest +import ../utils let # TODO: Unified with the CellSize specified in branch "data-sampler" @@ -38,6 +39,9 @@ proc new*( numberOfSlotBlocks: numberOfSlotBlocks )) +proc cellsPerBlock(self: SlotBuilder): int = + self.manifest.blockSize.int div CellSize + proc getTreeLeafCid(self: SlotBuilder, datasetTreeCid: Cid, datasetBlockIndex: int): Future[?!Cid] {.async.} = without slotBlockCid =? await self.blockStore.getCid(datasetTreeCid, datasetBlockIndex), err: error "Failed to get block for tree at index", index=datasetBlockIndex, tree=datasetTreeCid @@ -85,16 +89,63 @@ proc calculateNumberOfPaddingCells*(self: SlotBuilder, numberOfSlotBlocks: int): raiseAssert("BlockSize should always be divisable by Cell size (2kb).") let - cellsPerBlock = blockSize div CellSize - numberOfCells = numberOfSlotBlocks * cellsPerBlock + numberOfCells = numberOfSlotBlocks * self.cellsPerBlock nextPowerOfTwo = findNextPowerOfTwo(numberOfCells) return nextPowerOfTwo - numberOfCells -proc createAndSaveSlotTree*(self: SlotBuilder, datasetSlotIndex: int): Future[?!MerkleTree] {.async.} = - without var builder =? MerkleTreeBuilder.init(), err: +proc addSlotBlocksToTreeBuilder(builder: var MerkleTreeBuilder, slotBlocks: seq[Cid]): ?!void = + for slotBlockCid in slotBlocks: + without leafHash =? slotBlockCid.mhash: + error "Failed to get leaf hash from CID" + return failure("Failed to get leaf hash from CID") + + if builder.addLeaf(leafHash).isErr: + error "Failed to add slotBlockCid to slot tree builder" + return failure("Failed to add slotBlockCid to slot tree builder") + + return success() + +proc addPadBlocksToTreeBuilder(self: SlotBuilder, builder: var MerkleTreeBuilder, nBlocks: int): ?!void = + without cid =? emptyCid(self.manifest.version, self.manifest.hcodec, self.manifest.codec), err: + error "Unable to initialize empty cid" return failure(err) + without emptyLeaf =? cid.mhash: + error "Failed to get leaf hash from empty CID" + return failure("Failed to get leaf hash from empty CID") + + for i in 0 ..< nBlocks: + if builder.addLeaf(emptyLeaf).isErr: + error "Failed to add empty leaf to slot tree builder" + return failure("Failed to add empty leaf to slot tree builder") + + return success() + +proc buildSlotTree*(self: SlotBuilder, slotBlocks: seq[Cid], numberOfPaddingCells: int): Future[?!MerkleTree] {.async.} = + let numberOfPadBlocks = divUp(numberOfPaddingCells, self.cellsPerBlock) + + without var builder =? MerkleTreeBuilder.init(), err: + error "Failed to initialize merkle tree builder" + return failure(err) + + if addSlotBlocksToTreeBuilder(builder, slotBlocks).isErr: + error "Failed to add slot blocks to tree builder" + return failure("Failed to add slot blocks to tree builder") + + if self.addPadBlocksToTreeBuilder(builder, numberOfPadBlocks).isErr: + error "Failed to add padding blocks to tree builder" + return failure("Failed to add padding blocks to tree builder") + + without slotTree =? builder.build(), err: + error "Failed to build slot tree" + return failure(err) + + return success(slotTree) + +proc createAndSaveSlotTree*(self: SlotBuilder, datasetSlotIndex: int): Future[?!MerkleTree] {.async.} = + + raiseAssert("not implemented") # select slot blocks diff --git a/tests/codex/slotbuilder/testslotbuilder.nim b/tests/codex/slotbuilder/testslotbuilder.nim index 29a52675..f6c75421 100644 --- a/tests/codex/slotbuilder/testslotbuilder.nim +++ b/tests/codex/slotbuilder/testslotbuilder.nim @@ -18,6 +18,7 @@ import codex/slotbuilder/slotbuilder asyncchecksuite "Slot builder": let blockSize = 64 * 1024 + numberOfCellsPerBlock = blockSize div CellSize numberOfSlotBlocks = 6 numberOfSlots = 5 numberOfDatasetBlocks = numberOfSlotBlocks * numberOfSlots @@ -121,50 +122,63 @@ asyncchecksuite "Slot builder": check: expectedPadCells == nPadCells + for i in 0 ..< numberOfSlots: + test "Can select slot block CIDs (index: " & $i & ")": + let + steppedStrategy = SteppedIndexingStrategy.new(0, numberOfDatasetBlocks - 1, numberOfSlots) + expectedDatasetBlockIndicies = steppedStrategy.getIndicies(i) + expectedBlockCids = expectedDatasetBlockIndicies.mapIt(datasetBlocks[it].cid) + slotBlockCids = (await slotBuilder.selectSlotBlocks(i)).tryGet() - # for i in 0 ..< numberOfSlots: - # test "Can select slot block CIDs (index: " & $i & ")": - # let - # steppedStrategy = SteppedIndexingStrategy.new(0, numberOfDatasetBlocks - 1, numberOfSlots) - # expectedDatasetBlockIndicies = steppedStrategy.getIndicies(i) - # expectedBlockCids = expectedDatasetBlockIndicies.mapIt(datasetBlocks[it].cid) + check: + expectedBlockCids == slotBlockCids - # slotBlockCids = (await slotBuilder.selectSlotBlocks(i)).tryGet() + test "Can create slot tree": + let + slotBlockCids = datasetBlocks[0 ..< numberOfSlotBlocks].mapIt(it.cid) + numPadCells = numberOfCellsPerBlock div 2 # We expect 1 pad block. + expectedEmptyCid = emptyCid(protectedManifest.version, protectedManifest.hcodec, protectedManifest.codec) - # check: - # expectedBlockCids == slotBlockCids + slotTree = (await slotBuilder.buildSlotTree(slotBlockCids, numPadCells)).tryGet() + + check: + # Tree size + slotTree.leavesCount == slotBlockCids.len + 1 + + for i in 0 ..< numberOfSlotBlocks: + check: + # Each slot block + slotTree.getLeafCid(i).tryget() == slotBlockCids[i] + + check: + # 1 pad block + slotTree.getLeafCid(numberOfSlotBlocks) == expectedEmptyCid - - - - # for i in 0 ..< numberOfSlots: - # test "Can create slot tree given index (" & $i & ")": - # let - # selectStart = i * numberOfSlotBlocks - # selectEnd = selectStart + numberOfSlotBlocks - # expectedCids = datasetBlocks.mapIt(it.cid)[selectStart ..< selectEnd] - # m = (await slotBuilder.createSlotIntermediateManifest(i)).tryGet() - - # check: - # m.treeCid # check - # m.datasetSize == (numberOfSlotBlocks * blockSize).NBytes - # m.blockSize == blockSize - # m.version == manifest.version - # m.hcodec == manifest.hcodec - # m.codec == manifest.codec - # #m.ecK == ?? - # #m.ecM == ?? - # m.originalTreeCid == manifest.originalTreeCid - # m.originalDatasetSize = manifest.originalDatasetSize - - # m.isSlot == true - # m.datasetSlotIndex == i - # m.originalProtectedTreeCide == manifest.treeCid - # m.originalProtectedDatasetSize == manifest.datasetSize + # test "Can create slot tree (index: " & $i & ")": + # let + # slotBlockCids = + # m = (await slotBuilder.buildSlotTree(slotBlockCids, numPadCells)).tryGet() + + # check: + # m.treeCid # check + # m.datasetSize == (numberOfSlotBlocks * blockSize).NBytes + # m.blockSize == blockSize + # m.version == manifest.version + # m.hcodec == manifest.hcodec + # m.codec == manifest.codec + # #m.ecK == ?? + # #m.ecM == ?? + # m.originalTreeCid == manifest.originalTreeCid + # m.originalDatasetSize = manifest.originalDatasetSize + + # m.isSlot == true + # m.datasetSlotIndex == i + # m.originalProtectedTreeCide == manifest.treeCid + # m.originalProtectedDatasetSize == manifest.datasetSize