From 84993e89840015dd0d2fedac15482c32630ab414 Mon Sep 17 00:00:00 2001 From: benbierens Date: Mon, 27 Nov 2023 12:01:06 +0100 Subject: [PATCH] Refactoring to object-oriented --- codex/proof/datasampler.nim | 137 ++++++++------ codex/proof/indexing.nim | 24 --- codex/proof/slotblocks.nim | 59 ++++-- tests/codex/proof/testdatasampler.nim | 262 ++++++++++++++------------ tests/codex/proof/testindexing.nim | 29 --- tests/codex/proof/testslotblocks.nim | 39 +++- 6 files changed, 299 insertions(+), 251 deletions(-) delete mode 100644 codex/proof/indexing.nim diff --git a/codex/proof/datasampler.nim b/codex/proof/datasampler.nim index d30aa7c6..469b1c68 100644 --- a/codex/proof/datasampler.nim +++ b/codex/proof/datasampler.nim @@ -17,12 +17,63 @@ import pkg/poseidon2 import misc import slotblocks -import indexing import types +# Index naming convention: +# "Index" => The index of an ElementType within a ContainerType. +# Some examples: +# SlotBlockIndex => The index of a Block within a Slot. +# DatasetBlockIndex => The index of a Block within a Dataset. + logScope: topics = "codex datasampler" +type + DataSampler* = ref object of RootObj + slot: Slot + blockStore: BlockStore + slotBlocks: SlotBlocks + # The following data is invariant over time for a given slot: + slotRootHash: DSFieldElement + slotPoseidonTree: MerkleTree + datasetToSlotProof: MerkleProof + blockSize: uint64 + numberOfCellsInSlot: uint64 + numberOfCellsPerBlock: uint64 + +proc getNumberOfCellsInSlot*(slot: Slot): uint64 = + (slot.request.ask.slotSize.truncate(uint64) div CellSize) + +proc new*( + T: type DataSampler, + slot: Slot, + blockStore: BlockStore, + slotRootHash: DSFieldElement, + slotPoseidonTree: MerkleTree, + datasetToSlotProof: MerkleProof +): Future[?!DataSampler] {.async.} = + # Create a DataSampler for a slot. + # A DataSampler can create the input required for the proving circuit. + without slotBlocks =? await SlotBlocks.new(slot, blockStore), err: + error "Failed to create SlotBlocks object for slot" + return failure(err) + + let + numberOfCellsInSlot = getNumberOfCellsInSlot(slot) + blockSize = slotBlocks.manifest.blockSize.uint64 + + success(DataSampler( + slot: slot, + blockStore: blockStore, + slotBlocks: slotBlocks, + slotRootHash: slotRootHash, + slotPoseidonTree: slotPoseidonTree, + datasetToSlotProof: datasetToSlotProof, + blockSize: blockSize, + numberOfCellsInSlot: numberOfCellsInSlot, + numberOfCellsPerBlock: blockSize div CellSize + )) + func extractLowBits*[n: static int](A: BigInt[n], k: int): uint64 = assert(k > 0 and k <= 64) var r: uint64 = 0 @@ -36,58 +87,51 @@ func extractLowBits*[n: static int](A: BigInt[n], k: int): uint64 = r = bitor(r, 1'u64 shl i) return r -proc convertToSlotCellIndex(fe: DSFieldElement, numberOfCells: int): uint64 = - let log2 = ceilingLog2(numberOfCells) - assert((1 shl log2) == numberOfCells , "expected `numberOfCells` to be a power of two.") +proc convertToSlotCellIndex(self: DataSampler, fe: DSFieldElement): uint64 = + let + n = self.numberOfCellsInSlot.int + log2 = ceilingLog2(n) + assert((1 shl log2) == n , "expected `numberOfCellsInSlot` to be a power of two.") return extractLowBits(fe.toBig(), log2) -proc getNumberOfCellsInSlot*(slot: Slot): uint64 = - (slot.request.ask.slotSize.truncate(uint64) div CellSize) +proc getSlotBlockIndexForSlotCellIndex*(self: DataSampler, slotCellIndex: DSSlotCellIndex): uint64 = + return slotCellIndex div self.numberOfCellsPerBlock -proc findSlotCellIndex*( - slotRootHash: DSFieldElement, - challenge: DSFieldElement, - counter: DSFieldElement, - numberOfCells: uint64): DSSlotCellIndex = +proc getBlockCellIndexForSlotCellIndex*(self: DataSampler, slotCellIndex: DSSlotCellIndex): uint64 = + return slotCellIndex mod self.numberOfCellsPerBlock + +proc findSlotCellIndex*(self: DataSampler, challenge: DSFieldElement, counter: DSFieldElement): DSSlotCellIndex = # Computes the slot-cell index for a single sample. let - input = @[slotRootHash, challenge, counter] + input = @[self.slotRootHash, challenge, counter] hash = Sponge.digest(input, rate = 2) - index = convertToSlotCellIndex(hash, numberOfCells.int) + return convertToSlotCellIndex(self, hash) - return index - -func findSlotCellIndices*( - slot: Slot, - slotRootHash: DSFieldElement, - challenge: DSFieldElement, - nSamples: int): seq[DSSlotCellIndex] = +func findSlotCellIndices*(self: DataSampler, challenge: DSFieldElement, nSamples: int): seq[DSSlotCellIndex] = # Computes nSamples slot-cell indices. - let numberOfCells = getNumberOfCellsInSlot(slot) - return collect(newSeq, (for i in 1..nSamples: findSlotCellIndex(slotRootHash, challenge, toF(i), numberOfCells))) + return collect(newSeq, (for i in 1..nSamples: self.findSlotCellIndex(challenge, toF(i)))) -proc getCellFromBlock*(blk: bt.Block, slotCellIndex: DSSlotCellIndex, blockSize: uint64): DSCell = +proc getCellFromBlock*(self: DataSampler, blk: bt.Block, slotCellIndex: DSSlotCellIndex): DSCell = let - blockCellIndex = getBlockCellIndexForSlotCellIndex(slotCellIndex, blockSize) + blockCellIndex = self.getBlockCellIndexForSlotCellIndex(slotCellIndex) dataStart = (CellSize * blockCellIndex) dataEnd = dataStart + CellSize return blk.data[dataStart ..< dataEnd] -proc getBlockCells*(blk: bt.Block, blockSize: uint64): seq[DSCell] = - let numberOfCellsPerBlock = blockSize div CellSize +proc getBlockCells*(self: DataSampler, blk: bt.Block): seq[DSCell] = var cells: seq[DSCell] - for i in 0..Index" => The index of an ElementType within a ContainerType. -# Some examples: -# SlotBlockIndex => The index of a Block within a Slot. -# DatasetBlockIndex => The index of a Block within a Dataset. - -proc getSlotBlockIndexForSlotCellIndex*(slotCellIndex: DSSlotCellIndex, blockSize: uint64): uint64 = - let numberOfCellsPerBlock = blockSize div CellSize - return slotCellIndex div numberOfCellsPerBlock - -proc getBlockCellIndexForSlotCellIndex*(slotCellIndex: DSSlotCellIndex, blockSize: uint64): uint64 = - let numberOfCellsPerBlock = blockSize div CellSize - return slotCellIndex mod numberOfCellsPerBlock - -proc getDatasetBlockIndexForSlotBlockIndex*(slot: Slot, blockSize: uint64, slotBlockIndex: uint64): uint64 = - let - slotSize = slot.request.ask.slotSize.truncate(uint64) - blocksInSlot = slotSize div blockSize - datasetSlotIndex = slot.slotIndex.truncate(uint64) - return (datasetSlotIndex * blocksInSlot) + slotBlockIndex diff --git a/codex/proof/slotblocks.nim b/codex/proof/slotblocks.nim index a9ea8341..8a71bd50 100644 --- a/codex/proof/slotblocks.nim +++ b/codex/proof/slotblocks.nim @@ -10,14 +10,19 @@ import pkg/questionable/results import ../contracts/requests import ../stores/blockstore import ../manifest -import indexing -proc getManifestForSlot*(slot: Slot, blockstore: BlockStore): Future[?!Manifest] {.async.} = +type + SlotBlocks* = ref object of RootObj + slot: Slot + blockStore: BlockStore + manifest: Manifest + +proc getManifestForSlot(slot: Slot, blockStore: BlockStore): Future[?!Manifest] {.async.} = without manifestBlockCid =? Cid.init(slot.request.content.cid).mapFailure, err: error "Unable to init CID from slot.content.cid" return failure err - without manifestBlock =? await blockstore.getBlock(manifestBlockCid), err: + without manifestBlock =? await blockStore.getBlock(manifestBlockCid), err: error "Failed to fetch manifest block", cid = manifestBlockCid return failure err @@ -27,19 +32,39 @@ proc getManifestForSlot*(slot: Slot, blockstore: BlockStore): Future[?!Manifest] return success(manifest) -proc getSlotBlock*(slot: Slot, blockstore: BlockStore, manifest: Manifest, slotBlockIndex: uint64): Future[?!Block] {.async.} = - let - blocksInManifest = (manifest.datasetSize div manifest.blockSize).uint64 - datasetIndex = getDatasetBlockIndexForSlotBlockIndex(slot, manifest.blockSize.uint64, slotBlockIndex) - - if datasetIndex >= blocksInManifest: - return failure("Found slotBlockIndex that is out-of-range: " & $datasetIndex) - - return await blockstore.getBlock(manifest.treeCid, datasetIndex) - -proc getSlotBlock*(slot: Slot, blockstore: BlockStore, slotBlockIndex: uint64): Future[?!Block] {.async.} = - without manifest =? (await getManifestForSlot(slot, blockstore)), err: +proc new*( + T: type SlotBlocks, + slot: Slot, + blockStore: BlockStore +): Future[?!SlotBlocks] {.async.} = + # Create a SlotBlocks object for a slot. + # SlotBlocks lets you get the manifest of a slot and blocks by slotBlockIndex for a slot. + without manifest =? await getManifestForSlot(slot, blockStore): error "Failed to get manifest for slot" - return failure(err) + return failure("Failed to get manifest for slot") - return await getSlotBlock(slot, blockstore, manifest, slotBlockIndex) + success(SlotBlocks( + slot: slot, + blockStore: blockStore, + manifest: manifest + )) + +proc manifest*(self: SlotBlocks): Manifest = + self.manifest + +proc getDatasetBlockIndexForSlotBlockIndex*(self: SlotBlocks, slotBlockIndex: uint64): uint64 = + let + slotSize = self.slot.request.ask.slotSize.truncate(uint64) + blocksInSlot = slotSize div self.manifest.blockSize.uint64 + datasetSlotIndex = self.slot.slotIndex.truncate(uint64) + return (datasetSlotIndex * blocksInSlot) + slotBlockIndex + +proc getSlotBlock*(self: SlotBlocks, slotBlockIndex: uint64): Future[?!Block] {.async.} = + let + blocksInManifest = (self.manifest.datasetSize div self.manifest.blockSize).uint64 + datasetBlockIndex = self.getDatasetBlockIndexForSlotBlockIndex(slotBlockIndex) + + if datasetBlockIndex >= blocksInManifest: + return failure("Found datasetBlockIndex that is out-of-range: " & $datasetBlockIndex) + + return await self.blockStore.getBlock(self.manifest.treeCid, datasetBlockIndex) diff --git a/tests/codex/proof/testdatasampler.nim b/tests/codex/proof/testdatasampler.nim index bf80af19..c7a34d43 100644 --- a/tests/codex/proof/testdatasampler.nim +++ b/tests/codex/proof/testdatasampler.nim @@ -5,6 +5,7 @@ import std/random import pkg/questionable/results import pkg/constantine/math/arithmetic import pkg/poseidon2/types +import pkg/poseidon2/io import pkg/poseidon2 import pkg/chronos import pkg/asynctest @@ -20,7 +21,6 @@ import pkg/codex/stores/cachestore import pkg/codex/proof/datasampler import pkg/codex/proof/misc import pkg/codex/proof/types -import pkg/codex/proof/indexing import ../helpers import ../examples @@ -59,7 +59,7 @@ asyncchecksuite "Test proof datasampler - components": isPow2(numberOfCells) test "Extract low bits": - proc extract(value: int, nBits: int): uint64 = + proc extract(value: uint64, nBits: int): uint64 = let big = toF(value).toBig() return extractLowBits(big, nBits) @@ -70,10 +70,10 @@ asyncchecksuite "Test proof datasampler - components": extract(0x9A, 7) == 0x1A.uint64 extract(0x1248, 10) == 0x248.uint64 extract(0x1248, 12) == 0x248.uint64 - # extract(0x1248306A560C9AC0, 10) == 0x2C0.uint64 - # extract(0x1248306A560C9AC0, 12) == 0xAC0.uint64 - # extract(0x1248306A560C9AC0, 50) == 0x306A560C9AC0.uint64 - # extract(0x1248306A560C9AC0, 52) == 0x8306A560C9AC0.uint64 + extract(0x1248306A560C9AC0.uint64, 10) == 0x2C0.uint64 + extract(0x1248306A560C9AC0.uint64, 12) == 0xAC0.uint64 + extract(0x1248306A560C9AC0.uint64, 50) == 0x306A560C9AC0.uint64 + extract(0x1248306A560C9AC0.uint64, 52) == 0x8306A560C9AC0.uint64 test "Should calculate total number of cells in Slot": let @@ -84,89 +84,6 @@ asyncchecksuite "Test proof datasampler - components": expectedNumberOfCells == 512 expectedNumberOfCells == getNumberOfCellsInSlot(slot) - let knownIndices = @[178.uint64, 277.uint64, 366.uint64] - - test "Can find single slot-cell index": - let numberOfCells = getNumberOfCellsInSlot(slot) - - proc slotCellIndex(i: int): DSSlotCellIndex = - let counter: DSFieldElement = toF(i) - return findSlotCellIndex(slotRootHash, challenge, counter, numberOfCells) - - proc getExpectedIndex(i: int): DSSlotCellIndex = - let hash = Sponge.digest(@[slotRootHash, challenge, toF(i)], rate = 2) - return extractLowBits(hash.toBig(), ceilingLog2(numberOfCells.int)) - - check: - slotCellIndex(1) == getExpectedIndex(1) - slotCellIndex(1) == knownIndices[0] - slotCellIndex(2) == getExpectedIndex(2) - slotCellIndex(2) == knownIndices[1] - slotCellIndex(3) == getExpectedIndex(3) - slotCellIndex(3) == knownIndices[2] - - test "Can find sequence of slot-cell indices": - proc slotCellIndices(n: int): seq[DSSlotCellIndex] = - findSlotCellIndices(slot, slotRootHash, challenge, n) - - let numberOfCells = getNumberOfCellsInSlot(slot) - proc getExpectedIndices(n: int): seq[DSSlotCellIndex] = - return collect(newSeq, (for i in 1..n: findSlotCellIndex(slotRootHash, challenge, toF(i), numberOfCells))) - - check: - slotCellIndices(3) == getExpectedIndices(3) - slotCellIndices(3) == knownIndices - - test "Can get cell from block": - let - blockSize = CellSize * 3 - bytes = newSeqWith(blockSize.int, rand(uint8)) - blk = bt.Block.new(bytes).tryGet() - - sample0 = getCellFromBlock(blk, 0, blockSize.uint64) - sample1 = getCellFromBlock(blk, 1, blockSize.uint64) - sample2 = getCellFromBlock(blk, 2, blockSize.uint64) - - check: - sample0 == bytes[0.. " & $expected & ")": + let + slotCellIndex = input.uint64 + slotBlockIndex = dataSampler.getSlotBlockIndexForSlotCellIndex(slotCellIndex) - discard await localStore.putBlock(manifestBlock) + check: + slotBlockIndex == expected.uint64 - let a = (await getProofInput( - slot, - localStore, - slotRootHash, - slotPoseidonTree, - datasetToSlotProof, - challenge, - nSamples)).tryget() + for (input, expected) in [(10, 10), (31, 31), (32, 0), (63, 31), (64, 0)]: + test "Can get blockCellIndex from slotCellIndex (" & $input & " -> " & $expected & ")": + let + slotCellIndex = input.uint64 + blockCellIndex = dataSampler.getBlockCellIndexForSlotCellIndex(slotCellIndex) - proc toStr(proof: MerkleProof): string = - toHex(proof.nodesBuffer) + check: + blockCellIndex == expected.uint64 - let - expectedSlotToBlockProofs = getExpectedSlotToBlockProofs() - expectedBlockToCellProofs = getExpectedBlockToCellProofs() - expectedSampleData = getExpectedSampleData() + let knownIndices = @[178.uint64, 277.uint64, 366.uint64] + + test "Can find single slot-cell index": + proc slotCellIndex(i: int): DSSlotCellIndex = + let counter: DSFieldElement = toF(i) + return dataSampler.findSlotCellIndex(challenge, counter) + + proc getExpectedIndex(i: int): DSSlotCellIndex = + let + numberOfCellsInSlot = (bytesPerBlock * numberOfSlotBlocks) div CellSize.int + hash = Sponge.digest(@[slotRootHash, challenge, toF(i)], rate = 2) + return extractLowBits(hash.toBig(), ceilingLog2(numberOfCellsInSlot)) check: - a.datasetToSlotProof == datasetToSlotProof - a.slotToBlockProofs.mapIt(toStr(it)) == expectedSlotToBlockProofs - a.blockToCellProofs.mapIt(toStr(it)) == expectedBlockToCellProofs - toHex(a.sampleData) == expectedSampleData + slotCellIndex(1) == getExpectedIndex(1) + slotCellIndex(1) == knownIndices[0] + slotCellIndex(2) == getExpectedIndex(2) + slotCellIndex(2) == knownIndices[1] + slotCellIndex(3) == getExpectedIndex(3) + slotCellIndex(3) == knownIndices[2] + + # test "Can find sequence of slot-cell indices": + # proc slotCellIndices(n: int): seq[DSSlotCellIndex] = + # dataSampler.findSlotCellIndices(challenge, n) + + # proc getExpectedIndices(n: int): seq[DSSlotCellIndex] = + # return collect(newSeq, (for i in 1..n: dataSampler.findSlotCellIndex(challenge, toF(i)))) + + # check: + # slotCellIndices(3) == getExpectedIndices(3) + # slotCellIndices(3) == knownIndices + + # test "Can get cell from block": + # let + # bytes = newSeqWith(bytesPerBlock, rand(uint8)) + # blk = bt.Block.new(bytes).tryGet() + + # sample0 = dataSampler.getCellFromBlock(blk, 0) + # sample1 = dataSampler.getCellFromBlock(blk, 1) + # sample2 = dataSampler.getCellFromBlock(blk, 2) + + # check: + # sample0 == bytes[0.. " & $expected & ")": - let - slotCellIndex = input.uint64 - slotBlockIndex = getSlotBlockIndexForSlotCellIndex(slotCellIndex, blockSize) - - check: - slotBlockIndex == expected.uint64 - - for input in 0 ..< numberOfSlotBlocks: - test "Can get datasetBlockIndex from slotBlockIndex (" & $input & ")": - let - slotBlockIndex = input.uint64 - datasetBlockIndex = getDatasetBlockIndexForSlotBlockIndex(slot, blockSize, slotBlockIndex) - datasetSlotIndex = slot.slotIndex.truncate(uint64) - expectedIndex = (numberOfSlotBlocks.uint64 * datasetSlotIndex) + slotBlockIndex - - check: - datasetBlockIndex == expectedIndex - - for (input, expected) in [(10, 10), (31, 31), (32, 0), (63, 31), (64, 0)]: - test "Can get blockCellIndex from slotCellIndex (" & $input & " -> " & $expected & ")": - let - slotCellIndex = input.uint64 - - blockCellIndex = getBlockCellIndexForSlotCellIndex(slotCellIndex, blockSize) - - check: - blockCellIndex == expected.uint64 diff --git a/tests/codex/proof/testslotblocks.nim b/tests/codex/proof/testslotblocks.nim index 289444bd..e0301b8c 100644 --- a/tests/codex/proof/testslotblocks.nim +++ b/tests/codex/proof/testslotblocks.nim @@ -12,7 +12,6 @@ import pkg/codex/contracts import pkg/codex/merkletree import pkg/codex/proof/slotblocks -import pkg/codex/proof/indexing import ../helpers import ../examples @@ -47,15 +46,20 @@ asyncchecksuite "Test slotblocks - manifest": setup: discard await localStore.putBlock(manifestBlock) + proc getManifest(store: BlockStore): Future[?!Manifest] {.async.} = + without slotBlocks =? await SlotBlocks.new(slot, store), err: + return failure(err) + return success(slotBlocks.manifest) + test "Can get manifest for slot": - let m = (await getManifestForSlot(slot, localStore)).tryGet() + let m = (await getManifest(localStore)).tryGet() check: m.treeCid == manifest.treeCid test "Can fail to get manifest for invalid cid": slot.request.content.cid = "invalid" - let m = (await getManifestForSlot(slot, localStore)) + let m = (await getManifest(localStore)) check: m.isErr @@ -63,7 +67,7 @@ asyncchecksuite "Test slotblocks - manifest": test "Can fail to get manifest when manifest block not found": let emptyStore = CacheStore.new() - m = (await getManifestForSlot(slot, emptyStore)) + m = (await getManifest(emptyStore)) check: m.isErr @@ -71,7 +75,7 @@ asyncchecksuite "Test slotblocks - manifest": test "Can fail to get manifest when manifest fails to decode": manifestBlock.data = @[] - let m = (await getManifestForSlot(slot, localStore)) + let m = (await getManifest(localStore)) check: m.isErr @@ -90,6 +94,7 @@ asyncchecksuite "Test slotblocks - slot blocks by index": manifestBlock: bt.Block slot: Slot datasetBlocks: seq[bt.Block] + slotBlocks: SlotBlocks proc createDatasetBlocks(): Future[void] {.async.} = while true: @@ -129,17 +134,33 @@ asyncchecksuite "Test slotblocks - slot blocks by index": slotIndex: u256(datasetSlotIndex) ) + proc createSlotBlocks(): Future[void] {.async.} = + slotBlocks = (await SlotBlocks.new(slot, localStore)).tryGet() + setup: await createDatasetBlocks() await createManifest() createSlot() discard await localStore.putBlock(manifestBlock) + await createSlotBlocks() + + for input in 0 ..< numberOfSlotBlocks: + test "Can get datasetBlockIndex from slotBlockIndex (" & $input & ")": + let + slotBlockIndex = input.uint64 + datasetBlockIndex = slotBlocks.getDatasetBlockIndexForSlotBlockIndex(slotBlockIndex) + datasetSlotIndex = slot.slotIndex.truncate(uint64) + expectedIndex = (numberOfSlotBlocks.uint64 * datasetSlotIndex) + slotBlockIndex + + check: + datasetBlockIndex == expectedIndex for input in [0, 1, numberOfSlotBlocks-1]: test "Can get slot block by index (" & $input & ")": let - slotBlock = (await getSlotBlock(slot, localStore, input.uint64)).tryget() - expectedDatasetBlockIndex = getDatasetBlockIndexForSlotBlockIndex(slot, bytesPerBlock.uint64, input.uint64) + slotBlockIndex = input.uint64 + slotBlock = (await slotBlocks.getSlotBlock(slotBlockIndex)).tryget() + expectedDatasetBlockIndex = slotBlocks.getDatasetBlockIndexForSlotBlockIndex(slotBlockIndex) expectedBlock = datasetBlocks[expectedDatasetBlockIndex] check: @@ -148,8 +169,8 @@ asyncchecksuite "Test slotblocks - slot blocks by index": test "Can fail to get block when index is out of range": let - b1 = await getSlotBlock(slot, localStore, numberOfSlotBlocks.uint64) - b2 = await getSlotBlock(slot, localStore, (numberOfSlotBlocks + 1).uint64) + b1 = await slotBlocks.getSlotBlock(numberOfSlotBlocks.uint64) + b2 = await slotBlocks.getSlotBlock((numberOfSlotBlocks + 1).uint64) check: b1.isErr