Introduce Concurrent Merkle Tree Implementation (#3455)

Author: ngundotra

Cargo.lock

@@ -18,6 +18,8 @@ members = [
   "governance/tools",
   "governance/chat/program",
   "libraries/math",
+  "libraries/concurrent-merkle-tree",
+  "libraries/merkle-tree-reference",
   "memo/program",
   "name-service/program",
   "record/program",

Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "spl-concurrent-merkle-tree"
+version = "0.1.0"
+description = "Solana Program Library Concurrent Merkle Tree"
+authors = ["Solana Maintainers <[email protected]>"]
+repository = "https://github.com/solana-labs/solana-program-library"
+license = "Apache-2.0"
+edition = "2018"
+
+[features]
+log = []
+sol-log = [ "log" ]
+
+[dependencies]
+solana-program = "1.10.33"
+bytemuck = "1.8.0"
+thiserror = "1.0.30"
+
+[dev-dependencies]
+rand_distr = "0.4.3"
+rand = "0.7"
+spl-merkle-tree-reference = { path = "../merkle-tree-reference" }
+tokio = { version = "1.12", features = ["full"] }
+
+[lib]
+crate-type = ["cdylib", "lib"]
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]

libraries/concurrent-merkle-tree/Cargo.toml

@@ -0,0 +1,78 @@
+use crate::{
+    hash::hash_to_parent,
+    node::{Node, EMPTY},
+};
+
+/// Stores the path of nodes changed in a tree by a Merkle tree operation
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[repr(C)]
+pub struct ChangeLog<const MAX_DEPTH: usize> {
+    /// Historical root value before Path was applied
+    pub root: Node,
+    /// Nodes of off-chain merkle tree
+    pub path: [Node; MAX_DEPTH],
+    /// Bitmap of node parity (used when hashing)
+    pub index: u32,
+    pub _padding: u32,
+}
+
+impl<const MAX_DEPTH: usize> ChangeLog<MAX_DEPTH> {
+    pub fn default() -> Self {
+        Self {
+            root: EMPTY,
+            path: [EMPTY; MAX_DEPTH],
+            index: 0,
+            _padding: 0,
+        }
+    }
+
+    pub fn new(root: Node, path: [Node; MAX_DEPTH], index: u32) -> Self {
+        Self {
+            root,
+            path,
+            index,
+            _padding: 0,
+        }
+    }
+
+    /// Returns the leaf value modified when the change log was recorded
+    pub fn get_leaf(&self) -> Node {
+        self.path[0]
+    }
+
+    /// Sets all change log values from a leaf and valid proof
+    pub fn replace_and_recompute_path(
+        &mut self,
+        index: u32,
+        mut node: Node,
+        proof: &[Node],
+    ) -> Node {
+        self.index = index;
+        for (i, sibling) in proof.iter().enumerate() {
+            self.path[i] = node;
+            hash_to_parent(&mut node, sibling, self.index >> i & 1 == 0);
+        }
+        self.root = node;
+        node
+    }
+
+    /// Fast forwards the given proof and corresponding leaf by applying an update from
+    /// the current change log
+    pub fn update_proof_or_leaf(
+        &self,
+        leaf_index: u32,
+        proof: &mut [Node; MAX_DEPTH],
+        leaf: &mut Node,
+    ) {
+        let padding: usize = 32 - MAX_DEPTH;
+        if leaf_index != self.index {
+            // This bit math is used to identify which node in the proof
+            // we need to swap for a corresponding node in a saved change log
+            let common_path_len = ((leaf_index ^ self.index) << padding).leading_zeros() as usize;
+            let critbit_index = (MAX_DEPTH - 1) - common_path_len;
+            proof[critbit_index] = self.path[critbit_index];
+        } else {
+            *leaf = self.get_leaf();
+        }
+    }
+}