feat: Switch from sparse MT to concurrent MT (#449)

* feat: Switch from sparse MT to concurrent MT

Our concurrent Merkle tree implementation supports an append-only
case by ability to define the size of changelog as 0.

Therefore, there is no point in keeping multiple implementations. The
old implementation of sparse Merkle tree gets removed here.

* fix(programs): Store Merkle trees as byte arrays

Anchor is not smart enough to support const generics in account field
types. It also doesn't support using constants in array fields.

This change contains the following workaround:

* Storing Merkle trees as byte arrays.
* Defining their sizes as hard-coded integer...
* Providing methods which use unsafe Rust to cast byte slices to
  MT struct references.

Author: vadorovsky

Cargo.lock

@@ -6,6 +6,7 @@ description = "Concurrent Merkle tree implementation"
 license = "Apache-2.0"
 
 [dependencies]
+bytemuck = "1.14"
 light-hasher = { path = "../hasher", version = "0.1.0" }
 
 [dev-dependencies]

cli/accounts/merkle-tree-set-0.json

@@ -1,5 +1,6 @@
 use std::{cmp::Ordering, marker::PhantomData};
 
+use bytemuck::{Pod, Zeroable};
 use hash::compute_root;
 use light_hasher::{errors::HasherError, Hasher};
 
@@ -23,6 +24,7 @@ use crate::{changelog::ChangelogEntry, hash::compute_parent_node};
 /// Due to ability to make a decent number of concurrent update requests to be
 /// valid, no lock is necessary.
 #[repr(C)]
+#[derive(Copy, Clone)]
 pub struct ConcurrentMerkleTree<
     H,
     const HEIGHT: usize,
@@ -31,18 +33,18 @@ pub struct ConcurrentMerkleTree<
 > where
     H: Hasher,
 {
+    /// Index of the newest non-empty leaf.
+    pub next_index: u64,
+    /// History of roots.
+    pub roots: [[u8; 32]; MAX_ROOTS],
     /// History of Merkle proofs.
     pub changelog: [ChangelogEntry<HEIGHT>; MAX_CHANGELOG],
     /// Index of the newest changelog.
     pub current_changelog_index: u64,
-    /// History of roots.
-    pub roots: [[u8; 32]; MAX_ROOTS],
     /// Index of the newest root.
     pub current_root_index: u64,
     /// The newest Merkle proof.
     pub rightmost_proof: [[u8; 32]; HEIGHT],
-    /// Index of the newest non-empty leaf.
-    pub rightmost_index: u64,
     /// The newest non-empty leaf.
     pub rightmost_leaf: [u8; 32],
 
@@ -61,13 +63,54 @@ where
             roots: [[0u8; 32]; MAX_ROOTS],
             current_root_index: 0,
             rightmost_proof: [[0u8; 32]; HEIGHT],
-            rightmost_index: 0,
+            next_index: 0,
             rightmost_leaf: [0u8; 32],
             _hasher: PhantomData,
         }
     }
 }
 
+/// Mark `ConcurrentMerkleTree` as `Zeroable`, providing Anchor a guarantee
+/// that it can be always initialized with zeros.
+///
+/// # Safety
+///
+/// [`bytemuck`](bytemuck) is not able to ensure that our custom types (`Hasher`
+/// and `ConcurrentMerkleTree`) can be a subject of initializing with zeros. It
+/// also doesn't support structs with const generics (it would need to ensure
+/// alignment).
+///
+/// Therefore, it's our responsibility to guarantee that `ConcurrentMerkleTree`
+/// doesn't contain any fields which are not zeroable.
+unsafe impl<H, const HEIGHT: usize, const MAX_CHANGELOG: usize, const MAX_ROOTS: usize> Zeroable
+    for ConcurrentMerkleTree<H, HEIGHT, MAX_CHANGELOG, MAX_ROOTS>
+where
+    H: Hasher,
+{
+}
+
+/// Mark `ConcurrentMerkleTree` as `Pod` (Plain Old Data), providing Anchor a
+/// guarantee that it can be used in a zero-copy account.
+///
+/// # Safety
+///
+/// [`bytemuck`](bytemuck) is not able to ensure that our custom types (`Hasher`
+/// and `ConcurrentMerkleTree`) can be a subject of byte serialization. It also
+/// doesn't support structs with const generics (it would need to ensure
+/// alignment).
+///
+/// Therefore, it's our responsibility to guarantee that:
+///
+/// * `Hasher` and `ConcurrentMerkleTree` with given const generics are aligned.
+/// * They don't contain any fields which are not implementing `Copy` or are
+///   not an easy subject for byte serialization.
+unsafe impl<H, const HEIGHT: usize, const MAX_CHANGELOG: usize, const MAX_ROOTS: usize> Pod
+    for ConcurrentMerkleTree<H, HEIGHT, MAX_CHANGELOG, MAX_ROOTS>
+where
+    H: Hasher + Copy + 'static,
+{
+}
+
 impl<H, const HEIGHT: usize, const MAX_CHANGELOG: usize, const MAX_ROOTS: usize>
     ConcurrentMerkleTree<H, HEIGHT, MAX_CHANGELOG, MAX_ROOTS>
 where
@@ -236,14 +279,14 @@ where
             .ok_or(HasherError::RootsZero)? = node;
 
         // Update the rightmost proof. It has to be done only if tree is not full.
-        if self.rightmost_index < (1 << HEIGHT) {
-            if self.rightmost_index > 0 && leaf_index < self.rightmost_index as usize - 1 {
+        if self.next_index < (1 << HEIGHT) {
+            if self.next_index > 0 && leaf_index < self.next_index as usize - 1 {
                 // Update the rightmost proof with the current changelog entry when:
                 //
                 // * `rightmost_index` is greater than 0 (tree is non-empty).
                 // * The updated leaf is non-rightmost.
                 if let Some(proof) = changelog_entry
-                    .update_proof(self.rightmost_index as usize - 1, &self.rightmost_proof)
+                    .update_proof(self.next_index as usize - 1, &self.rightmost_proof)
                 {
                     self.rightmost_proof = proof;
                 }
@@ -272,7 +315,7 @@ where
         leaf_index: usize,
         proof: &[[u8; 32]; HEIGHT],
     ) -> Result<(), HasherError> {
-        let updated_proof = if self.rightmost_index > 0 && MAX_CHANGELOG > 0 {
+        let updated_proof = if self.next_index > 0 && MAX_CHANGELOG > 0 {
             match self.update_proof_or_leaf(changelog_index, leaf_index, proof) {
                 Some(proof) => proof,
                 // This case means that the leaf we are trying to update was
@@ -284,7 +327,7 @@ where
                 }
             }
         } else {
-            if leaf_index != self.rightmost_index as usize {
+            if leaf_index != self.next_index as usize {
                 return Err(HasherError::AppendOnly);
             }
             proof.to_owned()
@@ -296,11 +339,11 @@ where
 
     /// Appends a new leaf to the tree.
     pub fn append(&mut self, leaf: &[u8; 32]) -> Result<(), HasherError> {
-        if self.rightmost_index >= 1 << HEIGHT {
+        if self.next_index >= 1 << HEIGHT {
             return Err(HasherError::TreeFull);
         }
 
-        if self.rightmost_index == 0 {
+        if self.next_index == 0 {
             // NOTE(vadorovsky): This is not mentioned in the whitepaper, but
             // appending to an empty Merkle tree is a special case, where
             // `computer_parent_node` can't be called, because the usual
@@ -321,7 +364,7 @@ where
         } else {
             let mut current_node = *leaf;
             let mut intersection_node = self.rightmost_leaf;
-            let intersection_index = self.rightmost_index.trailing_zeros() as usize;
+            let intersection_index = self.next_index.trailing_zeros() as usize;
             let mut changelog_path = [[0u8; 32]; HEIGHT];
 
             for (i, item) in changelog_path.iter_mut().enumerate() {
@@ -334,7 +377,7 @@ where
                         intersection_node = compute_parent_node::<H>(
                             &intersection_node,
                             &self.rightmost_proof[i],
-                            self.rightmost_index as usize - 1,
+                            self.next_index as usize - 1,
                             i,
                         )?;
                         self.rightmost_proof[i] = empty_node;
@@ -347,7 +390,7 @@ where
                         current_node = compute_parent_node::<H>(
                             &current_node,
                             &self.rightmost_proof[i],
-                            self.rightmost_index as usize - 1,
+                            self.next_index as usize - 1,
                             i,
                         )?;
                     }
@@ -360,7 +403,7 @@ where
                 .get_mut(self.current_changelog_index as usize)
             {
                 *changelog_element =
-                    ChangelogEntry::new(current_node, changelog_path, self.rightmost_index as usize)
+                    ChangelogEntry::new(current_node, changelog_path, self.next_index as usize)
             }
             self.inc_current_root_index();
             *self
@@ -369,9 +412,21 @@ where
                 .ok_or(HasherError::RootsZero)? = current_node;
         }
 
-        self.rightmost_index += 1;
+        self.next_index += 1;
         self.rightmost_leaf = *leaf;
 
         Ok(())
     }
+
+    /// Appends a new pair of leaves to the tree.
+    pub fn append_two(
+        &mut self,
+        leaf_left: &[u8; 32],
+        leaf_right: &[u8; 32],
+    ) -> Result<(), HasherError> {
+        // TODO(vadorovsky): Instead of this naive double append, implement an
+        // optimized insertion of two leaves.
+        self.append(leaf_left)?;
+        self.append(leaf_right)
+    }
 }

merkle-tree/concurrent/Cargo.toml

@@ -80,7 +80,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 1);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 1);
+    assert_eq!(merkle_tree.next_index, 1);
     assert_eq!(merkle_tree.rightmost_leaf, leaf1);
 
     // Appending the 2nd leaf should result in recomputing the root due to the
@@ -123,7 +123,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 2);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 2);
+    assert_eq!(merkle_tree.next_index, 2);
     assert_eq!(merkle_tree.rightmost_leaf, leaf2);
 
     // Appending the 3rd leaf alters the next subtree on the right.
@@ -166,7 +166,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 3);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 3);
+    assert_eq!(merkle_tree.next_index, 3);
     assert_eq!(merkle_tree.rightmost_leaf, leaf3);
 
     // Appending the 4th leaf alters the next subtree on the right.
@@ -204,7 +204,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 4);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 4);
+    assert_eq!(merkle_tree.next_index, 4);
     assert_eq!(merkle_tree.rightmost_leaf, leaf4);
 }
 
@@ -257,7 +257,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 4);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 4);
+    assert_eq!(merkle_tree.next_index, 4);
     assert_eq!(merkle_tree.rightmost_leaf, leaf4);
 
     // Replace `leaf1`.
@@ -304,7 +304,7 @@ where
     // Note that we didn't create a new `expected_proof` here. We just re-used
     // the previous one.
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 4);
+    assert_eq!(merkle_tree.next_index, 4);
     assert_eq!(merkle_tree.rightmost_leaf, leaf4);
 
     // Replace `leaf2`.
@@ -348,7 +348,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 6);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 4);
+    assert_eq!(merkle_tree.next_index, 4);
     assert_eq!(merkle_tree.rightmost_leaf, leaf4);
 
     // Replace `leaf3`.
@@ -392,7 +392,7 @@ where
     assert_eq!(merkle_tree.root().unwrap(), expected_root);
     assert_eq!(merkle_tree.current_root_index, 7);
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 4);
+    assert_eq!(merkle_tree.next_index, 4);
     assert_eq!(merkle_tree.rightmost_leaf, leaf4);
 
     // Replace `leaf4`.
@@ -436,7 +436,7 @@ where
     assert_eq!(merkle_tree.current_root_index, 8);
     // This time `rightmost_*` fields should be changed.
     assert_eq!(merkle_tree.rightmost_proof, expected_proof);
-    assert_eq!(merkle_tree.rightmost_index, 4);
+    assert_eq!(merkle_tree.next_index, 4);
     assert_eq!(merkle_tree.rightmost_leaf, new_leaf4);
 }
 
@@ -711,7 +711,7 @@ fn compare_trees<H, const HEIGHT: usize, const MAX_CHANGELOG: usize, const MAX_R
         spl_concurrent_mt.rightmost_proof.proof
     );
     assert_eq!(
-        concurrent_mt.rightmost_index,
+        concurrent_mt.next_index,
         spl_concurrent_mt.rightmost_proof.index as u64
     );
     assert_eq!(