From a7ad30f49a7031bcc383cdaaa8b3003ef6bef434 Mon Sep 17 00:00:00 2001
From: Csaba Kiraly <csaba.kiraly@gmail.com>
Date: Thu, 16 Nov 2023 11:52:02 +0100
Subject: [PATCH] encoding: adding mult-dimensional encoding API

Signed-off-by: Csaba Kiraly <csaba.kiraly@gmail.com>
---
 codex/erasure/erasure.nim   | 39 +++++++++++++++++++++++++++++++++++++
 codex/manifest/manifest.nim |  5 +++++
 tests/codex/testerasure.nim | 33 +++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+)

diff --git a/codex/erasure/erasure.nim b/codex/erasure/erasure.nim
index 6996353c..2803a9c8 100644
--- a/codex/erasure/erasure.nim
+++ b/codex/erasure/erasure.nim
@@ -385,6 +385,25 @@ proc encode*(
 
   return success encodedManifest
 
+proc encodeMulti*(
+  self: Erasure,
+  manifest: Manifest,
+  code: seq[(int,int)]): Future[?!Manifest] {.async.} =
+  ## Encode a manifest into one that is erasure protected in multiple dimensions.
+  ##
+  ## code: list of (K,M) RS code parameter pairs
+
+  var
+    mfest = manifest
+    i = 1
+
+  for (k,m) in code:
+    mfest = (await self.encode(mfest, k, m, i)).tryGet()
+    i *= (k+m)
+
+  return success mfest
+
+
 proc decode*(
     self: Erasure,
     encoded: Manifest
@@ -495,6 +514,26 @@ proc decode*(
 
   return decoded.success
 
+proc decodeMulti*(
+    self: Erasure,
+    encoded: Manifest
+): Future[?!Manifest] {.async.} =
+  ## Decode a protected manifest into it's original unprotected
+  ## manifest
+  ##
+  ## `encoded` - the encoded (protected) manifest to
+  ##             be recovered
+  ##
+
+  var
+    mfest = encoded
+
+  while mfest.protected:
+    mfest = (await self.decode(mfest)).tryGet()
+
+  return success mfest
+
+
 proc start*(self: Erasure) {.async.} =
   return
 
diff --git a/codex/manifest/manifest.nim b/codex/manifest/manifest.nim
index 8499ff17..270ed0a6 100644
--- a/codex/manifest/manifest.nim
+++ b/codex/manifest/manifest.nim
@@ -85,6 +85,11 @@ proc originalTreeCid*(self: Manifest): Cid =
 proc originalBlocksCount*(self: Manifest): int =
   divUp(self.originalManifest.datasetSize.int, self.blockSize.int)
 
+proc unprotectedBlocksCount*(self: Manifest): int =
+  var mfest = self
+  while mfest.protected:
+    mfest = mfest.originalManifest
+  divUp(mfest.datasetSize.int, self.blockSize.int)
 
 proc originalDatasetSize*(self: Manifest): NBytes =
   self.originalDatasetSize
diff --git a/tests/codex/testerasure.nim b/tests/codex/testerasure.nim
index 1852e0e9..25a8cf35 100644
--- a/tests/codex/testerasure.nim
+++ b/tests/codex/testerasure.nim
@@ -356,3 +356,36 @@ asyncchecksuite "Erasure encode/decode":
     for d in dropped:
       let present = await store.hasBlock(manifest.treeCid, d)
       check present.tryGet()
+
+  test "3D encode: test multi-dimensional API":
+    const
+      encoding = @[(7, 3),(5, 2),(3, 1)]
+
+    let
+      encoded = (await erasure.encodeMulti(manifest, encoding)).tryGet()
+      decoded = (await erasure.decodeMulti(encoded)).tryGet()
+
+    check:
+      decoded.treeCid == manifest.treeCid
+      decoded.blocksCount == encoded.unprotectedBlocksCount
+
+  test "3D encode: test multi-dimensional API with drop":
+    const
+      encoding = @[(7, 3),(5, 2),(3, 1)]
+
+    let encoded = (await erasure.encodeMulti(manifest, encoding)).tryGet()
+
+    var
+      idx = rng.rand(encoded.steps - 1) # random row
+      dropped: seq[int]
+
+    for _ in 0..<encoded.ecM:
+      dropped.add(idx)
+      (await store.delBlock(encoded.treeCid, idx)).tryGet()
+      idx.inc(encoded.interleave)
+
+    let decoded = (await erasure.decodeMulti(encoded)).tryGet()
+
+    check:
+      decoded.treeCid == manifest.treeCid
+      decoded.blocksCount == encoded.unprotectedBlocksCount