From 4642174c6d23054f31beea87f180a19956577743 Mon Sep 17 00:00:00 2001
From: benbierens <thatbenbierens@gmail.com>
Date: Thu, 23 Nov 2023 11:51:26 +0100
Subject: [PATCH] Adds method to get dataset block index from slot block index

---
 codex/proof/datasampler.nim           | 68 ++++++++++++++++++++++++++-
 tests/codex/proof/testdatasampler.nim | 23 +++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/codex/proof/datasampler.nim b/codex/proof/datasampler.nim
index 34fa4110..40c6e5ed 100644
--- a/codex/proof/datasampler.nim
+++ b/codex/proof/datasampler.nim
@@ -1,11 +1,14 @@
 import ../contracts/requests
 import ../blocktype as bt
 import ../merkletree
+import ../manifest
+import ../stores/blockstore
 
 import std/bitops
 import std/sugar
 
 import pkg/chronicles
+import pkg/chronos
 import pkg/questionable
 import pkg/questionable/results
 import pkg/constantine/math/arithmetic
@@ -13,6 +16,7 @@ import pkg/poseidon2/types
 import pkg/poseidon2
 
 import misc
+import slotblocks
 
 const
   # Size of a cell.
@@ -23,7 +27,10 @@ type
   DSFieldElement* = F
   DSCellIndex* = uint64
   DSCell* = seq[byte]
-
+  ProofInput* = ref object
+    blockInclProofs*: seq[MerkleProof]
+    cellInclProofs*: seq[MerkleProof]
+    sampleData*: seq[byte]
 
 func extractLowBits*[n: static int](A: BigInt[n], k: int): uint64 =
   assert(k > 0 and k <= 64)
@@ -73,6 +80,14 @@ proc getSlotBlockIndex*(cellIndex: DSCellIndex, blockSize: uint64): uint64 =
   let numberOfCellsPerBlock = blockSize div CellSize
   return cellIndex div numberOfCellsPerBlock
 
+proc getDatasetBlockIndex*(slot: Slot, slotBlockIndex: uint64, blockSize: uint64): uint64 =
+  let
+    slotIndex = slot.slotIndex.truncate(uint64)
+    slotSize = slot.request.ask.slotSize.truncate(uint64)
+    blocksInSlot = slotSize div blockSize
+
+  return (blocksInSlot * slotIndex) + slotBlockIndex
+
 proc getCellIndexInBlock*(cellIndex: DSCellIndex, blockSize: uint64): uint64 =
   let numberOfCellsPerBlock = blockSize div CellSize
   return cellIndex mod numberOfCellsPerBlock
@@ -104,3 +119,54 @@ proc getBlockCellMiniTree*(blk: bt.Block, blockSize: uint64): ?!MerkleTree =
       return failure("Failed to add cell data to tree")
 
   return builder.build()
+
+proc getProofInput*(
+  slot: Slot,
+  blockStore: BlockStore,
+  slotRootHash: DSFieldElement,
+  dataSetPoseidonTree: MerkleTree,
+  challenge: DSFieldElement,
+  nSamples: int
+): Future[?!ProofInput] {.async.} =
+  var
+    blockProofs: seq[MerkleProof]
+    cellProofs: seq[MerkleProof]
+    sampleData: seq[byte]
+
+  without manifest =? await getManifestForSlot(slot, blockStore), err:
+    error "Failed to get manifest for slot"
+    return failure(err)
+
+  let
+    blockSize = manifest.blockSize.uint64
+    cellIndices = findCellIndices(slot, slotRootHash, challenge, nSamples)
+
+  for cellIndex in cellIndices:
+    let slotBlockIndex = getSlotBlockIndex(cellIndex, blockSize)
+    without blk =? await getSlotBlock(slot, blockStore, manifest, slotBlockIndex), err:
+      error "Failed to get slot block"
+      return failure(err)
+
+    without miniTree =? getBlockCellMiniTree(blk, blockSize), err:
+      error "Failed to calculate minitree for block"
+      return failure(err)
+
+    # without blockProof =? dataSetPoseidonTree.getProof(???block index in dataset!), err:
+    #   error "Failed to get dataset inclusion proof"
+    #   return failure(err)
+    # blockProofs.add(blockProof)
+
+    without cellProof =? miniTree.getProof(cellIndex), err:
+      error "Failed to get cell inclusion proof"
+      return failure(err)
+    cellProofs.add(cellProof)
+
+    let cell = getCellFromBlock(blk, cellIndex, blockSize)
+    sampleData = sampleData & cell
+
+  trace "Successfully collected proof input data"
+  success(ProofInput(
+    blockInclProofs: blockProofs,
+    cellInclProofs: cellProofs,
+    sampleData: sampleData
+  ))
diff --git a/tests/codex/proof/testdatasampler.nim b/tests/codex/proof/testdatasampler.nim
index fba6e2aa..a2add9fa 100644
--- a/tests/codex/proof/testdatasampler.nim
+++ b/tests/codex/proof/testdatasampler.nim
@@ -24,6 +24,7 @@ import pkg/codex/utils/asynciter
 import pkg/codex/contracts/requests
 import pkg/codex/contracts
 import pkg/codex/merkletree
+import pkg/codex/stores/cachestore
 
 import pkg/codex/proof/datasampler
 import pkg/codex/proof/misc
@@ -157,6 +158,17 @@ asyncchecksuite "Test proof datasampler":
       check:
         slotBlockIndex == expected.uint64
 
+  for input in 0 ..< numberOfSlotBlocks:
+    test "Can get datasetBlockIndex from slotBlockIndex (" & $input & ")":
+      let
+        slotBlockIndex = input.uint64
+        datasetBlockIndex = getDatasetBlockIndex(slot, slotBlockIndex, bytesPerBlock.uint64)
+        slotIndex = slot.slotIndex.truncate(uint64)
+        expectedIndex = (numberOfSlotBlocks.uint64 * slotIndex) + slotBlockIndex
+
+      check:
+        datasetBlockIndex == expectedIndex
+
   for (input, expected) in [(10, 10), (31, 31), (32, 0), (63, 31), (64, 0)]:
     test "Can get cellIndexInBlock from cell index (" & $input & " -> " & $expected & ")":
       let
@@ -217,3 +229,14 @@ asyncchecksuite "Test proof datasampler":
       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
+  #     localStore = CacheStore.new()
+  #     dataSetPoseidonTree = MerkleTree.init(@[Cid.example])
+  #     a = (await getProofInput(slot, localStore, slotRootHash, dataSetPoseidonTree, challenge, 3)).tryget()
+
+  #   echo "a.blockInclProofs: " & $a.blockInclProofs.len
+  #   echo "a.cellInclProofs: " & $a.cellInclProofs.len
+  #   echo "a.sampleData: " & $a.sampleData.len