From 24be5b2788d47eb813476143ce1cd6eef83668cb Mon Sep 17 00:00:00 2001
From: benbierens <thatbenbierens@gmail.com>
Date: Mon, 27 Nov 2023 14:45:24 +0100
Subject: [PATCH] setting up

---
 codex/proof/datasampler.nim           | 18 +++++++++++---
 tests/codex/proof/testdatasampler.nim | 34 ++++++++++++++++++++++-----
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/codex/proof/datasampler.nim b/codex/proof/datasampler.nim
index 15ea6f6d..97b823dd 100644
--- a/codex/proof/datasampler.nim
+++ b/codex/proof/datasampler.nim
@@ -35,10 +35,12 @@ type
     slotBlocks: SlotBlocks
     # The following data is invariant over time for a given slot:
     datasetRoot: FieldElement
+    slotRootHash: FieldElement
     slotPoseidonTree: MerkleTree
     datasetToSlotProof: MerkleProof
     blockSize: uint64
     numberOfCellsInSlot: uint64
+    datasetSlotIndex: uint64
     numberOfCellsPerBlock: uint64
 
 proc getNumberOfCellsInSlot*(slot: Slot): uint64 =
@@ -67,10 +69,12 @@ proc new*(
     blockStore: blockStore,
     slotBlocks: slotBlocks,
     datasetRoot: datasetRoot,
+    slotRootHash: toF(1234), # TODO - when slotPoseidonTree is a poseidon tree, its root should be a FieldElement.
     slotPoseidonTree: slotPoseidonTree,
     datasetToSlotProof: datasetToSlotProof,
     blockSize: blockSize,
     numberOfCellsInSlot: numberOfCellsInSlot,
+    datasetSlotIndex: slot.slotIndex.truncate(uint64),
     numberOfCellsPerBlock: blockSize div CellSize
   ))
 
@@ -175,8 +179,16 @@ proc getProofInput*(self: DataSampler, challenge: FieldElement, nSamples: int):
 
   trace "Successfully collected proof input"
   success(ProofInput(
+    datasetRoot: self.datasetRoot,
+    entropy: challenge,
+    numberOfCellsInSlot: self.numberOfCellsInSlot,
+    numberOfSlots: self.slot.request.ask.slots,
+    datasetSlotIndex: self.datasetSlotIndex,
+    slotRoot: self.slotRootHash,
     datasetToSlotProof: self.datasetToSlotProof,
-    slotToBlockProofs: slotToBlockProofs,
-    blockToCellProofs: blockToCellProofs,
-    sampleData: sampleData
+    #proofSamples: yeah!
+
+    # slotToBlockProofs: slotToBlockProofs,
+    # blockToCellProofs: blockToCellProofs,
+    # sampleData: sampleData
   ))
diff --git a/tests/codex/proof/testdatasampler.nim b/tests/codex/proof/testdatasampler.nim
index b4ea71ed..1d25a05a 100644
--- a/tests/codex/proof/testdatasampler.nim
+++ b/tests/codex/proof/testdatasampler.nim
@@ -4,6 +4,7 @@ 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
@@ -29,7 +30,7 @@ import testdatasampler_expected
 let
   bytesPerBlock = 64 * 1024
   challenge: FieldElement = toF(12345)
-  slotRootHash: FieldElement = toF(6789)
+  datasetRootHash: FieldElement = toF(6789)
 
 asyncchecksuite "Test proof datasampler - components":
   let
@@ -168,7 +169,7 @@ asyncchecksuite "Test proof datasampler - main":
     dataSampler = (await DataSampler.new(
       slot,
       localStore,
-      slotRootHash,
+      datasetRootHash,
       slotPoseidonTree,
       datasetToSlotProof
     )).tryGet()
@@ -202,7 +203,7 @@ asyncchecksuite "Test proof datasampler - main":
     proc getExpectedIndex(i: int): uint64 =
       let
         numberOfCellsInSlot = (bytesPerBlock * numberOfSlotBlocks) div CellSize.int
-        hash = Sponge.digest(@[slotRootHash, challenge, toF(i)], rate = 2)
+        hash = Sponge.digest(@[datasetRootHash, challenge, toF(i)], rate = 2)
       return extractLowBits(hash.toBig(), ceilingLog2(numberOfCellsInSlot))
 
     check:
@@ -278,11 +279,32 @@ asyncchecksuite "Test proof datasampler - main":
       expectedBlockToCellProofs = getExpectedBlockToCellProofs()
       expectedSampleData = getExpectedSampleData()
 
+    proc equal(a: FieldElement, b: FieldElement): bool =
+      a.toDecimal() == b.toDecimal()
+
     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
-      input.slotToBlockProofs.mapIt(toStr(it)) == expectedSlotToBlockProofs
-      input.blockToCellProofs.mapIt(toStr(it)) == expectedBlockToCellProofs
-      toHex(input.sampleData) == expectedSampleData
+      # proofSamples*: seq[ProofSample]
+      # yeah
+
+
+
+      # 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 & ")":