Sets up getting sample from block

codex/proof/datasampler.nim

@@ -1,4 +1,5 @@
 import ../contracts/requests
+import ../blocktype as bt
 
 import std/bitops
 import std/sugar
@@ -9,14 +10,16 @@ import pkg/poseidon2
 
 import misc
 
-type
-  DSFieldElement* = F
-  DSCellIndex* = uint64
-
 const
   # Size of a cell.
   # A cell is a sample of storage-data selected for proving.
-  CellSize* = u256(2048)
+  CellSize* = 2048.uint64
+
+type
+  DSFieldElement* = F
+  DSCellIndex* = uint64
+  DSSample* = array[CellSize, byte]
+
 
 func extractLowBits*[n: static int](A: BigInt[n], k: int): uint64 =
   assert(k > 0 and k <= 64)
@@ -37,19 +40,19 @@ proc getCellIndex(fe: DSFieldElement, numberOfCells: int): uint64 =
 
   return extractLowBits(fe.toBig(), log2)
 
-proc getNumberOfCellsInSlot*(slot: Slot): int =
-  (slot.request.ask.slotSize div CellSize).truncate(int)
+proc getNumberOfCellsInSlot*(slot: Slot): uint64 =
+  (slot.request.ask.slotSize.truncate(uint64) div CellSize)
 
 proc findCellIndex*(
   slotRootHash: DSFieldElement,
   challenge: DSFieldElement,
   counter: DSFieldElement,
-  numberOfCells: int): DSCellIndex =
+  numberOfCells: uint64): DSCellIndex =
   # Computes the cell index for a single sample.
   let
     input = @[slotRootHash, challenge, counter]
     hash = Sponge.digest(input, rate = 2)
-    index = getCellIndex(hash, numberOfCells)
+    index = getCellIndex(hash, numberOfCells.int)
 
   return index
 
@@ -61,3 +64,14 @@ func findCellIndices*(
   # Computes nSamples cell indices.
   let numberOfCells = getNumberOfCellsInSlot(slot)
   return collect(newSeq, (for i in 1..nSamples: findCellIndex(slotRootHash, challenge, toF(i), numberOfCells)))
+
+proc getSlotBlockIndex*(cellIndex: DSCellIndex, blockSize: uint64): uint64 =
+  let numberOfCellsPerBlock = blockSize div CellSize
+  return cellIndex div numberOfCellsPerBlock
+
+proc getCellIndexInBlock*(cellIndex: DSCellIndex, blockSize: uint64): uint64 =
+  let numberOfCellsPerBlock = blockSize div CellSize
+  return cellIndex mod numberOfCellsPerBlock
+
+proc getSampleFromBlock*(blk: bt.Block, cellIndex: DSCellIndex, blockSize: uint64): DSSample =
+  raiseAssert("a")

codex/proof/slotblocks.nim

@@ -37,11 +37,7 @@ proc getIndexForSlotBlock*(slot: Slot, blockSize: NBytes, slotBlockIndex: uint64
   trace "Converted slotBlockIndex to datasetIndex", slotBlockIndex, datasetIndex
   return datasetIndex
 
-proc getSlotBlock*(slot: Slot, blockstore: BlockStore, slotBlockIndex: uint64): Future[?!Block] {.async.} =
-  without manifest =? (await getManifestForSlot(slot, blockstore)), err:
-    error "Failed to get manifest for slot"
-    return failure(err)
-
+proc getSlotBlock*(slot: Slot, blockstore: BlockStore, manifest: Manifest, slotBlockIndex: uint64): Future[?!Block] {.async.} =
   let
     blocksInManifest = (manifest.datasetSize div manifest.blockSize).uint64
     datasetIndex = getIndexForSlotBlock(slot, manifest.blockSize, slotBlockIndex)
@@ -50,3 +46,10 @@ proc getSlotBlock*(slot: Slot, blockstore: BlockStore, slotBlockIndex: uint64):
     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:
+    error "Failed to get manifest for slot"
+    return failure(err)
+
+  return await getSlotBlock(slot, blockstore, manifest, slotBlockIndex)

tests/codex/proof/testdatasampler.nim

@@ -2,6 +2,7 @@ import std/os
 import std/strutils
 import std/sequtils
 import std/sugar
+import std/random
 
 import pkg/questionable
 import pkg/questionable/results
@@ -80,7 +81,7 @@ asyncchecksuite "Test proof datasampler":
       let log2 = ceilingLog2(value)
       return (1 shl log2) == value
 
-    let numberOfCells = getNumberOfCellsInSlot(slot)
+    let numberOfCells = getNumberOfCellsInSlot(slot).int
 
     check:
       isPow2(numberOfCells)
@@ -97,11 +98,15 @@ asyncchecksuite "Test proof datasampler":
       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
 
   test "Should calculate total number of cells in Slot":
     let
-      slotSizeInBytes = slot.request.ask.slotSize
-      expectedNumberOfCells = (slotSizeInBytes div CellSize).truncate(int)
+      slotSizeInBytes = (slot.request.ask.slotSize).truncate(uint64)
+      expectedNumberOfCells = slotSizeInBytes div CellSize
 
     check:
       expectedNumberOfCells == 512
@@ -118,7 +123,7 @@ asyncchecksuite "Test proof datasampler":
 
     proc getExpectedIndex(i: int): DSCellIndex =
       let hash = Sponge.digest(@[slotRootHash, challenge, toF(i)], rate = 2)
-      return extractLowBits(hash.toBig(), ceilingLog2(numberOfCells))
+      return extractLowBits(hash.toBig(), ceilingLog2(numberOfCells.int))
 
     check:
       cellIndex(1) == getExpectedIndex(1)
@@ -139,3 +144,45 @@ asyncchecksuite "Test proof datasampler":
     check:
       cellIndices(3) == getExpectedIndices(3)
       cellIndices(3) == knownIndices
+
+  for (input, expected) in [(10, 0), (31, 0), (32, 1), (63, 1), (64, 2)]:
+    test "Can get slotBlockIndex from cell index (" & $input & " -> " & $expected & ")":
+      let
+        cellIndex = input.uint64
+        blockSize = (64 * 1024).uint64
+
+        slotBlockIndex = getSlotBlockIndex(cellIndex, blockSize)
+
+      check:
+        slotBlockIndex == expected.uint64
+
+  for (input, expected) in [(10, 10), (31, 31), (32, 0), (63, 31), (64, 0)]:
+    test "Can get cellIndexInBlock from cell index (" & $input & " -> " & $expected & ")":
+      let
+        cellIndex = input.uint64
+        blockSize = (64 * 1024).uint64
+
+        cellIndexInBlock = getCellIndexInBlock(cellIndex, blockSize)
+
+      check:
+        cellIndexInBlock == expected.uint64
+
+  test "Can get sample from block":
+    let
+      blockSize = CellSize * 2
+      bytes = newSeqWith(blockSize.int, rand(uint8))
+      blk = bt.Block.new(bytes).tryGet()
+
+      sample0 = getSampleFromBlock(blk, 0, blockSize.uint64)
+      sample1 = getSampleFromBlock(blk, 1, blockSize.uint64)
+
+    check:
+      sample0 == bytes[0..<CellSize]
+      sample1 == bytes[CellSize..^1]
+
+    # proc getSampleFromBlock(blk: bt.Block, cellIndex: DSCellIndex, blockSize: uint64): DSSample
+
+
+    # let length = rand(4096)
+  # let bytes = newSeqWith(length, rand(uint8))
+  # bt.Block.new(bytes).tryGet()