From 07ed066b1eec78d6b9b1f70b764384dafa7f0221 Mon Sep 17 00:00:00 2001
From: benbierens <thatbenbierens@gmail.com>
Date: Wed, 22 Nov 2023 16:28:05 +0100
Subject: [PATCH] Implements selecting a cell sample from a block

---
 codex/proof/datasampler.nim           | 9 +++++++--
 tests/codex/proof/testdatasampler.nim | 6 ++++--
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/codex/proof/datasampler.nim b/codex/proof/datasampler.nim
index c317e183..51670dbd 100644
--- a/codex/proof/datasampler.nim
+++ b/codex/proof/datasampler.nim
@@ -18,7 +18,7 @@ const
 type
   DSFieldElement* = F
   DSCellIndex* = uint64
-  DSSample* = array[CellSize, byte]
+  DSSample* = seq[byte]
 
 
 func extractLowBits*[n: static int](A: BigInt[n], k: int): uint64 =
@@ -74,4 +74,9 @@ proc getCellIndexInBlock*(cellIndex: DSCellIndex, blockSize: uint64): uint64 =
   return cellIndex mod numberOfCellsPerBlock
 
 proc getSampleFromBlock*(blk: bt.Block, cellIndex: DSCellIndex, blockSize: uint64): DSSample =
-  raiseAssert("a")
+  let
+    inBlockCellIndex = getCellIndexInBlock(cellIndex, blockSize)
+    dataStart = (CellSize * inBlockCellIndex)
+    dataEnd = dataStart + CellSize
+
+  return blk.data[dataStart ..< dataEnd]
diff --git a/tests/codex/proof/testdatasampler.nim b/tests/codex/proof/testdatasampler.nim
index 813f4ba2..3320c205 100644
--- a/tests/codex/proof/testdatasampler.nim
+++ b/tests/codex/proof/testdatasampler.nim
@@ -169,16 +169,18 @@ asyncchecksuite "Test proof datasampler":
 
   test "Can get sample from block":
     let
-      blockSize = CellSize * 2
+      blockSize = CellSize * 3
       bytes = newSeqWith(blockSize.int, rand(uint8))
       blk = bt.Block.new(bytes).tryGet()
 
       sample0 = getSampleFromBlock(blk, 0, blockSize.uint64)
       sample1 = getSampleFromBlock(blk, 1, blockSize.uint64)
+      sample2 = getSampleFromBlock(blk, 2, blockSize.uint64)
 
     check:
       sample0 == bytes[0..<CellSize]
-      sample1 == bytes[CellSize..^1]
+      sample1 == bytes[CellSize..<(CellSize*2)]
+      sample2 == bytes[(CellSize*2)..^1]
 
     # proc getSampleFromBlock(blk: bt.Block, cellIndex: DSCellIndex, blockSize: uint64): DSSample