feat: Namespace merkle tree (#121)

* Implement namespace merkle tree and tests

* Clean up todos

* Improve Namespace import path and add Default::default for Namespace type

* Scarb fmt

* Rebase and scarb fmt

src/lib.cairo

@@ -28,9 +28,14 @@ mod tree {
         }
     }
     mod namespace {
+        mod hasher;
         mod merkle_tree;
+        mod namespace;
+        use namespace::Namespace;
         #[cfg(test)]
         mod tests {
+            mod test_hasher;
+            mod test_merkle_multi_proof;
             mod test_merkle_tree;
         }
     }

src/tree/consts.cairo

@@ -1,4 +1,4 @@
-use blobstream_sn::tree::namespace::merkle_tree::Namespace;
+use blobstream_sn::tree::namespace::Namespace;
 use core::bytes_31::bytes31_const;
 
 const MAX_HEIGHT: u256 = 256;

src/tree/namespace/hasher.cairo

@@ -1,9 +1,9 @@
-use alexandria_bytes::Bytes;
-use alexandria_bytes::BytesTrait;
+use alexandria_bytes::{Bytes, BytesTrait};
 use alexandria_math::U256BitShift;
 use blobstream_sn::tree::consts::{LEAF_PREFIX, NODE_PREFIX, parity_share_namespace};
+use blobstream_sn::tree::namespace::Namespace;
 use blobstream_sn::tree::namespace::merkle_tree::{
-    Namespace, NamespaceNode, NamespaceMerkleMultiproof, NamespaceMerkleProof
+    NamespaceNode, NamespaceMerkleMultiproof, NamespaceMerkleProof
 };
 
 fn leaf_digest(namespace: Namespace, data: @Bytes) -> NamespaceNode {
@@ -60,4 +60,3 @@ fn append_bytes28(ref self: Bytes, value: bytes31) {
     let mut bytes28Bytes: Bytes = BytesTrait::new(28, array![value_u256.high, value_u256.low]);
     self.concat(@bytes28Bytes);
 }
-

src/tree/namespace/merkle_tree.cairo

@@ -1,5 +1,5 @@
-// // Celestia-app namespace ID and its version
-// // See: https://celestiaorg.github.io/celestia-app/specs/namespace.html
+use blobstream_sn::tree::namespace::Namespace;
+
 #[derive(Serde, Drop, Copy, PartialEq)]
 struct NamespaceNode {
     min: Namespace,
@@ -8,12 +8,6 @@ struct NamespaceNode {
     digest: u256,
 }
 
-#[derive(Serde, Drop, Copy)]
-struct Namespace {
-    version: u8,
-    id: bytes31,
-}
-
 #[derive(Drop, PartialEq, PartialOrd)]
 struct NamespaceMerkleMultiproof {
     begin_key: u256,
@@ -28,86 +22,262 @@ struct NamespaceMerkleProof {
     num_leaves: u256,
 }
 
-// mod NamespaceMerkleTree {
-//     use super::{Namespace, NamespaceNode, NamespaceMerkleMultiproof, NamespaceMerkleProof};
-//     use alexandria_bytes::BytesTrait;
-//     use blobstream_sn::tree::binary::hasher::leafDigest;
-
-//     fn verify(
-//         root: NamespaceNode,
-//         proof: NamespaceMerkleProof,
-//         namespace: Namespace,
-//         data: alexandria_bytes::bytes::Bytes
-//     ) -> bool {
-
-//         node : NamespaceNode = leafDigest(namespace, data);
-//         true
-//     }
-// }
-
-impl NamespacePartialOrd of PartialOrd<Namespace> {
-    #[inline(always)]
-    fn le(lhs: Namespace, rhs: Namespace) -> bool {
-        let lhs_id: u256 = lhs.id.into();
-        let rhs_id: u256 = rhs.id.into();
-        if (lhs_id <= rhs_id && lhs.version <= rhs.version) {
-            return true;
-        } else {
+mod NamespaceMerkleTree {
+    use alexandria_bytes::bytes::Bytes;
+    use alexandria_math::U256BitShift;
+    use blobstream_sn::tree::namespace::hasher;
+    use blobstream_sn::tree::utils;
+    use super::{Namespace, NamespaceNode, NamespaceMerkleMultiproof, NamespaceMerkleProof};
+
+    fn verify(
+        root: NamespaceNode, proof: NamespaceMerkleProof, namespace: Namespace, data: Bytes
+    ) -> bool {
+        // Create a sibling leaf at height 1.
+        let node: NamespaceNode = hasher::leaf_digest(namespace, @data);
+
+        // Since we're verifying a leaf, height is 1.
+        return verify_inner(root, proof, node, 1);
+    }
+
+    fn verify_inner(
+        root: NamespaceNode,
+        proof: NamespaceMerkleProof,
+        mut node: NamespaceNode,
+        starting_height: u256
+    ) -> bool {
+        if starting_height < 1 {
             return false;
         }
-    }
-    #[inline(always)]
-    fn ge(lhs: Namespace, rhs: Namespace) -> bool {
-        let lhs_id: u256 = lhs.id.into();
-        let rhs_id: u256 = rhs.id.into();
-        if (lhs_id >= rhs_id && lhs.version >= rhs.version) {
-            return true;
-        } else {
+        let height_offset: u256 = starting_height - 1;
+
+        let proof_side_nodes_len: u256 = proof.side_nodes.len().into();
+
+        // Check proof is correct length for the key it is proving.
+        if proof.num_leaves <= 1 {
+            if proof_side_nodes_len != 0 {
+                return false;
+            }
+        } else if proof_side_nodes_len.into()
+            + height_offset != utils::path_length_from_key(proof.key, proof.num_leaves) {
             return false;
         }
-    }
-    #[inline(always)]
-    fn lt(lhs: Namespace, rhs: Namespace) -> bool {
-        let lhs_id: u256 = lhs.id.into();
-        let rhs_id: u256 = rhs.id.into();
-        if (lhs_id < rhs_id && lhs.version < rhs.version) {
-            return true;
-        } else {
+
+        // Check key is in tree
+        if proof.key >= proof.num_leaves {
             return false;
         }
-    }
-    #[inline(always)]
-    fn gt(lhs: Namespace, rhs: Namespace) -> bool {
-        let lhs_id: u256 = lhs.id.into();
-        let rhs_id: u256 = rhs.id.into();
-        if (lhs_id > rhs_id && lhs.version > rhs.version) {
-            return true;
-        } else {
+        // Handle case where proof is empty: i.e, only one leaf exists, so verify hash(data) is root
+        if proof_side_nodes_len == 0 {
+            if proof.num_leaves == 1 {
+                return super::namespace_node_eq(root, node);
+            } else {
+                return false;
+            }
+        }
+
+        let mut height: u256 = starting_height;
+        let mut stable_end: u256 = proof.key;
+
+        let mut exit_after: bool = false;
+        let side_nodes_span: Span<NamespaceNode> = proof.side_nodes.span();
+        while true {
+            let rounding_factor: u256 = U256BitShift::shl(1, height);
+            let sub_tree_start_index: u256 = (proof.key / rounding_factor) * rounding_factor;
+            let sub_tree_end_index: u256 = sub_tree_start_index + rounding_factor - 1;
+
+            if sub_tree_end_index >= proof.num_leaves {
+                break;
+            }
+            stable_end = sub_tree_end_index;
+
+            // Check if key is in the first or second half of the sub-tree.
+            if proof_side_nodes_len.into() + height_offset <= height - 1 {
+                exit_after = true;
+                break;
+            }
+            let side_node: NamespaceNode = *side_nodes_span
+                .at((height - height_offset - 1).try_into().unwrap());
+            if proof.key - sub_tree_start_index < rounding_factor / 2 {
+                node = hasher::node_digest(node, side_node);
+            } else {
+                node = hasher::node_digest(side_node, node);
+            }
+
+            height += 1;
+        };
+        if exit_after {
             return false;
         }
+
+        if stable_end != proof.num_leaves - 1 {
+            if proof_side_nodes_len.into() <= height - height_offset - 1 {
+                return false;
+            }
+            node =
+                hasher::node_digest(
+                    node, *proof.side_nodes.at((height - height_offset - 1).try_into().unwrap())
+                );
+            height += 1;
+        }
+
+        // All remaining elements in proof set will belong to a left sibling.
+        while height
+            - height_offset
+            - 1 < proof_side_nodes_len
+                .into() {
+                    node =
+                        hasher::node_digest(
+                            *proof.side_nodes.at((height - height_offset - 1).try_into().unwrap()),
+                            node
+                        );
+                    height += 1;
+                };
+
+        return super::namespace_node_eq(root, node);
     }
-}
 
-impl NamespacePartialEq of PartialEq<Namespace> {
-    #[inline(always)]
-    fn eq(lhs: @Namespace, rhs: @Namespace) -> bool {
-        let lhs_id: u256 = (*lhs.id).into();
-        let rhs_id: u256 = (*rhs.id).into();
-        if ((lhs_id == rhs_id) && (lhs.version == rhs.version)) {
-            return true;
-        } else {
-            return false;
+    fn verify_multi(
+        root: NamespaceNode,
+        proof: NamespaceMerkleMultiproof,
+        namespace: Namespace,
+        data: Array<Bytes>
+    ) -> bool {
+        // Hash all the leaves to get leaf nodes.
+        let mut nodes: Array<NamespaceNode> = array![];
+        let mut i: u32 = 0;
+        while i < data.len() {
+            nodes.append(hasher::leaf_digest(namespace, data.at(i)));
+            i += 1;
+        };
+
+        // Verify inclusion of leaf nodes.
+        return verify_multi_hashes(root, @proof, nodes);
+    }
+
+    fn verify_multi_hashes(
+        root: NamespaceNode, proof: @NamespaceMerkleMultiproof, leaf_nodes: Array<NamespaceNode>
+    ) -> bool {
+        let mut leaf_index: u256 = 0;
+        let mut left_subtrees: Array<NamespaceNode> = array![];
+
+        let mut i: u32 = 0;
+        while leaf_index != *proof.begin_key
+            && i < proof
+                .side_nodes
+                .len() {
+                    let subtree_size = _next_subtree_size(leaf_index, *proof.begin_key);
+                    left_subtrees.append(*proof.side_nodes.at(i));
+                    leaf_index += subtree_size;
+                    i += 1;
+                };
+
+        // estimate the leaf size of the subtree containing the proof range
+        let mut proof_range_subtree_estimate = utils::get_split_point(*proof.end_key) * 2;
+        if proof_range_subtree_estimate < 1 {
+            proof_range_subtree_estimate = 1;
         }
+
+        let (mut root_hash, proof_head, _, _) = _compute_root(
+            proof, leaf_nodes.span(), 0, proof_range_subtree_estimate, 0, 0
+        );
+        let mut i: u32 = proof_head.try_into().unwrap();
+        while i < proof
+            .side_nodes
+            .len() {
+                root_hash = hasher::node_digest(root_hash, *proof.side_nodes.at(i));
+                i += 1;
+            };
+
+        return super::namespace_node_eq(root_hash, root);
     }
-    #[inline(always)]
-    fn ne(lhs: @Namespace, rhs: @Namespace) -> bool {
-        let lhs_id: u256 = (*lhs.id).into();
-        let rhs_id: u256 = (*rhs.id).into();
-        if (lhs_id != rhs_id && lhs.version != rhs.version) {
-            return true;
+
+    // Gives the size of the subtree adjacent to `begin` that does not overlap `end`.
+    fn _next_subtree_size(begin: u256, end: u256) -> u256 {
+        let ideal: u256 = _bits_trailing_zeroes(begin);
+        let max: u256 = utils::bits_len(end - begin) - 1;
+        if ideal > max {
+            return U256BitShift::shl(1, max);
         } else {
-            return false;
+            return U256BitShift::shl(1, ideal);
+        }
+    }
+
+    // Returns the number of trailing zero bits in `x`.
+    fn _bits_trailing_zeroes(mut x: u256) -> u256 {
+        let mask: u256 = 1;
+        let mut count: u256 = 0;
+
+        while (x != 0 && mask & x == 0) {
+            count += 1;
+            x = U256BitShift::shr(x, 1);
+        };
+
+        return count;
+    }
+
+    // Computes the NMT root recursively.
+    fn _compute_root(
+        proof: @NamespaceMerkleMultiproof,
+        leaf_nodes: Span<NamespaceNode>,
+        begin: u256,
+        end: u256,
+        head_proof: u256,
+        head_leaves: u256
+    ) -> (NamespaceNode, u256, u256, bool) {
+        // Reached a leaf
+        if end - begin == 1 {
+            if *proof.begin_key <= begin && begin < *proof.end_key {
+                return _pop_leaves_if_non_empty(
+                    leaf_nodes, head_leaves, leaf_nodes.len().into(), head_proof
+                );
+            }
+
+            return _pop_proof_if_non_empty(proof.side_nodes.span(), head_proof, end, head_leaves);
+        }
+
+        if end <= *proof.begin_key || begin >= *proof.end_key {
+            return _pop_proof_if_non_empty(proof.side_nodes.span(), head_proof, end, head_leaves);
+        }
+
+        // Recursively get left and right subtree
+        let k: u256 = utils::get_split_point(end - begin);
+        let (left, new_head_proof_left, new_head_leaves_left, _) = _compute_root(
+            proof, leaf_nodes, begin, begin + k, head_proof, head_leaves
+        );
+        let (right, new_head_proof, new_head_leaves, right_is_nil) = _compute_root(
+            proof, leaf_nodes, begin + k, end, new_head_proof_left, new_head_leaves_left
+        );
+
+        if right_is_nil {
+            return (left, new_head_proof, new_head_leaves, false);
+        }
+        let hash = hasher::node_digest(left, right);
+        return (hash, new_head_proof, new_head_leaves, false);
+    }
+
+    fn _pop_leaves_if_non_empty(
+        nodes: Span<NamespaceNode>, head_leaves: u256, end: u256, head_proof: u256
+    ) -> (NamespaceNode, u256, u256, bool) {
+        let (node, new_head, is_nil) = _pop_if_non_empty(nodes, head_leaves, end);
+        return (node, head_proof, new_head, is_nil);
+    }
+
+    fn _pop_proof_if_non_empty(
+        nodes: Span<NamespaceNode>, head_proof: u256, end: u256, head_leaves: u256
+    ) -> (NamespaceNode, u256, u256, bool) {
+        let (node, new_head, is_nil) = _pop_if_non_empty(nodes, head_proof, end);
+        return (node, new_head, head_leaves, is_nil);
+    }
+
+    fn _pop_if_non_empty(
+        nodes: Span<NamespaceNode>, head: u256, end: u256
+    ) -> (NamespaceNode, u256, bool) {
+        if nodes.len() == 0 || head >= nodes.len().into() || head >= end {
+            let nid: Namespace = Default::default();
+            let node: NamespaceNode = NamespaceNode { min: nid, max: nid, digest: 0 };
+            return (node, head, true);
         }
+        return (*nodes.at(head.try_into().unwrap()), head + 1, false);
     }
 }