From 6e9fcc632c9599679bca8788c9c4221eabdb5e21 Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Thu, 11 Jul 2024 15:47:23 +0200
Subject: [PATCH] fix the incorrect handling of singleton Merkle trees

---
 circuit/codex/merkle.circom       | 13 ++++++++++++-
 circuit/codex/sample_cells.circom | 23 ++++++++++++++++++-----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/circuit/codex/merkle.circom b/circuit/codex/merkle.circom
index 6af6587..9e00bea 100644
--- a/circuit/codex/merkle.circom
+++ b/circuit/codex/merkle.circom
@@ -50,6 +50,17 @@ template RootFromMerklePath( maxDepth ) {
   signal input  merklePath[ maxDepth ];
   signal output recRoot;
 
+  // in case of a singleton tree, we receive maskBits = [0,0,0,...,0]
+  // but what we really need is [1,0,0,0,...,0]
+  // maybe it's the best to fix that here
+  //
+  // this is a bit of hackish, but because we always expect [1,1,...,1,0,0,...,0] anyway,
+  // we can just set the first entry to 1 and that should fix this issue.
+  //
+  signal maskBitsCorrected[ maxDepth + 1];
+  maskBitsCorrected[0] <== 1;
+  for(var i=1; i<=maxDepth; i++) { maskBitsCorrected[i] <== maskBits[i]; }
+
   // the sequence of reconstructed hashes along the path
   signal aux[ maxDepth+1 ];
   aux[0] <== leaf;
@@ -96,7 +107,7 @@ template RootFromMerklePath( maxDepth ) {
   var sum = 0;
   signal prods[maxDepth];
   for(var i=0; i<maxDepth; i++) {
-    prods[i] <== (maskBits[i] - maskBits[i+1]) * aux[i+1];
+    prods[i] <== (maskBitsCorrected[i] - maskBitsCorrected[i+1]) * aux[i+1];
     sum += prods[i];
   }
   recRoot <== sum;
diff --git a/circuit/codex/sample_cells.circom b/circuit/codex/sample_cells.circom
index fbaa6f4..cb03e92 100644
--- a/circuit/codex/sample_cells.circom
+++ b/circuit/codex/sample_cells.circom
@@ -74,6 +74,19 @@ template SampleAndProve( maxDepth, maxLog2NSlots, blockTreeDepth, nFieldElemsPer
   signal input cellData[nSamples][nFieldElemsPerCell];   // private input
   signal input merklePaths[nSamples][maxDepth];          // private input
 
+  // -------------------------------------------------------
+  // sanity check for singleton trees
+  // we don't have to do anything here, it's fixed in `merkle.circom`
+
+  // component eqsing = IsEqual();
+  // eqsing.A <== nCellsPerSlot;
+  // eqsing.B <== (1<<blockTreeDepth);     // normally 32
+  // signal slotTreeIsSingleton <== eqsing.out;
+
+log("block tree (bottom) is singleton = ", blockTreeDepth == 0 );
+log("slot tree (middle)  is singleton = ", nCellsPerSlot == (1<<blockTreeDepth) );
+log("dataset tree (top)  is singleton = ", nSlotsPerDataSet == 1 );
+
   // -------------------------------------------------------
   //
   // first we prove the inclusion of the slot root in the dataset-level
@@ -125,11 +138,11 @@ log("top root check = ", mtop.recRoot == dataSetRoot);
     calci[cnt].counter          <== cnt + 1;
     calci[cnt].indexBits        ==> prove[cnt].indexBits;
 
-    prove[cnt].slotRoot   <== slotRoot;
-    prove[cnt].lastBits   <== lastBits;
-    prove[cnt].maskBits   <== lg.mask;
-    prove[cnt].data       <== cellData[cnt];
-    prove[cnt].merklePath <== merklePaths[cnt];
+    prove[cnt].slotRoot     <== slotRoot;
+    prove[cnt].lastBits     <== lastBits;
+    prove[cnt].maskBits     <== lg.mask;
+    prove[cnt].data         <== cellData[cnt];
+    prove[cnt].merklePath   <== merklePaths[cnt];
 
   }
 }