From 949bbd83217522e6bb9a6fedc62b77f20f37bf38 Mon Sep 17 00:00:00 2001 From: Arch0125 Date: Fri, 12 Jul 2024 12:08:08 +0530 Subject: [PATCH 1/2] feat(imt): add imt implementation in rust Refs: #3 --- Cargo.lock | 29 ++++ crates/imt/Cargo.toml | 10 ++ crates/imt/LICENSE | 21 +++ crates/imt/src/hash.rs | 14 ++ crates/imt/src/imt.rs | 317 +++++++++++++++++++++++++++++++++++++++++ crates/imt/src/lib.rs | 2 + 6 files changed, 393 insertions(+) create mode 100644 crates/imt/Cargo.toml create mode 100644 crates/imt/LICENSE create mode 100644 crates/imt/src/hash.rs create mode 100644 crates/imt/src/imt.rs create mode 100644 crates/imt/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 21d186b..4669ae2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,35 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "crunchy" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" + [[package]] name = "example" version = "0.1.0" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "imt" +version = "0.0.1" +dependencies = [ + "hex", + "tiny-keccak", +] + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] diff --git a/crates/imt/Cargo.toml b/crates/imt/Cargo.toml new file mode 100644 index 0000000..3fdf351 --- /dev/null +++ b/crates/imt/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "imt" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true + +[dependencies] +hex = "0.4.3" +tiny-keccak = { version = "2.0.0", features = ["keccak"] } diff --git a/crates/imt/LICENSE b/crates/imt/LICENSE new file mode 100644 index 0000000..1763b4d --- /dev/null +++ b/crates/imt/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Pinco + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/crates/imt/src/hash.rs b/crates/imt/src/hash.rs new file mode 100644 index 0000000..63304d1 --- /dev/null +++ b/crates/imt/src/hash.rs @@ -0,0 +1,14 @@ +use tiny_keccak::{Hasher, Keccak}; + +pub fn keccak256_hash_function(nodes: Vec) -> String { + let mut keccak = Keccak::v256(); + let mut result = [0u8; 32]; + + for node in nodes { + keccak.update(node.as_bytes()); + } + + keccak.finalize(&mut result); + + hex::encode(result) +} diff --git a/crates/imt/src/imt.rs b/crates/imt/src/imt.rs new file mode 100644 index 0000000..9f5943f --- /dev/null +++ b/crates/imt/src/imt.rs @@ -0,0 +1,317 @@ +pub struct IMT { + nodes: Vec>, + zeroes: Vec, + hash: IMTHashFunction, + depth: usize, + arity: usize, +} + +pub struct IMTMerkleProof { + root: IMTNode, + leaf: IMTNode, + path_indices: Vec, + siblings: Vec>, +} + +pub type IMTNode = String; +pub type IMTHashFunction = fn(Vec) -> IMTNode; + +impl IMT { + pub fn new( + hash: IMTHashFunction, + depth: usize, + zero_value: IMTNode, + arity: usize, + leaves: Vec, + ) -> Result { + if leaves.len() > arity.pow(depth as u32) { + return Err("The tree cannot contain more than arity^depth leaves"); + } + + let mut imt = IMT { + nodes: vec![vec![]; depth + 1], + zeroes: vec![], + hash, + depth, + arity, + }; + + let mut current_zero = zero_value; + for _ in 0..depth { + imt.zeroes.push(current_zero.clone()); + current_zero = (imt.hash)(vec![current_zero; arity]); + } + + imt.nodes[0] = leaves; + + for level in 0..depth { + for index in 0..((imt.nodes[level].len() as f64 / arity as f64).ceil() as usize) { + let position = index * arity; + let children: Vec<_> = (0..arity) + .map(|i| { + imt.nodes[level] + .get(position + i) + .cloned() + .unwrap_or_else(|| imt.zeroes[level].clone()) + }) + .collect(); + + if let Some(next_level) = imt.nodes.get_mut(level + 1) { + next_level.push((imt.hash)(children)); + } + } + } + + Ok(imt) + } + + pub fn root(&mut self) -> Option { + self.nodes[self.depth].first().cloned() + } + + pub fn depth(&self) -> usize { + self.depth + } + + pub fn leaves(&self) -> Vec { + self.nodes[0].clone() + } + + pub fn arity(&self) -> usize { + self.arity + } + + pub fn insert(&mut self, leaf: IMTNode) -> Result<(), &'static str> { + if self.nodes[0].len() >= self.arity.pow(self.depth as u32) { + return Err("The tree is full"); + } + + let index = self.nodes[0].len(); + self.nodes[0].push(leaf); + self.update(index, self.nodes[0][index].clone()) + } + + pub fn update(&mut self, mut index: usize, new_leaf: IMTNode) -> Result<(), &'static str> { + if index >= self.nodes[0].len() { + return Err("The leaf does not exist in this tree"); + } + + let mut node = new_leaf; + self.nodes[0][index] = node.clone(); + + for level in 0..self.depth { + let position = index % self.arity; + let level_start_index = index - position; + let level_end_index = level_start_index + self.arity; + + let children: Vec<_> = (level_start_index..level_end_index) + .map(|i| { + self.nodes[level] + .get(i) + .cloned() + .unwrap_or_else(|| self.zeroes[level].clone()) + }) + .collect(); + + node = (self.hash)(children); + index /= self.arity; + + if self.nodes[level + 1].len() <= index { + self.nodes[level + 1].push(node.clone()); + } else { + self.nodes[level + 1][index] = node.clone(); + } + } + + Ok(()) + } + + pub fn delete(&mut self, index: usize) -> Result<(), &'static str> { + self.update(index, self.zeroes[0].clone()) + } + + pub fn create_proof(&self, index: usize) -> Result { + if index >= self.nodes[0].len() { + return Err("The leaf does not exist in this tree"); + } + + let mut siblings = Vec::with_capacity(self.depth); + let mut path_indices = Vec::with_capacity(self.depth); + let mut current_index = index; + + for level in 0..self.depth { + let position = current_index % self.arity; + let level_start_index = current_index - position; + let level_end_index = level_start_index + self.arity; + + path_indices.push(position); + let mut level_siblings = Vec::new(); + + for i in level_start_index..level_end_index { + if i != current_index { + level_siblings.push( + self.nodes[level] + .get(i) + .cloned() + .unwrap_or_else(|| self.zeroes[level].clone()), + ); + } + } + + siblings.push(level_siblings); + current_index /= self.arity; + } + + Ok(IMTMerkleProof { + root: self.nodes[self.depth][0].clone(), + leaf: self.nodes[0][index].clone(), + path_indices, + siblings, + }) + } + + pub fn verify_proof(&self, proof: &IMTMerkleProof) -> bool { + let mut node = proof.leaf.clone(); + + for (i, sibling) in proof.siblings.iter().enumerate() { + let mut children = sibling.clone(); + children.insert(proof.path_indices[i], node); + + node = (self.hash)(children); + } + + node == proof.root + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn simple_hash_function(nodes: Vec) -> String { + nodes.join(",") + } + + #[test] + fn test_new_imt() { + let hash: IMTHashFunction = simple_hash_function; + let imt = IMT::new(hash, 3, "zero".to_string(), 2, vec![]); + + assert!(imt.is_ok()); + } + + #[test] + fn test_insertion() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec![]).unwrap(); + + assert!(imt.insert("leaf1".to_string()).is_ok()); + } + + #[test] + fn test_delete() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec!["leaf1".to_string()]).unwrap(); + + assert!(imt.delete(0).is_ok()); + } + + #[test] + fn test_update() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec!["leaf1".to_string()]).unwrap(); + + assert!(imt.update(0, "new_leaf".to_string()).is_ok()); + } + + #[test] + fn test_create_and_verify_proof() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec!["leaf1".to_string()]).unwrap(); + imt.insert("leaf2".to_string()).unwrap(); + + let proof = imt.create_proof(0); + assert!(proof.is_ok()); + + let proof = proof.unwrap(); + assert!(imt.verify_proof(&proof)); + } + + #[test] + fn should_not_initialize_with_too_many_leaves() { + let hash: IMTHashFunction = simple_hash_function; + let leaves = vec![ + "leaf1".to_string(), + "leaf2".to_string(), + "leaf3".to_string(), + "leaf4".to_string(), + "leaf5".to_string(), + ]; + let imt = IMT::new(hash, 2, "zero".to_string(), 2, leaves); + assert!(imt.is_err()); + } + + #[test] + fn should_not_insert_in_full_tree() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new( + hash, + 1, + "zero".to_string(), + 2, + vec!["leaf1".to_string(), "leaf2".to_string()], + ) + .unwrap(); + + let result = imt.insert("leaf3".to_string()); + assert!(result.is_err()); + } + + #[test] + fn should_not_delete_nonexistent_leaf() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec!["leaf1".to_string()]).unwrap(); + + let result = imt.delete(1); + assert!(result.is_err()); + } + + #[test] + fn test_root() { + let hash: IMTHashFunction = simple_hash_function; + let mut imt = IMT::new( + hash, + 2, + "zero".to_string(), + 2, + vec!["leaf1".to_string(), "leaf2".to_string()], + ) + .unwrap(); + + assert_eq!(imt.root(), Some("leaf1,leaf2,zero,zero".to_string())); + } + + #[test] + fn test_leaves() { + let hash: IMTHashFunction = simple_hash_function; + let imt = IMT::new( + hash, + 2, + "zero".to_string(), + 2, + vec!["leaf1".to_string(), "leaf2".to_string()], + ) + .unwrap(); + + assert_eq!(imt.leaves(), vec!["leaf1".to_string(), "leaf2".to_string()]); + } + + #[test] + fn test_depth_and_arity() { + let hash: IMTHashFunction = simple_hash_function; + let imt = IMT::new(hash, 3, "zero".to_string(), 2, vec![]).unwrap(); + + assert_eq!(imt.depth(), 3); + assert_eq!(imt.arity(), 2); + } +} diff --git a/crates/imt/src/lib.rs b/crates/imt/src/lib.rs new file mode 100644 index 0000000..3ca61dd --- /dev/null +++ b/crates/imt/src/lib.rs @@ -0,0 +1,2 @@ +pub mod hash; +pub mod imt; From e9cbfd2b1f37239795c69a04e767af42ffad9c97 Mon Sep 17 00:00:00 2001 From: Arch0125 Date: Fri, 12 Jul 2024 12:12:46 +0530 Subject: [PATCH 2/2] style(imt): lint fixes Lint fixes done after make fix command Refs: #3 --- crates/imt/src/imt.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/imt/src/imt.rs b/crates/imt/src/imt.rs index 9f5943f..8ae7d01 100644 --- a/crates/imt/src/imt.rs +++ b/crates/imt/src/imt.rs @@ -97,7 +97,7 @@ impl IMT { } let mut node = new_leaf; - self.nodes[0][index] = node.clone(); + self.nodes[0][index].clone_from(&node); for level in 0..self.depth { let position = index % self.arity; @@ -119,7 +119,7 @@ impl IMT { if self.nodes[level + 1].len() <= index { self.nodes[level + 1].push(node.clone()); } else { - self.nodes[level + 1][index] = node.clone(); + self.nodes[level + 1][index].clone_from(&node); } }