From c320dd2e7d493d7330b90454d8756210be906624 Mon Sep 17 00:00:00 2001
From: benbierens <thatbenbierens@gmail.com>
Date: Wed, 22 Nov 2023 09:19:41 +0100
Subject: [PATCH] Sets up tests for bitwise modulo

---
 codex/proof/datasampler.nim           | 57 ++++++++++++++++++++++++++-
 codex/proof/misc.nim                  | 13 ++++++
 tests/codex/proof/testdatasampler.nim | 44 ++++++++++++++++++---
 3 files changed, 107 insertions(+), 7 deletions(-)
 create mode 100644 codex/proof/misc.nim

diff --git a/codex/proof/datasampler.nim b/codex/proof/datasampler.nim
index c38739fc..249395d0 100644
--- a/codex/proof/datasampler.nim
+++ b/codex/proof/datasampler.nim
@@ -1,9 +1,62 @@
 import ../contracts/requests
 
+import std/bitops
+
+import pkg/constantine/math/arithmetic
+import pkg/poseidon2/types
+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)
 
-proc getNumberOfCellsInSlot*(slot: Slot): Uint256 =
-  slot.request.ask.slotSize div CellSize
+func extractLowBits*[n: static int](A: BigInt[n], k: int): uint64 =
+  assert(k > 0 and k <= 64)
+  var r: uint64 = 0
+  for i in 0..<k:
+    # A is big-endian. Run index backwards: n-1-i
+    #let b = bit[n](A, n-1-i)
+    let b = bit[n](A, i)
+
+    let y = uint64(b)
+    if (y != 0):
+      r = bitor(r, 1'u64 shl i)
+  return r
+
+proc getCellIndex(fe: DSFieldElement, numberOfCells: int): uint64 =
+  let log2 = ceilingLog2(numberOfCells)
+  assert((1 shl log2) == numberOfCells , "expected `numberOfCells` to be a power of two.")
+
+  return extractLowBits(fe.toBig(), log2)
+
+proc getNumberOfCellsInSlot*(slot: Slot): int =
+  (slot.request.ask.slotSize div CellSize).truncate(int)
+
+proc findCellIndex*(
+  slotRootHash: DSFieldElement,
+  challenge: DSFieldElement,
+  counter: DSFieldElement,
+  numberOfCells: int): DSCellIndex =
+  # Computes the cell index for a single sample.
+  let
+    input = @[slotRootHash, challenge, counter]
+    hash = Sponge.digest(input, rate = 2)
+    index = getCellIndex(hash, numberOfCells)
+
+  return index
+
+# func findCellIndices*(
+#   slot: Slot,
+#   slotRootHash: DSFieldElement,
+#   challenge: DSFieldElement,
+#   nSamples: int): seq[int] =
+#   # Computes nSamples cell indices.
+#   let numberOfCells = getNumberOfCellsInSlot(slot)
+#   return collect(newSeq, (for i in 1..nSamples: findCellIndex(slotRootHash, challenge, toF(i), numberOfCells)))
diff --git a/codex/proof/misc.nim b/codex/proof/misc.nim
new file mode 100644
index 00000000..58fa25b1
--- /dev/null
+++ b/codex/proof/misc.nim
@@ -0,0 +1,13 @@
+func floorLog2* (x : int) : int =
+  var k = -1
+  var y = x
+  while (y > 0):
+    k += 1
+    y = y shr 1
+  return k
+
+func ceilingLog2* (x : int) : int =
+  if (x==0):
+    return -1
+  else:
+    return (floorLog2(x-1) + 1)
diff --git a/tests/codex/proof/testdatasampler.nim b/tests/codex/proof/testdatasampler.nim
index 97b170e5..93652621 100644
--- a/tests/codex/proof/testdatasampler.nim
+++ b/tests/codex/proof/testdatasampler.nim
@@ -4,13 +4,14 @@ import std/sequtils
 
 import pkg/questionable
 import pkg/questionable/results
-
+import pkg/constantine/math/arithmetic
+import pkg/poseidon2/types
+import pkg/poseidon2
 import pkg/chronos
 import pkg/asynctest
 import pkg/stew/byteutils
 import pkg/stew/endians2
 import pkg/datastore
-
 import pkg/codex/rng
 import pkg/codex/stores/cachestore
 import pkg/codex/chunker
@@ -22,13 +23,16 @@ import pkg/codex/contracts/requests
 import pkg/codex/contracts
 
 import pkg/codex/proof/datasampler
+import pkg/codex/proof/misc
 
 import ../helpers
 import ../examples
 
 let
   bytesPerBlock = 64 * 1024
-  numberOfSlotBlocks = 10
+  numberOfSlotBlocks = 16
+  challenge: DSFieldElement = toF(12345)
+  slotRootHash: DSFieldElement = toF(6789)
   slot = Slot(
     request: StorageRequest(
       client: Address.example,
@@ -69,13 +73,43 @@ asyncchecksuite "Test proof datasampler":
   setup:
     await createSlotBlocks()
 
+  test "number of cells is a power of two":
+    proc isPow2(value: int): bool =
+      let log2 = ceilingLog2(value)
+      return (1 shl log2) == value
+
+    let numberOfCells = getNumberOfCellsInSlot(slot)
+
+    check:
+      isPow2(numberOfCells)
+
+  test "Extract low bits":
+    proc extract(value: int, 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
+
   test "Should calculate total number of cells in Slot":
     let
       slotSizeInBytes = slot.request.ask.slotSize
-      expectedNumberOfCells = slotSizeInBytes div CellSize
+      expectedNumberOfCells = (slotSizeInBytes div CellSize).truncate(int)
 
     check:
-      expectedNumberOfCells == 320
+      expectedNumberOfCells == 512
       expectedNumberOfCells == getNumberOfCellsInSlot(slot)
 
+  test "Can find single cell index":
+    let
+      counter: DSFieldElement = toF(1)
+      numberOfCells = getNumberOfCellsInSlot(slot)
+      cellIndex = findCellIndex(slotRootHash, challenge, counter, numberOfCells)
 
+    check:
+      cellIndex == 2