From 9e0e9752794a09bb99b1142059150fa9a4e0737d Mon Sep 17 00:00:00 2001 From: Dmitriy Ryajov Date: Mon, 22 Jan 2024 13:54:44 -0600 Subject: [PATCH] generify sampler and builder --- codex/node.nim | 8 +-- codex/slots.nim | 5 ++ codex/slots/builder/builder.nim | 98 +++++++++++++-------------- codex/slots/sampler/sampler.nim | 64 ++++++++--------- tests/codex/node/testcontracts.nim | 4 +- tests/codex/node/testnode.nim | 2 +- tests/codex/slots/testsampler.nim | 10 +-- tests/codex/slots/testslotbuilder.nim | 36 +++++----- 8 files changed, 116 insertions(+), 111 deletions(-) diff --git a/codex/node.nim b/codex/node.nim index 40e8ee26..17933055 100644 --- a/codex/node.nim +++ b/codex/node.nim @@ -388,7 +388,7 @@ proc setupRequest( trace "Unable to erasure code dataset" return failure(error) - without builder =? SlotsBuilder.new(self.blockStore, encoded), err: + without builder =? Poseidon2Builder.new(self.blockStore, encoded), err: trace "Unable to create slot builder" return failure(err) @@ -501,7 +501,7 @@ proc onStore( trace "Unable to fetch manifest for cid", cid, err = err.msg return failure(err) - without builder =? SlotsBuilder.new(self.blockStore, manifest), err: + without builder =? Poseidon2Builder.new(self.blockStore, manifest), err: trace "Unable to create slots builder", err = err.msg return failure(err) @@ -570,11 +570,11 @@ proc onProve( error "Unable to fetch manifest for cid", err = err.msg return failure(err) - without builder =? SlotsBuilder.new(self.blockStore, manifest), err: + without builder =? Poseidon2Builder.new(self.blockStore, manifest), err: error "Unable to create slots builder", err = err.msg return failure(err) - without sampler =? DataSampler.new(slotIdx, self.blockStore, builder), err: + without sampler =? Poseidon2Sampler.new(slotIdx, self.blockStore, builder), err: error "Unable to create data sampler", err = err.msg return failure(err) diff --git a/codex/slots.nim b/codex/slots.nim index fb0827a4..262b0a6e 100644 --- a/codex/slots.nim +++ b/codex/slots.nim @@ -1,4 +1,9 @@ import ./slots/builder import ./slots/sampler +import ./merkletree export builder, sampler + +type + Poseidon2Builder* = SlotsBuilder[Poseidon2Tree, Poseidon2Hash] + Poseidon2Sampler* = DataSampler[Poseidon2Tree, Poseidon2Hash, Poseidon2Proof] diff --git a/codex/slots/builder/builder.nim b/codex/slots/builder/builder.nim index 1f15c3a1..f7084f85 100644 --- a/codex/slots/builder/builder.nim +++ b/codex/slots/builder/builder.nim @@ -32,7 +32,7 @@ import ../../utils/digest import ../../utils/poseidon2digest import ../converters -export converters +export converters, asynciter logScope: topics = "codex slotsbuilder" @@ -47,31 +47,31 @@ const type # TODO: should be a generic type that # supports all merkle trees - SlotsBuilder* = ref object of RootObj + SlotsBuilder*[T, H] = ref object of RootObj store: BlockStore manifest: Manifest strategy: IndexingStrategy cellSize: NBytes - emptyDigestTree: Poseidon2Tree + emptyDigestTree: T blockPadBytes: seq[byte] - slotsPadLeafs: seq[Poseidon2Hash] - rootsPadLeafs: seq[Poseidon2Hash] - slotRoots: seq[Poseidon2Hash] - verifyTree: ?Poseidon2Tree + slotsPadLeafs: seq[H] + rootsPadLeafs: seq[H] + slotRoots: seq[H] + verifyTree: ?T -func slotRoots*(self: SlotsBuilder): seq[Poseidon2Hash] {.inline.} = +func slotRoots*[T, H](self: SlotsBuilder[T, H]): seq[H] {.inline.} = ## Returns the slot roots. ## self.slotRoots -func verifyTree*(self: SlotsBuilder): ?Poseidon2Tree {.inline.} = +func verifyTree*[T, H](self: SlotsBuilder[T, H]): ?H {.inline.} = ## Returns the slots tree (verification tree). ## self.verifyTree -func verifyRoot*(self: SlotsBuilder): ?Poseidon2Hash {.inline.} = +func verifyRoot*[T, H](self: SlotsBuilder[T, H]): ?H {.inline.} = ## Returns the slots root (verification root). ## @@ -84,70 +84,70 @@ func nextPowerOfTwoPad*(a: int): int = nextPowerOfTwo(a) - a -func numBlockPadBytes*(self: SlotsBuilder): Natural = +func numBlockPadBytes*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of padding bytes required for a pow2 ## merkle tree for each block. ## self.blockPadBytes.len -func numSlotsPadLeafs*(self: SlotsBuilder): Natural = +func numSlotsPadLeafs*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of padding field elements required for a pow2 ## merkle tree for each slot. ## self.slotsPadLeafs.len -func numRootsPadLeafs*(self: SlotsBuilder): Natural = +func numRootsPadLeafs*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of padding field elements required for a pow2 ## merkle tree for the slot roots. ## self.rootsPadLeafs.len -func numSlots*(self: SlotsBuilder): Natural = +func numSlots*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of slots. ## self.manifest.numSlots -func numSlotBlocks*(self: SlotsBuilder): Natural = +func numSlotBlocks*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of blocks per slot. ## self.manifest.blocksCount div self.manifest.numSlots -func slotBytes*(self: SlotsBuilder): NBytes = +func slotBytes*[T, H](self: SlotsBuilder[T, H]): NBytes = ## Number of bytes per slot. ## (self.manifest.blockSize.int * self.numSlotBlocks).NBytes -func numBlockCells*(self: SlotsBuilder): Natural = +func numBlockCells*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of cells per block. ## (self.manifest.blockSize div self.cellSize).Natural -func cellSize*(self: SlotsBuilder): NBytes = +func cellSize*[T, H](self: SlotsBuilder[T, H]): NBytes = ## Cell size. ## self.cellSize -func numSlotCells*(self: SlotsBuilder): Natural = +func numSlotCells*[T, H](self: SlotsBuilder[T, H]): Natural = ## Number of cells per slot. ## self.numBlockCells * self.numSlotBlocks -func slotIndiciesIter*(self: SlotsBuilder, slot: Natural): ?!Iter[int] = +func slotIndiciesIter*[T, H](self: SlotsBuilder[T, H], slot: Natural): ?!Iter[int] = ## Returns the slot indices. ## self.strategy.getIndicies(slot).catch -func slotIndicies*(self: SlotsBuilder, slot: Natural): seq[int] = +func slotIndicies*[T, H](self: SlotsBuilder[T, H], slot: Natural): seq[int] = ## Returns the slot indices. ## @@ -157,15 +157,15 @@ func slotIndicies*(self: SlotsBuilder, slot: Natural): seq[int] = trace "Failed to get slot indicies" newSeq[int]() -func manifest*(self: SlotsBuilder): Manifest = +func manifest*[T, H](self: SlotsBuilder[T, H]): Manifest = ## Returns the manifest. ## self.manifest -proc buildBlockTree*( - self: SlotsBuilder, - blkIdx: Natural): Future[?!(seq[byte], Poseidon2Tree)] {.async.} = +proc buildBlockTree*[T, H]( + self: SlotsBuilder[T, H], + blkIdx: Natural): Future[?!(seq[byte], T)] {.async.} = without blk =? await self.store.getBlock(self.manifest.treeCid, blkIdx), err: error "Failed to get block CID for tree at index" return failure(err) @@ -174,15 +174,15 @@ proc buildBlockTree*( success (DefaultEmptyBlock & self.blockPadBytes, self.emptyDigestTree) else: without tree =? - Poseidon2Tree.digestTree(blk.data & self.blockPadBytes, self.cellSize.int), err: + T.digestTree(blk.data & self.blockPadBytes, self.cellSize.int), err: error "Failed to create digest for block" return failure(err) success (blk.data, tree) -proc getCellHashes*( - self: SlotsBuilder, - slotIndex: Natural): Future[?!seq[Poseidon2Hash]] {.async.} = +proc getCellHashes*[T, H]( + self: SlotsBuilder[T, H], + slotIndex: Natural): Future[?!seq[H]] {.async.} = let treeCid = self.manifest.treeCid @@ -197,7 +197,7 @@ proc getCellHashes*( slotIndex = slotIndex let - hashes: seq[Poseidon2Hash] = collect(newSeq): + hashes: seq[H] = collect(newSeq): for blkIdx in self.strategy.getIndicies(slotIndex): trace "Getting block CID for tree at index" @@ -210,18 +210,18 @@ proc getCellHashes*( success hashes -proc buildSlotTree*( - self: SlotsBuilder, - slotIndex: Natural): Future[?!Poseidon2Tree] {.async.} = +proc buildSlotTree*[T, H]( + self: SlotsBuilder[T, H], + slotIndex: Natural): Future[?!T] {.async.} = without cellHashes =? (await self.getCellHashes(slotIndex)), err: error "Failed to select slot blocks", err = err.msg return failure(err) - Poseidon2Tree.init(cellHashes & self.slotsPadLeafs) + T.init(cellHashes & self.slotsPadLeafs) -proc buildSlot*( - self: SlotsBuilder, - slotIndex: Natural): Future[?!Poseidon2Hash] {.async.} = +proc buildSlot*[T, H]( + self: SlotsBuilder[T, H], + slotIndex: Natural): Future[?!H] {.async.} = ## Build a slot tree and store it in the block store. ## @@ -254,12 +254,12 @@ proc buildSlot*( tree.root() -func buildVerifyTree*( - self: SlotsBuilder, - slotRoots: openArray[Poseidon2Hash]): ?!Poseidon2Tree = - Poseidon2Tree.init(@slotRoots & self.rootsPadLeafs) +func buildVerifyTree*[T, H]( + self: SlotsBuilder[T, H], + slotRoots: openArray[H]): ?!T = + T.init(@slotRoots & self.rootsPadLeafs) -proc buildSlots*(self: SlotsBuilder): Future[?!void] {.async.} = +proc buildSlots*[T, H](self: SlotsBuilder[T, H]): Future[?!void] {.async.} = ## Build all slot trees and store them in the block store. ## @@ -289,7 +289,7 @@ proc buildSlots*(self: SlotsBuilder): Future[?!void] {.async.} = success() -proc buildManifest*(self: SlotsBuilder): Future[?!Manifest] {.async.} = +proc buildManifest*[T, H](self: SlotsBuilder[T, H]): Future[?!Manifest] {.async.} = if err =? (await self.buildSlots()).errorOption: error "Failed to build slot roots", err = err.msg return failure(err) @@ -305,12 +305,12 @@ proc buildManifest*(self: SlotsBuilder): Future[?!Manifest] {.async.} = Manifest.new(self.manifest, rootProvingCid, rootCids) -proc new*( - T: type SlotsBuilder, +proc new*[T, H]( + _: type SlotsBuilder[T, H], store: BlockStore, manifest: Manifest, strategy: ?IndexingStrategy = none IndexingStrategy, - cellSize = DefaultCellSize): ?!SlotsBuilder = + cellSize = DefaultCellSize): ?!SlotsBuilder[T, H] = if not manifest.protected: return failure("Can only create SlotsBuilder using protected manifests.") @@ -334,9 +334,9 @@ proc new*( numSlotLeafs = (manifest.blocksCount div manifest.numSlots) slotsPadLeafs = newSeqWith(numSlotLeafs.nextPowerOfTwoPad, Poseidon2Zero) # power of two padding for block roots rootsPadLeafs = newSeqWith(manifest.numSlots.nextPowerOfTwoPad, Poseidon2Zero) - emptyDigestTree = ? Poseidon2Tree.digestTree(DefaultEmptyBlock & blockPadBytes, DefaultCellSize.int) + emptyDigestTree = ? T.digestTree(DefaultEmptyBlock & blockPadBytes, DefaultCellSize.int) - var self = SlotsBuilder( + var self = SlotsBuilder[T, H]( store: store, manifest: manifest, strategy: strategy, @@ -350,7 +350,7 @@ proc new*( if manifest.slotRoots.len == 0 or manifest.slotRoots.len != manifest.numSlots: return failure "Manifest is verifiable but slot roots are missing or invalid." - let slotRoots = manifest.slotRoots.mapIt( (? it.fromSlotCid())) + let slotRoots = manifest.slotRoots.mapIt( (? it.fromSlotCid() )) without tree =? self.buildVerifyTree(slotRoots), err: error "Failed to build slot roots tree", err = err.msg diff --git a/codex/slots/sampler/sampler.nim b/codex/slots/sampler/sampler.nim index 08dd8ef9..7e53f8bc 100644 --- a/codex/slots/sampler/sampler.nim +++ b/codex/slots/sampler/sampler.nim @@ -36,59 +36,44 @@ logScope: type Cell* = seq[byte] - Sample* = object + Sample*[P] = object data*: Cell - slotProof*: Poseidon2Proof - cellProof*: Poseidon2Proof + slotProof*: P + cellProof*: P slotBlockIdx*: Natural blockCellIdx*: Natural - ProofInput* = object - entropy*: Poseidon2Hash - verifyRoot*: Poseidon2Hash - verifyProof*: Poseidon2Proof + ProofInput*[H, P] = object + entropy*: H + verifyRoot*: H + verifyProof*: P numSlots*: Natural numCells*: Natural slotIndex*: Natural - samples*: seq[Sample] + samples*: seq[Sample[P]] - DataSampler* = ref object of RootObj + DataSampler*[T, H, P] = ref object of RootObj index: Natural blockStore: BlockStore # The following data is invariant over time for a given slot: - builder: SlotsBuilder + builder: SlotsBuilder[T, H] -proc new*( - T: type DataSampler, - index: Natural, - blockStore: BlockStore, - builder: SlotsBuilder): ?!DataSampler = - - if index > builder.slotRoots.high: - error "Slot index is out of range" - return failure("Slot index is out of range") - - success DataSampler( - index: index, - blockStore: blockStore, - builder: builder) - -proc getCell*(self: DataSampler, blkBytes: seq[byte], blkCellIdx: Natural): Cell = +proc getCell*[T, H, P](self: DataSampler[T, H, P], blkBytes: seq[byte], blkCellIdx: Natural): Cell = let cellSize = self.builder.cellSize.uint64 dataStart = cellSize * blkCellIdx.uint64 dataEnd = dataStart + cellSize return blkBytes[dataStart ..< dataEnd] -proc getProofInput*( - self: DataSampler, +proc getProofInput*[T, H, P]( + self: DataSampler[T, H, P], entropy: ProofChallenge, - nSamples: Natural): Future[?!ProofInput] {.async.} = + nSamples: Natural): Future[?!ProofInput[H, P]] {.async.} = ## Generate proofs as input to the proving circuit. ## let - entropy = Poseidon2Hash.fromBytes( + entropy = H.fromBytes( array[31, byte].initCopyFrom(entropy[0..30])) # truncate to 31 bytes, otherwise it _might_ be greater than mod without verifyTree =? self.builder.verifyTree and @@ -146,14 +131,14 @@ proc getProofInput*( error "Failed to get proof from block tree", err = err.msg return failure(err) - Sample( + Sample[P]( data: self.getCell(bytes, blkCellIdx), slotProof: slotProof, cellProof: blockProof, slotBlockIdx: slotCellIdx.Natural, blockCellIdx: blkCellIdx.Natural) - success ProofInput( + success ProofInput[H, P]( entropy: entropy, verifyRoot: verifyRoot, verifyProof: verifyProof, @@ -161,3 +146,18 @@ proc getProofInput*( numCells: self.builder.numSlotCells, slotIndex: self.index, samples: samples) + +proc new*[T, H, P]( + _: type DataSampler[T, H, P], + index: Natural, + blockStore: BlockStore, + builder: SlotsBuilder[T, H]): ?!DataSampler[T, H, P] = + + if index > builder.slotRoots.high: + error "Slot index is out of range" + return failure("Slot index is out of range") + + success DataSampler[T, H, P]( + index: index, + blockStore: blockStore, + builder: builder) diff --git a/tests/codex/node/testcontracts.nim b/tests/codex/node/testcontracts.nim index 05e5a047..d31e6280 100644 --- a/tests/codex/node/testcontracts.nim +++ b/tests/codex/node/testcontracts.nim @@ -54,7 +54,7 @@ asyncchecksuite "Test Node - Host contracts": manifestCidStr: string manifestCid: Cid market: MockMarket - builder: SlotsBuilder + builder: Poseidon2Builder verifiable: Manifest verifiableBlock: bt.Block protected: Manifest @@ -84,7 +84,7 @@ asyncchecksuite "Test Node - Host contracts": (await localStore.putBlock(manifestBlock)).tryGet() protected = (await erasure.encode(manifest, 3, 2)).tryGet() - builder = SlotsBuilder.new(localStore, protected).tryGet() + builder = Poseidon2Builder.new(localStore, protected).tryGet() verifiable = (await builder.buildManifest()).tryGet() verifiableBlock = bt.Block.new( verifiable.encode().tryGet(), diff --git a/tests/codex/node/testnode.nim b/tests/codex/node/testnode.nim index 4827d530..7bfbfe35 100644 --- a/tests/codex/node/testnode.nim +++ b/tests/codex/node/testnode.nim @@ -126,7 +126,7 @@ asyncchecksuite "Test Node - Basic": codec = ManifestCodec).tryGet() protected = (await erasure.encode(manifest, 3, 2)).tryGet() - builder = SlotsBuilder.new(localStore, protected).tryGet() + builder = Poseidon2Builder.new(localStore, protected).tryGet() verifiable = (await builder.buildManifest()).tryGet() verifiableBlock = bt.Block.new( verifiable.encode().tryGet(), diff --git a/tests/codex/slots/testsampler.nim b/tests/codex/slots/testsampler.nim index 2ca2b40f..b4e3afc6 100644 --- a/tests/codex/slots/testsampler.nim +++ b/tests/codex/slots/testsampler.nim @@ -19,7 +19,7 @@ import pkg/codex/contracts import pkg/codex/merkletree import pkg/codex/stores/cachestore -import pkg/codex/slots/sampler +import pkg/codex/slots import pkg/codex/slots/builder/builder import ../helpers @@ -28,20 +28,20 @@ import ../merkletree/helpers import testsampler_expected import ./provingtestenv -asyncchecksuite "Test DataSampler": +asyncchecksuite "Test Poseidon2Sampler": var env: ProvingTestEnvironment - dataSampler: DataSampler + dataSampler: Poseidon2Sampler blk: bt.Block cell0Bytes: seq[byte] cell1Bytes: seq[byte] cell2Bytes: seq[byte] proc createDataSampler(): Future[void] {.async.} = - dataSampler = DataSampler.new( + dataSampler = Poseidon2Sampler.new( datasetSlotIndex, env.localStore, - SlotsBuilder.new(env.localStore, env.manifest).tryGet()).tryGet() + Poseidon2Builder.new(env.localStore, env.manifest).tryGet()).tryGet() setup: randomize() diff --git a/tests/codex/slots/testslotbuilder.nim b/tests/codex/slots/testslotbuilder.nim index 8022e07a..869023d4 100644 --- a/tests/codex/slots/testslotbuilder.nim +++ b/tests/codex/slots/testslotbuilder.nim @@ -61,7 +61,7 @@ suite "Slot builder": manifest: Manifest protectedManifest: Manifest expectedEmptyCid: Cid - slotBuilder: SlotsBuilder + slotBuilder: Poseidon2Builder chunker: Chunker proc createBlocks(): Future[void] {.async.} = @@ -153,7 +153,7 @@ suite "Slot builder": datasetSize = originalDatasetSize.NBytes) check: - SlotsBuilder.new(localStore, unprotectedManifest, cellSize = cellSize) + Poseidon2Builder.new(localStore, unprotectedManifest, cellSize = cellSize) .error.msg == "Can only create SlotsBuilder using protected manifests." test "Number of blocks must be devisable by number of slots": @@ -169,7 +169,7 @@ suite "Slot builder": ecM = ecM) check: - SlotsBuilder.new(localStore, mismatchManifest, cellSize = cellSize) + Poseidon2Builder.new(localStore, mismatchManifest, cellSize = cellSize) .error.msg == "Number of blocks must be divisable by number of slots." test "Block size must be divisable by cell size": @@ -185,11 +185,11 @@ suite "Slot builder": ecM = ecM) check: - SlotsBuilder.new(localStore, mismatchManifest, cellSize = cellSize) + Poseidon2Builder.new(localStore, mismatchManifest, cellSize = cellSize) .error.msg == "Block size must be divisable by cell size." test "Should build correct slot builder": - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -202,7 +202,7 @@ suite "Slot builder": test "Should build slot hashes for all slots": let steppedStrategy = SteppedIndexingStrategy.new(0, numTotalBlocks - 1, numSlots) - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -225,7 +225,7 @@ suite "Slot builder": test "Should build slot trees for all slots": let steppedStrategy = SteppedIndexingStrategy.new(0, numTotalBlocks - 1, numSlots) - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -248,7 +248,7 @@ suite "Slot builder": test "Should persist trees for all slots": let - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -273,7 +273,7 @@ suite "Slot builder": test "Should build correct verification root": let steppedStrategy = SteppedIndexingStrategy.new(0, numTotalBlocks - 1, numSlots) - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -302,7 +302,7 @@ suite "Slot builder": test "Should build correct verification root manifest": let steppedStrategy = SteppedIndexingStrategy.new(0, numTotalBlocks - 1, numSlots) - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -331,21 +331,21 @@ suite "Slot builder": test "Should not build from verifiable manifest with 0 slots": var - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() verifyManifest = (await slotBuilder.buildManifest()).tryGet() verifyManifest.slotRoots = @[] - check SlotsBuilder.new( + check Poseidon2Builder.new( localStore, verifyManifest, cellSize = cellSize).isErr test "Should not build from verifiable manifest with incorrect number of slots": var - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -356,14 +356,14 @@ suite "Slot builder": verifyManifest.slotRoots.len - 1 ) - check SlotsBuilder.new( + check Poseidon2Builder.new( localStore, verifyManifest, cellSize = cellSize).isErr test "Should not build from verifiable manifest with invalid verify root": let - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() @@ -375,21 +375,21 @@ suite "Slot builder": Rng.instance, verifyManifest.verifyRoot.data.buffer) - check SlotsBuilder.new( + check Poseidon2Builder.new( localStore, verifyManifest, cellSize = cellSize).isErr test "Should build from verifiable manifest": let - slotBuilder = SlotsBuilder.new( + slotBuilder = Poseidon2Builder.new( localStore, protectedManifest, cellSize = cellSize).tryGet() verifyManifest = (await slotBuilder.buildManifest()).tryGet() - verificationBuilder = SlotsBuilder.new( + verificationBuilder = Poseidon2Builder.new( localStore, verifyManifest, cellSize = cellSize).tryGet()