bug fix

EspressoSystems · Feb 27, 2024 · 66d9481 · 66d9481
1 parent 41fe9e5
commit 66d9481
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 6 deletions.
diff --git a/primitives/src/merkle_tree/append_only.rs b/primitives/src/merkle_tree/append_only.rs
@@ -167,6 +167,11 @@ mod mt_tests {
     }
 
     fn test_mt_lookup_helper<F: RescueParameter>() {
+        // singleton merkle tree test (#499)
+        let mt = RescueMerkleTree::<F>::from_elems(None, [F::from(0u64)]).unwrap();
+        let (elem, _) = mt.lookup(0).expect_ok().unwrap();
+        assert_eq!(elem, &F::from(0u64));
+
         let mt =
             RescueMerkleTree::<F>::from_elems(Some(2), [F::from(3u64), F::from(1u64)]).unwrap();
         let root = mt.commitment().digest();

diff --git a/primitives/src/merkle_tree/internal.rs b/primitives/src/merkle_tree/internal.rs
@@ -200,7 +200,19 @@ where
         Err(PrimitivesError::ParameterError(
             "Too many data for merkle tree".to_string(),
         ))
-    } else if num_leaves > 0 {
+    } else if num_leaves == 0 {
+        Ok((Arc::new(MerkleNode::<E, u64, T>::Empty), height, 0))
+    } else if height == 0usize {
+        Ok((
+            Arc::new(MerkleNode::Leaf {
+                value: H::digest_leaf(&0, leaves[0].borrow())?,
+                pos: 0,
+                elem: leaves[0].borrow().clone(),
+            }),
+            height,
+            1,
+        ))
+    } else {
         let mut cur_nodes = leaves
             .into_iter()
             .enumerate()
@@ -243,8 +255,6 @@ where
                 .collect::<Result<Vec<_>, PrimitivesError>>()?;
         }
         Ok((cur_nodes[0].clone(), height, num_leaves))
-    } else {
-        Ok((Arc::new(MerkleNode::<E, u64, T>::Empty), height, 0))
     }
 }
 
@@ -278,7 +288,19 @@ where
         Err(PrimitivesError::ParameterError(
             "Too many data for merkle tree".to_string(),
         ))
-    } else if num_leaves > 0 {
+    } else if num_leaves == 0 {
+        Ok((Arc::new(MerkleNode::<E, u64, T>::Empty), height, 0))
+    } else if height == 0usize {
+        Ok((
+            Arc::new(MerkleNode::Leaf {
+                value: H::digest_leaf(&0, leaves[0].borrow())?,
+                pos: 0,
+                elem: leaves[0].borrow().clone(),
+            }),
+            height,
+            1,
+        ))
+    } else {
         let mut cur_nodes = leaves
             .into_iter()
             .enumerate()
@@ -337,8 +359,6 @@ where
             }
         }
         Ok((cur_nodes[0].clone(), height, num_leaves))
-    } else {
-        Ok((Arc::new(MerkleNode::<E, u64, T>::Empty), height, 0))
     }
 }
 

diff --git a/primitives/src/merkle_tree/light_weight.rs b/primitives/src/merkle_tree/light_weight.rs
@@ -174,6 +174,11 @@ mod mt_tests {
     }
 
     fn test_light_mt_lookup_helper<F: RescueParameter>() {
+        // singleton merkle tree test (#499)
+        let mt = RescueLightWeightMerkleTree::<F>::from_elems(None, [F::from(0u64)]).unwrap();
+        let (elem, _) = mt.lookup(0).expect_ok().unwrap();
+        assert_eq!(elem, &F::from(0u64));
+
         let mut mt =
             RescueLightWeightMerkleTree::<F>::from_elems(Some(2), [F::from(3u64), F::from(1u64)])
                 .unwrap();