logos-storage-nim/tests/codex/slots/proof/testdatasampler.nim
2024-01-15 21:48:54 -06:00

330 lines
11 KiB
Nim

# import std/sequtils
# import std/sugar
# import std/random
# import pkg/questionable/results
# import pkg/constantine/math/arithmetic
# import pkg/constantine/math/io/io_fields
# import pkg/poseidon2/types
# import pkg/poseidon2/io
# import pkg/poseidon2
# import pkg/chronos
# import pkg/asynctest
# import pkg/codex/stores/cachestore
# import pkg/codex/chunker
# import pkg/codex/stores
# import pkg/codex/blocktype as bt
# import pkg/codex/contracts/requests
# import pkg/codex/contracts
# import pkg/codex/merkletree
# import pkg/codex/stores/cachestore
# import pkg/codex/slots
# import pkg/codex/proof/misc
# import pkg/codex/proof/types
# import ../helpers
# import ../examples
# import testdatasampler_expected
# let
# bytesPerBlock = 64 * 1024
# challenge: FieldElement = toF(12345)
# datasetRootHash: FieldElement = toF(6789)
# asyncchecksuite "Test proof datasampler - components":
# let
# numberOfSlotBlocks = 16
# slot = Slot(
# request: StorageRequest(
# ask: StorageAsk(
# slots: 10,
# slotSize: u256(bytesPerBlock * numberOfSlotBlocks),
# ),
# content: StorageContent(
# cid: $Cid.example
# )
# ),
# slotIndex: u256(3)
# )
# test "Number of cells is a power of two":
# # This is to check that the data used for testing is sane.
# proc isPow2(value: int): bool =
# let log2 = ceilingLog2(value)
# return (1 shl log2) == value
# let numberOfCells = getNumberOfCellsInSlot(slot).int
# check:
# isPow2(numberOfCells)
# test "Extract low bits":
# proc extract(value: uint64, nBits: int): uint64 =
# let big = toF(value).toBig()
# return extractLowBits(big, nBits)
# check:
# extract(0x88, 4) == 0x8.uint64
# extract(0x88, 7) == 0x8.uint64
# extract(0x9A, 5) == 0x1A.uint64
# extract(0x9A, 7) == 0x1A.uint64
# extract(0x1248, 10) == 0x248.uint64
# extract(0x1248, 12) == 0x248.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
# slotSizeInBytes = (slot.request.ask.slotSize).truncate(uint64)
# expectedNumberOfCells = slotSizeInBytes div CellSize
# check:
# expectedNumberOfCells == 512
# expectedNumberOfCells == getNumberOfCellsInSlot(slot)
# asyncchecksuite "Test proof datasampler - main":
# let
# # The number of slot blocks and number of slots, combined with
# # the bytes per block, make it so that there are exactly 256 cells
# # in the dataset.
# numberOfSlotBlocks = 4
# totalNumberOfSlots = 2
# datasetSlotIndex = 1
# localStore = CacheStore.new()
# datasetToSlotProof = MerkleProof.example
# var
# manifest: Manifest
# manifestBlock: bt.Block
# slot: Slot
# datasetBlocks: seq[bt.Block]
# slotPoseidonTree: MerkleTree
# dataSampler: DataSampler
# proc createDatasetBlocks(): Future[void] {.async.} =
# let numberOfCellsNeeded = (numberOfSlotBlocks * totalNumberOfSlots * bytesPerBlock).uint64 div CellSize
# var data: seq[byte] = @[]
# # This generates a number of blocks that have different data, such that
# # Each cell in each block is unique, but nothing is random.
# for i in 0 ..< numberOfCellsNeeded:
# data = data & (i.byte).repeat(CellSize)
# let chunker = MockChunker.new(
# dataset = data,
# chunkSize = bytesPerBlock)
# while true:
# let chunk = await chunker.getBytes()
# if chunk.len <= 0:
# break
# let b = bt.Block.new(chunk).tryGet()
# datasetBlocks.add(b)
# discard await localStore.putBlock(b)
# proc createManifest(): Future[void] {.async.} =
# let
# cids = datasetBlocks.mapIt(it.cid)
# tree = MerkleTree.init(cids).tryGet()
# treeCid = tree.rootCid().tryGet()
# for index, cid in cids:
# let proof = tree.getProof(index).tryget()
# discard await localStore.putBlockCidAndProof(treeCid, index, cid, proof)
# manifest = Manifest.new(
# treeCid = treeCid,
# blockSize = bytesPerBlock.NBytes,
# datasetSize = (bytesPerBlock * numberOfSlotBlocks * totalNumberOfSlots).NBytes)
# manifestBlock = bt.Block.new(manifest.encode().tryGet(), codec = DagPBCodec).tryGet()
# proc createSlot(): void =
# slot = Slot(
# request: StorageRequest(
# ask: StorageAsk(
# slotSize: u256(bytesPerBlock * numberOfSlotBlocks)
# ),
# content: StorageContent(
# cid: $manifestBlock.cid
# ),
# ),
# slotIndex: u256(datasetSlotIndex)
# )
# proc createSlotPoseidonTree(): void =
# let
# slotSize = slot.request.ask.slotSize.truncate(uint64)
# blocksInSlot = slotSize div bytesPerBlock.uint64
# datasetSlotIndex = slot.slotIndex.truncate(uint64)
# datasetBlockIndexFirst = datasetSlotIndex * blocksInSlot
# datasetBlockIndexLast = datasetBlockIndexFirst + numberOfSlotBlocks.uint64
# slotBlocks = datasetBlocks[datasetBlockIndexFirst ..< datasetBlockIndexLast]
# slotBlockCids = slotBlocks.mapIt(it.cid)
# slotPoseidonTree = MerkleTree.init(slotBlockCids).tryGet()
# proc createDataSampler(): Future[void] {.async.} =
# dataSampler = (await DataSampler.new(
# slot,
# localStore,
# datasetRootHash,
# slotPoseidonTree,
# datasetToSlotProof
# )).tryGet()
# setup:
# await createDatasetBlocks()
# await createManifest()
# createSlot()
# discard await localStore.putBlock(manifestBlock)
# createSlotPoseidonTree()
# await createDataSampler()
# test "Number of cells is a power of two":
# # This is to check that the data used for testing is sane.
# proc isPow2(value: int): bool =
# let log2 = ceilingLog2(value)
# return (1 shl log2) == value
# let numberOfCells = getNumberOfCellsInSlot(slot).int
# check:
# isPow2(numberOfCells)
# let knownIndices = @[74.uint64, 41.uint64, 51.uint64]
# test "Can find single slot-cell index":
# proc slotCellIndex(i: int): uint64 =
# let counter: FieldElement = toF(i)
# return dataSampler.findSlotCellIndex(challenge, counter)
# proc getExpectedIndex(i: int): uint64 =
# let
# numberOfCellsInSlot = (bytesPerBlock * numberOfSlotBlocks) div CellSize.int
# slotRootHash = toF(1234) # TODO - replace with slotPoseidonTree.root when it is a poseidon tree.
# hash = Sponge.digest(@[slotRootHash, challenge, toF(i)], rate = 2)
# return extractLowBits(hash.toBig(), ceilingLog2(numberOfCellsInSlot))
# 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[uint64] =
# dataSampler.findSlotCellIndices(challenge, n)
# proc getExpectedIndices(n: int): seq[uint64] =
# return collect(newSeq, (for i in 1..n: dataSampler.findSlotCellIndex(challenge, toF(i))))
# check:
# slotCellIndices(3) == getExpectedIndices(3)
# slotCellIndices(3) == knownIndices
# let
# bytes = newSeqWith(bytesPerBlock, rand(uint8))
# blk = bt.Block.new(bytes).tryGet()
# cell0Bytes = bytes[0..<CellSize]
# cell1Bytes = bytes[CellSize..<(CellSize*2)]
# cell2Bytes = bytes[(CellSize*2)..<(CellSize*3)]
# test "Can get cell from block":
# let
# sample0 = dataSampler.getCellFromBlock(blk, 0)
# sample1 = dataSampler.getCellFromBlock(blk, 1)
# sample2 = dataSampler.getCellFromBlock(blk, 2)
# check:
# sample0 == cell0Bytes
# sample1 == cell1Bytes
# sample2 == cell2Bytes
# test "Can convert block into cells":
# let cells = dataSampler.getBlockCells(blk)
# check:
# cells.len == (bytesPerBlock div CellSize.int)
# cells[0] == cell0Bytes
# cells[1] == cell1Bytes
# cells[2] == cell2Bytes
# test "Can create mini tree for block cells":
# let miniTree = dataSampler.getBlockCellMiniTree(blk).tryGet()
# let
# cell0Proof = miniTree.getProof(0).tryGet()
# cell1Proof = miniTree.getProof(1).tryGet()
# cell2Proof = miniTree.getProof(2).tryGet()
# check:
# cell0Proof.verifyDataBlock(cell0Bytes, miniTree.root).tryGet()
# cell1Proof.verifyDataBlock(cell1Bytes, miniTree.root).tryGet()
# cell2Proof.verifyDataBlock(cell2Bytes, miniTree.root).tryGet()
# test "Can gather proof input":
# # This is the main entry point for this module, and what it's all about.
# let
# nSamples = 3
# input = (await dataSampler.getProofInput(challenge, nSamples)).tryget()
# proc equal(a: FieldElement, b: FieldElement): bool =
# a.toDecimal() == b.toDecimal()
# proc toStr(proof: MerkleProof): string =
# toHex(proof.nodesBuffer)
# let
# expectedMerkleProofs = getExpectedSlotToBlockProofs()
# expectedCellData = getExpectedCellData()
# check:
# # datasetRoot*: FieldElement
# equal(input.datasetRoot, datasetRootHash)
# # entropy*: FieldElement
# equal(input.entropy, challenge)
# # numberOfCellsInSlot*: uint64
# input.numberOfCellsInSlot == (bytesPerBlock * numberOfSlotBlocks).uint64 div CellSize
# # numberOfSlots*: uint64
# input.numberOfSlots == slot.request.ask.slots
# # datasetSlotIndex*: uint64
# input.datasetSlotIndex == slot.slotIndex.truncate(uint64)
# # slotRoot*: FieldElement
# equal(input.slotRoot, toF(1234)) # TODO - when slotPoseidonTree is a poseidon tree, its root should be a FieldElement.
# # datasetToSlotProof*: MerkleProof
# input.datasetToSlotProof == datasetToSlotProof
# # proofSamples*: seq[ProofSample]
# toStr(input.proofSamples[0].merkleProof) == expectedMerkleProofs[0]
# toStr(input.proofSamples[1].merkleProof) == expectedMerkleProofs[1]
# toStr(input.proofSamples[2].merkleProof) == expectedMerkleProofs[2]
# # cell data
# toHex(input.proofSamples[0].cellData) == expectedCellData[0]
# toHex(input.proofSamples[1].cellData) == expectedCellData[1]
# toHex(input.proofSamples[2].cellData) == expectedCellData[2]
# # input.slotToBlockProofs.mapIt(toStr(it)) == expectedSlotToBlockProofs
# # input.blockToCellProofs.mapIt(toStr(it)) == expectedBlockToCellProofs
# # toHex(input.sampleData) == expectedSampleData
# for (input, expected) in [(10, 0), (31, 0), (32, 1), (63, 1), (64, 2)]:
# test "Can get slotBlockIndex from slotCellIndex (" & $input & " -> " & $expected & ")":
# let
# slotCellIndex = input.uint64
# slotBlockIndex = dataSampler.getSlotBlockIndexForSlotCellIndex(slotCellIndex)
# check:
# slotBlockIndex == expected.uint64
# 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)
# check:
# blockCellIndex == expected.uint64