diff --git a/Cargo.toml b/Cargo.toml index ff78b0c2..759af077 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,6 +47,5 @@ perf = ["circuits-batcher/perf"] mongo-std-sync = ["mongodb/sync"] mongo-tokio-sync = ["mongodb/tokio-sync"] cuda = ["halo2_proofs/cuda"] -complex-leaf = [] profile = ["ark-std/print-trace", "halo2_proofs/profile", "circuits-batcher/profile"] diff --git a/src/circuits/host.rs b/src/circuits/host.rs index 7a4df423..79ae95f3 100644 --- a/src/circuits/host.rs +++ b/src/circuits/host.rs @@ -154,6 +154,7 @@ impl HostOpConfig { indicator: Fr, enable: bool, ) -> Result<(Limb, Limb), Error> { + /* values = [v0, v1, v2] in little endian */ let mut rev = values.clone(); let len = values.len(); rev.reverse(); diff --git a/src/host/db.rs b/src/host/db.rs index 3b8ff272..f1bf4adc 100644 --- a/src/host/db.rs +++ b/src/host/db.rs @@ -7,6 +7,8 @@ use mongodb::{ options::DropCollectionOptions, sync::{Client, Collection}, }; +use mongodb::options::IndexOptions; +use mongodb::IndexModel; use crate::host::datahash::DataHashRecord; use crate::host::mongomerkle::MerkleRecord; @@ -67,6 +69,19 @@ impl MongoDB { collection.drop(options) } + pub fn create_index(&self) -> Result<(), mongodb::error::Error> { + let cname = get_collection_name(MONGODB_MERKLE_NAME_PREFIX.to_string(), self.cname_id); + let collection = self.get_collection::(MONGODB_DATABASE.to_string(), cname.to_string())?; + let index_model = IndexModel::builder() + .keys(doc! { "_id": 1 }) + .options(IndexOptions::builder().unique(true).build()) + .build(); + // Create the index + collection + .create_index(index_model, None)?; + Ok(()) + } + pub fn merkel_collection(&self) -> Result, mongodb::error::Error> { let cname = get_collection_name(MONGODB_MERKLE_NAME_PREFIX.to_string(), self.cname_id); self.get_collection::(MONGODB_DATABASE.to_string(), cname.to_string()) diff --git a/src/host/merkle.rs b/src/host/merkle.rs index e99b336d..93a5d4d9 100644 --- a/src/host/merkle.rs +++ b/src/host/merkle.rs @@ -1,14 +1,6 @@ -use crate::host::db::TreeDB; -use std::cell::RefCell; use std::error::Error; use std::fmt; use std::fmt::Debug; -use std::rc::Rc; - -/* -const LEAF_SIG: u8 = 0u8; -const INTERNAL_SIG: u8 = 1u8; -*/ #[derive(Debug)] pub enum MerkleErrorCode { @@ -49,12 +41,14 @@ impl fmt::Display for MerkleError { impl Error for MerkleError {} -pub trait MerkleNode { +pub trait MerkleNode: Clone { + fn new_leaf(index: usize, data: &Vec) -> Self; fn hash(&self) -> H; + fn set_hash(&mut self, hash: H); fn index(&self) -> u64; fn set(&mut self, data: &Vec); - fn left(&self) -> Option; // hash of left child - fn right(&self) -> Option; // hash of right child + fn descendant(&self, offset: usize) -> Option; + fn set_descendant(&mut self, index: usize, hash: H); } #[derive(Debug, Clone)] @@ -65,33 +59,121 @@ pub struct MerkleProof { pub index: u64, } +/// get the offset of the node at its current level +/// 0 +/// 1 2 +/// 3 4 5 6 +/// 7 8 9 10 11 12 13 14 +/// get_offset(12) = 12 - 7 = 5 pub fn get_offset(index: u64) -> u64 { let height = (index + 1).ilog2(); let full = (1u64 << height) - 1; index - full } +/// get the index from leaf to the root +/// root index is not included in the result as root index is always 0 +/// Example: Given D=3 and a merkle tree as follows: +/// 0 +/// 1 2 +/// 3 4 5 6 +/// 7 8 9 10 11 12 13 14 +/// get_path(7) = [3, 1] +/// get_path(14) = [6, 2] +fn get_path(index: u64) -> Vec { + let mut height = (index + 1).ilog2(); + let round = height; + let full = (1u64 << height) - 1; + let mut p = index - full; + let mut path = vec![]; + for _ in 0..round { + let full = (1u64 << height) - 1; + // Calculate the index of current node + let i = full + p; + path.insert(0, i); + height = height - 1; + // Caculate the offset of parent + p = p / 2; + } + assert!(p == 0); + path +} + +fn get_path_binary(index: u64) -> Vec { + let height = (index + 1).ilog2(); + let round = height; + let left_most = (1u64 << height) - 1; + let mut c = index - left_most; + let mut binary_path = vec![]; + for _ in 0..round { + binary_path.push(c & 1); + c = c >> 1; + } + binary_path.reverse(); + binary_path +} + + + +pub fn binary_path_to_index(path: &[u64]) -> u64 { + let offset = path.iter().fold(0, |acc, e| (acc << 1) + e); + offset + (1u64 << path.len()) - 1 +} + +/// get a list of node index based on the binary path from root to leaf +pub fn binary_path_to_path(path: &[u64]) -> Vec { + let mut ret = vec![]; + let mut acc = vec![]; + for c in path.iter() { + acc.push(*c); + ret.push(binary_path_to_index(acc.as_slice())); + } + ret +} + +pub fn get_sibling_index(index: u64) -> u64 { + if index % 2 == 1 { + index + 1 + } else { + index - 1 + } +} + + + pub trait MerkleTree { type Node: MerkleNode; - type Id; - type Root; - - /// Create a new merkletree and connect it with a given merkle root. - /// If the root is None then the default root with all leafs are empty is used. - fn construct(addr: Self::Id, id: Self::Root, db: Option>>) -> Self; fn hash(a: &H, b: &H) -> H; - fn set_parent(&mut self, index: u64, hash: &H, left: &H, right: &H) -> Result<(), MerkleError>; - fn set_leaf(&mut self, leaf: &Self::Node) -> Result<(), MerkleError>; - fn set_leaf_and_parents( - &mut self, - leaf: &Self::Node, - parents: [(u64, H, H, H); D], - ) -> Result<(), MerkleError>; - fn get_node_with_hash(&self, index: u64, hash: &H) -> Result; - fn get_root_hash(&self) -> H; + fn get_node_with_hash(&self, index: u64, hash: &H) -> Result; fn update_root_hash(&mut self, hash: &H); + fn chunk_depth() -> usize; + + fn hash_with_index(primary: &H, assist: &H, index: u64) -> H { + let offset = get_offset(index); + let (a, b) = if offset % 2 == 1 { + (assist, primary) + } else { + (primary, assist) + }; + Self::hash(a, b) + } + + fn update_leaf(node: &mut Self::Node, path: &[u64], hash: H) { + let mut node_idxs = binary_path_to_path(path); + node_idxs.reverse(); + let mut new_hash = hash; + for index in node_idxs { + let offset = (index - 1) as usize; + let siboffset = (get_sibling_index(index) - 1) as usize; + node.set_descendant(offset as usize, new_hash.clone()); + let primary = node.descendant(offset); + let assist = node.descendant(siboffset); + new_hash = Self::hash_with_index(&primary.unwrap(), &assist.unwrap(), index); + }; + node.set_hash(new_hash); + } fn boundary_check(&self, index: u64) -> Result<(), MerkleError> { if index >= (2_u64.pow(D as u32 + 1) - 1) { @@ -105,17 +187,17 @@ pub trait MerkleTree { } } - /* - * Check that an index is a leaf. - * Example: Given D=2 and a merkle tree as follows: - * 0 - * 1 2 - * 3 4 5 6 - * then leaf index >= 3 which is (2^D - 1) - * - * Moreover, nodes at depth k start at - * first = 2^k-1, last = 2^{k+1}-2 - */ + /// + /// Check that an index is a leaf. + /// Example: Given D=2 and a merkle tree as follows: + /// 0 + /// 1 2 + /// 3 4 5 6 + /// then leaf index >= 3 which is (2^D - 1) + /// + /// Moreover, nodes at depth k start at + /// first = 2^k-1, last = 2^{k+1}-2 + /// fn leaf_check(&self, index: u64) -> Result<(), MerkleError> { if (index as u64) >= (2_u64.pow(D as u32) - 1) && (index as u64) < (2_u64.pow((D as u32) + 1) - 1) @@ -130,71 +212,82 @@ pub trait MerkleTree { } } - fn get_sibling_index(&self, index: u64) -> u64 { - if index % 2 == 1 { - index + 1 - } else { - index - 1 - } + fn get_path(&self, index: u64) -> Result <[u64; D], MerkleError> { + self.leaf_check(index)?; + let r = get_path(index); + assert!(r.len() == D); + Ok(r.try_into().unwrap()) } - /// get the index from leaf to the root - /// root index is not included in the result as root index is always 0 - /// Example: Given D=3 and a merkle tree as follows: - /// 0 - /// 1 2 - /// 3 4 5 6 - /// 7 8 9 10 11 12 13 14 - /// get_path(7) = [3, 1] - /// get_path(15) = [6, 2] - fn get_path(&self, index: u64) -> Result<[u64; D], MerkleError> { + + fn get_path_binary(&self, index: u64) -> Result <[u64; D], MerkleError> { self.leaf_check(index)?; - let mut height = (index + 1).ilog2(); - let round = height; - let full = (1u64 << height) - 1; - let mut p = index - full; - let mut path = vec![]; - for _ in 0..round { - let full = (1u64 << height) - 1; - // Calculate the index of current node - let i = full + p; - path.insert(0, i); - height = height - 1; - // Caculate the offset of parent - p = p / 2; - } - assert!(p == 0); - Ok(path.try_into().unwrap()) + let binary_path = get_path_binary(index); + assert!(binary_path.len() == D); + Ok(binary_path.try_into().unwrap()) } - fn get_leaf_with_proof( + /// Get the nodes and proof related to a leaf node with index + /// the nodes starts from the root and all the way to the leaf + fn trace_leaf_with_proof( &self, index: u64, - ) -> Result<(Self::Node, MerkleProof), MerkleError> { + update: Option + ) -> Result<(Vec, MerkleProof), MerkleError> { self.leaf_check(index)?; - let paths = self.get_path(index)?.to_vec(); - // We push the search from the top + let paths = self.get_path_binary(index)?.to_vec(); + let index_paths = self.get_path(index)?.to_vec(); + assert!((paths.len() % Self::chunk_depth()) == 0); + let mut path_chunks = paths.chunks_exact(Self::chunk_depth()).collect::>(); + + // We start the search from the root hash let hash = self.get_root_hash(); - let mut acc = 0; - let mut acc_node = self.get_node_with_hash(acc, &hash)?; - let assist: Vec = paths - .into_iter() - .map(|child| { - let (hash, sibling_hash) = if (acc + 1) * 2 == child + 1 { - // left child - (acc_node.left().unwrap(), acc_node.right().unwrap()) - } else { - assert!((acc + 1) * 2 == child); - (acc_node.right().unwrap(), acc_node.left().unwrap()) - }; - acc = child; - acc_node = self.get_node_with_hash(acc, &hash)?; - Ok(sibling_hash) - }) - .collect::, _>>()?; - let hash = acc_node.hash(); + let root_node = self.get_node_with_hash(0, &hash)?; + // nodes are from top to bottom + let (assist, mut nodes) = path_chunks + .iter() + .fold( + (vec![], vec![root_node]), + |(mut assist, mut nodes), chunk| { + let acc_node = nodes.last().unwrap().clone(); + let node_idxs = binary_path_to_path(chunk.to_vec().as_slice()); + //println!("node_idxs is {:?}", node_idxs); + let primary_hashs = node_idxs + .iter() + .map(|x| acc_node.descendant(*x as usize - 1).unwrap()) + .collect::>(); + let mut sibling_hashs = node_idxs + .iter() + .map(|x| acc_node.descendant(get_sibling_index(*x) as usize - 1).unwrap()) + .collect::>(); + + //println!("primary hash {:?}", primary_hashs); + let last_hash = primary_hashs.last().unwrap(); + let height = nodes.len() * Self::chunk_depth(); + let index = index_paths[height - 1]; + //println!("index is {}", index); + let acc_node = self.get_node_with_hash(index as u64, last_hash).unwrap(); + nodes.push(acc_node); + assist.append(&mut sibling_hashs); + (assist, nodes) + }); + + let hash = nodes.last().unwrap().hash(); + + + if let Some(mut update_hash) = update { + path_chunks.reverse(); + path_chunks.iter().fold(nodes.len() - 1, |acc, chunk| { + let mut node = nodes.get_mut(acc - 1).unwrap(); + Self::update_leaf(&mut node, chunk, update_hash.clone()); + update_hash = node.hash(); + acc - 1 + }); + } + + assert!(assist.len() == D); Ok(( - acc_node, + nodes, MerkleProof { source: hash, root: self.get_root_hash(), @@ -204,32 +297,30 @@ pub trait MerkleTree { )) } - fn set_leaf_with_proof(&mut self, leaf: &Self::Node) -> Result, MerkleError> { + fn get_leaf_with_proof( + &self, + index: u64, + ) -> Result<(Self::Node, MerkleProof), MerkleError> { + let (mut nodes, proof) = self.trace_leaf_with_proof(index, None)?; + return Ok((nodes.pop().unwrap(), proof)) + } + + + fn update_nodes(&mut self, nodes: Vec) -> Result<(), MerkleError>; + + fn set_leaf_with_proof( + &mut self, + leaf: &Self::Node, + ) -> Result, MerkleError> { let index = leaf.index(); - let mut hash = leaf.hash(); - let (_, mut proof) = self.get_leaf_with_proof(index)?; - proof.source = hash.clone(); - let mut p = get_offset(index); - //self.set_leaf(leaf)?; - let mut parents = vec![]; - for i in 0..D { - let cur_hash = hash; - let depth = D - i - 1; - let (left, right) = if p % 2 == 1 { - (&proof.assist[depth], &cur_hash) - } else { - (&cur_hash, &proof.assist[depth]) - }; - hash = Self::hash(left, right); - p = p / 2; - let index = p + (1 << depth) - 1; - //self.set_parent(index, &hash, left, right)?; - parents.push((index, hash.clone(), left.clone(), right.clone())); - } - self.set_leaf_and_parents(leaf, parents.try_into().unwrap())?; + let (mut nodes, mut proof) = self.trace_leaf_with_proof(index, Some(leaf.hash()))?; + let hash = nodes[0].hash(); self.update_root_hash(&hash); proof.root = hash; - Ok(proof) + proof.source = leaf.hash(); + nodes.push(leaf.clone()); + self.update_nodes(nodes)?; + return Ok(proof) } fn update_leaf_data_with_proof( @@ -237,9 +328,8 @@ pub trait MerkleTree { index: u64, data: &Vec, ) -> Result, MerkleError> { - let (mut leaf, _) = self.get_leaf_with_proof(index)?; - leaf.set(data); - self.set_leaf_with_proof(&leaf) + let node = Self::Node::new_leaf(index as usize, data); + self.set_leaf_with_proof(&node) } fn verify_proof(&self, proof: &MerkleProof) -> Result { @@ -261,16 +351,22 @@ pub trait MerkleTree { #[cfg(test)] mod tests { - use crate::host::db::TreeDB; - use crate::host::merkle::{MerkleError, MerkleNode, MerkleTree}; - use std::cell::RefCell; - use std::rc::Rc; + use crate::host::merkle::{get_path, get_path_binary, MerkleError, MerkleNode, MerkleTree}; + + use super::binary_path_to_path; struct MerkleAsArray { data: [u64; 127], // 2^7-1 and depth = 6 root_hash: u64, } impl MerkleAsArray { + fn construct() -> Self { + MerkleAsArray { + data: [0 as u64; 127], + root_hash: 0, + } + } + fn debug(&self) { let mut start = 0; for i in 0..6 { @@ -294,6 +390,7 @@ mod tests { } } + #[derive(Clone, Debug)] struct MerkleU64Node { pub value: u64, pub index: u64, @@ -302,6 +399,17 @@ mod tests { } impl MerkleNode for MerkleU64Node { + fn set_descendant(&mut self, index: usize, v: u64) { + if index == 0 { + self.left = v; + } else { + self.right = v; + } + } + fn set_hash(&mut self, v: u64) { + self.value = v; + } + fn index(&self) -> u64 { self.index } @@ -312,38 +420,54 @@ mod tests { let v: [u8; 8] = value.clone().try_into().unwrap(); self.value = u64::from_le_bytes(v); } - fn right(&self) -> Option { - Some(self.right) + + fn new_leaf(index: usize, value: &Vec) -> MerkleU64Node { + let v: [u8; 8] = value.clone().try_into().unwrap(); + MerkleU64Node { + value: u64::from_le_bytes(v), + index: index as u64, + left: 0, + right: 0, + } } - fn left(&self) -> Option { - Some(self.left) + + fn descendant(&self, i: usize) -> Option { + if i==0 { + Some(self.left) + } else if i==1 { + Some(self.right) + } else { + None + } } } impl MerkleTree for MerkleAsArray { - type Id = String; - type Root = String; type Node = MerkleU64Node; - fn construct( - _addr: Self::Id, - _id: Self::Root, - _db: Option>>, - ) -> Self { - MerkleAsArray { - data: [0 as u64; 127], - root_hash: 0, - } - } + fn hash(a: &u64, b: &u64) -> u64 { a + b } + + fn chunk_depth() -> usize { + 1 + } + fn get_root_hash(&self) -> u64 { return self.data[0]; } + fn update_root_hash(&mut self, hash: &u64) { self.root_hash = hash.clone(); } + fn update_nodes(&mut self, nodes: Vec) -> Result<(), MerkleError> { + for node in nodes.into_iter() { + self.data[node.index() as usize] = node.value; + } + Ok (()) + } + fn get_node_with_hash(&self, index: u64, _hash: &u64) -> Result { self.boundary_check(index)?; @@ -363,64 +487,40 @@ mod tests { right: right_val, }) } + } - fn set_parent( - &mut self, - index: u64, - hash: &u64, - left: &u64, - right: &u64, - ) -> Result<(), MerkleError> { - self.boundary_check(index)?; - self.data[index as usize] = *hash; - - let height = (index + 1).ilog2(); - if height < Self::height() { - // not leaf node - let (left_child_index, right_child_index) = MerkleAsArray::get_child_index(index); - self.data[left_child_index as usize] = *left; - self.data[right_child_index as usize] = *right; - } - - Ok(()) - } - fn set_leaf(&mut self, leaf: &Self::Node) -> Result<(), MerkleError> { - self.leaf_check(leaf.index())?; - self.data[leaf.index() as usize] = leaf.value; - Ok(()) + #[test] + fn test_binary_path_to_path() { + for index in 1..128 { + let path_binary = get_path_binary(index); + let path_b2p = binary_path_to_path(&path_binary); + let path = get_path(index); + assert_eq!(path_b2p, path); } + } - fn set_leaf_and_parents( - &mut self, - leaf: &Self::Node, - parents: [(u64, u64, u64, u64); 6], - ) -> Result<(), MerkleError> { - // set leaf - self.leaf_check(leaf.index())?; - self.data[leaf.index() as usize] = leaf.value; - - // set parent - for (index, hash, left, right) in parents.iter() { - self.boundary_check(*index)?; - self.data[*index as usize] = *hash; - - let height = (index + 1).ilog2(); - if height < Self::height() { - // not leaf node - let (left_child_index, right_child_index) = - MerkleAsArray::get_child_index(*index); - self.data[left_child_index as usize] = *left; - self.data[right_child_index as usize] = *right; - } - } - Ok(()) - } + #[test] + fn test_path() { + assert_eq!(get_path(0), vec![] as Vec); + assert_eq!(get_path(1), vec![1]); + assert_eq!(get_path(2), vec![2]); + assert_eq!(get_path(3), vec![1,3]); + assert_eq!(get_path(4), vec![1,4]); + assert_eq!(get_path(5), vec![2,5]); + assert_eq!(get_path(6), vec![2,6]); + assert_eq!(get_path(7), vec![1,3,7]); } + + #[test] fn test_merkle_path() { - let mut mt = MerkleAsArray::construct("test".to_string(), "test".to_string(), None); - let (mut leaf, _) = mt.get_leaf_with_proof(2_u64.pow(6) - 1).unwrap(); + let mut mt = MerkleAsArray::construct(); + let (traces, proof) = mt.trace_leaf_with_proof(2_u64.pow(6) - 1, None).unwrap(); + println!("traces is {:?}", traces); + let (mut leaf, proof) = mt.get_leaf_with_proof(2_u64.pow(6) - 1).unwrap(); + println!("leaf is {:?}", leaf); + println!("proof is {:?}", proof); leaf.value = 1; let _proof = mt.set_leaf_with_proof(&leaf).unwrap(); diff --git a/src/host/mongomerkle.rs b/src/host/mongomerkle.rs index 40889a9e..1adf437d 100644 --- a/src/host/mongomerkle.rs +++ b/src/host/mongomerkle.rs @@ -27,6 +27,17 @@ where } } +fn deserialize_bytes_from_binary<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + match Bson::deserialize(deserializer) { + Ok(Bson::Binary(bytes)) => Ok(bytes.bytes.to_vec()), + Ok(..) => Err(Error::invalid_value(Unexpected::Enum, &"Bson::Binary")), + Err(e) => Err(e), + } +} + fn serialize_bytes_as_binary(bytes: &[u8], serializer: S) -> Result where S: Serializer, @@ -111,26 +122,38 @@ impl MongoMerkle { //the input records must be in one leaf path pub fn update_records(&mut self, records: &Vec) -> Result<(), anyhow::Error> { + /* + for r in records.iter() { + println!("update record: {:?}", r.hash()); + } + */ self.db.borrow_mut().set_merkle_records(records)?; Ok(()) } + pub fn get_node_default_hash(&self, height: usize) -> [[u8; 32]; CHUNK_DEPTH] { + let mut default_hash = vec![]; + for i in 0..CHUNK_DEPTH { + let hash = self.get_default_hash(height + i + 1).unwrap(); + default_hash.push(hash); + } + default_hash.try_into().unwrap() + } + pub fn generate_default_node(&self, index: u64) -> Result { - let height = (index + 1).ilog2(); - let default = self.get_default_hash(height as usize)?; - let child_hash = if height == Self::height() as u32 { - None + let height = (index + 1).ilog2() as usize; + let default_hash = if height < Self::height() { + Some(self.get_node_default_hash(height as usize)) } else { - let hash = self.get_default_hash((height + 1) as usize)?; - Some(hash) + assert!(height == Self::height()); + None }; - Ok(MerkleRecord { index, - hash: default, + default_hash, + hash: self.get_default_hash(height)?, data: None, - left: child_hash, - right: child_hash, + descendants: vec![] }) } @@ -152,8 +175,9 @@ impl MongoMerkle { index: u64, hash: &[u8; 32], ) -> Result { - let height = (index + 1).ilog2(); - let default_hash = self.get_default_hash(height as usize)?; + let height = (index + 1).ilog2() as usize; + assert!(height <= Self::height()); + let default_hash = self.get_default_hash(height)?; if &default_hash == hash { self.generate_default_node(index) } else { @@ -161,6 +185,9 @@ impl MongoMerkle { Ok(Some(mut node)) => { // The index of MerkleRecord was not stored in db, so it needs to be assigned here. node.index = index; + if height != Self::height() { + node.default_hash = Some(self.get_node_default_hash(height as usize)) + }; Ok(node) } Ok(None) => Err(MerkleError::new( @@ -178,29 +205,23 @@ impl MongoMerkle { } } +const CHUNK_DEPTH: usize = 4; +const CHUNK_VOLUME: usize = 31; + #[derive(Debug, Serialize, Deserialize, Clone)] pub struct MerkleRecord { // The index will not to be stored in db. #[serde(skip_serializing, skip_deserializing, default)] pub index: u64, + #[serde(skip_serializing, skip_deserializing, default)] + pub default_hash: Option<[[u8; 32]; CHUNK_DEPTH]>, #[serde(serialize_with = "self::serialize_bytes_as_binary")] #[serde(deserialize_with = "self::deserialize_u256_as_binary")] #[serde(rename = "_id")] pub hash: [u8; 32], - #[serde( - skip_serializing_if = "Option::is_none", - serialize_with = "self::serialize_option_bytes_as_binary", - default - )] - #[serde(deserialize_with = "self::deserialize_option_u256_as_binary")] - pub left: Option<[u8; 32]>, - #[serde( - skip_serializing_if = "Option::is_none", - serialize_with = "self::serialize_option_bytes_as_binary", - default - )] - #[serde(deserialize_with = "self::deserialize_option_u256_as_binary")] - pub right: Option<[u8; 32]>, + #[serde(serialize_with = "self::serialize_bytes_as_binary")] + #[serde(deserialize_with = "self::deserialize_bytes_from_binary")] + pub descendants: Vec, // 2^4 - 1 - 1 (root) #[serde( skip_serializing_if = "Option::is_none", serialize_with = "self::serialize_option_bytes_as_binary", @@ -214,9 +235,38 @@ impl MerkleNode<[u8; 32]> for MerkleRecord { fn index(&self) -> u64 { self.index } + fn hash(&self) -> [u8; 32] { self.hash } + + fn set_hash(&mut self, hash: [u8; 32]) { + self.hash = hash; + } + + fn set_descendant(&mut self, index: usize, hash: [u8; 32]) { + for c in 0..CHUNK_VOLUME { + if c * 33 < self.descendants.len() { + if self.descendants[c*33] == index as u8 { + for i in 0..32 { + let offset = c * 33 + i + 1; + self.descendants[offset] = hash[i]; + } + return; + } + } + } + self.descendants.push(index as u8); + for h in hash.to_vec().iter() { + self.descendants.push(*h); + } + } + + fn new_leaf(index: usize, data: &Vec) -> Self{ + let mut leaf = MerkleRecord::new(index as u64); + leaf.set(data); + return leaf + } fn set(&mut self, data: &Vec) { let mut hasher = MERKLE_LEAF_HASHER.clone(); self.data = Some(data.clone().try_into().unwrap()); @@ -232,22 +282,26 @@ impl MerkleNode<[u8; 32]> for MerkleRecord { .collect::>(); let values: [Fr; 2] = batchdata.try_into().unwrap(); - cfg_if::cfg_if! { - if #[cfg(feature="complex-leaf")] { - hasher.update(&values); - self.hash = hasher.squeeze().to_repr(); - } else { - self.hash = hasher.update_exact(&values).to_repr(); - } - } + self.hash = hasher.update_exact(&values).to_repr(); //println!("update with values {:?}", values); //println!("update with new hash {:?}", self.hash); } - fn right(&self) -> Option<[u8; 32]> { - self.right - } - fn left(&self) -> Option<[u8; 32]> { - self.left + fn descendant(&self, offset: usize ) -> Option<[u8; 32]> { + for c in 0..CHUNK_VOLUME { + if c * 33 < self.descendants.len() { + if self.descendants[c*33] == offset as u8 { + let offset = c * 33 + 1; + let r = self.descendants[offset .. offset + 32].to_vec(); + return Some(r.try_into().unwrap()) + } + } + } + let index = offset + 1; + let height = (index + 1).ilog2() as usize; + //println!("offset is {} {} {}", offset, height, self.descendants.len()); + return self.default_hash.map( + |d| d[height - 1] + ); } } @@ -255,10 +309,10 @@ impl MerkleRecord { fn new(index: u64) -> Self { MerkleRecord { index, + default_hash: None, hash: [0; 32], data: None, - left: None, - right: None, + descendants: vec![], } } @@ -277,11 +331,21 @@ impl MongoMerkle { pub fn height() -> usize { return DEPTH; } + + pub fn construct(addr: [u8; 32], root: [u8; 32], db: Option>>) -> Self { + MongoMerkle { + root_hash: root, + default_hash: DEFAULT_HASH_VEC.clone(), + db: db.unwrap_or_else(|| Rc::new(RefCell::new(MongoDB::new(addr, None)))), + } + } + fn empty_leaf(index: u64) -> MerkleRecord { let mut leaf = MerkleRecord::new(index); leaf.set(&[0; 32].to_vec()); leaf } + /// depth start from 0 up to Self::height(). Example 20 height MongoMerkle, root depth=0, leaf depth=20 pub fn get_default_hash(&self, depth: usize) -> Result<[u8; 32], MerkleError> { if depth <= Self::height() { @@ -327,22 +391,16 @@ lazy_static::lazy_static! { } impl MerkleTree<[u8; 32], DEPTH> for MongoMerkle { - type Id = [u8; 32]; - type Root = [u8; 32]; type Node = MerkleRecord; - fn construct(addr: Self::Id, root: Self::Root, db: Option>>) -> Self { - MongoMerkle { - root_hash: root, - default_hash: DEFAULT_HASH_VEC.clone(), - db: db.unwrap_or_else(|| Rc::new(RefCell::new(MongoDB::new(addr, None)))), - } - } - fn get_root_hash(&self) -> [u8; 32] { self.root_hash } + fn chunk_depth() -> usize { + CHUNK_DEPTH + } + fn update_root_hash(&mut self, hash: &[u8; 32]) { self.root_hash = hash.clone(); } @@ -354,104 +412,18 @@ impl MerkleTree<[u8; 32], DEPTH> for MongoMerkle { hasher.update_exact(&[a, b]).to_repr() } - fn set_parent( + fn update_nodes( &mut self, - index: u64, - hash: &[u8; 32], - left: &[u8; 32], - right: &[u8; 32], + nodes: Vec ) -> Result<(), MerkleError> { - self.boundary_check(index)?; - let record = MerkleRecord { - index, - data: None, - left: Some(*left), - right: Some(*right), - hash: *hash, - }; - //println!("set_node_with_hash {} {:?}", index, hash); - self.update_record(record).expect("Unexpected DB Error"); - Ok(()) - } - - fn set_leaf_and_parents( - &mut self, - leaf: &MerkleRecord, - parents: [(u64, [u8; 32], [u8; 32], [u8; 32]); DEPTH], - ) -> Result<(), MerkleError> { - self.leaf_check(leaf.index)?; - let mut records: Vec = parents - .iter() - .filter_map(|(index, hash, left, right)| { - let height = (index + 1).ilog2(); - match self.get_default_hash(height as usize) { - // There is no need to set default nodes in db. - Ok(default_hash) if hash != &default_hash => Some(MerkleRecord { - index: *index, - data: None, - left: Some(*left), - right: Some(*right), - hash: *hash, - }), - _ => None, - } - }) - .collect(); - - records.push(leaf.clone()); - self.update_records(&records) + self.update_records(&nodes) .expect("Unexpected DB Error when update records."); - Ok(()) } fn get_node_with_hash(&self, index: u64, hash: &[u8; 32]) -> Result { self.generate_or_get_node(index, hash) } - - fn set_leaf(&mut self, leaf: &MerkleRecord) -> Result<(), MerkleError> { - self.leaf_check(leaf.index())?; - self.update_record(leaf.clone()) - .expect("Unexpected DB Error"); - Ok(()) - } - - fn get_leaf_with_proof( - &self, - index: u64, - ) -> Result<(Self::Node, MerkleProof<[u8; 32], DEPTH>), MerkleError> { - self.leaf_check(index)?; - let paths = self.get_path(index)?.to_vec(); - // We push the search from the top - let hash = self.get_root_hash(); - let mut acc = 0; - let mut acc_node = self.generate_or_get_node(acc, &hash)?; - let assist: Vec<[u8; 32]> = paths - .into_iter() - .map(|child| { - let (hash, sibling_hash) = if (acc + 1) * 2 == child + 1 { - // left child - (acc_node.left().unwrap(), acc_node.right().unwrap()) - } else { - assert_eq!((acc + 1) * 2, child); - (acc_node.right().unwrap(), acc_node.left().unwrap()) - }; - acc = child; - acc_node = self.generate_or_get_node(acc, &hash)?; - Ok(sibling_hash) - }) - .collect::, _>>()?; - let hash = acc_node.hash(); - Ok(( - acc_node, - MerkleProof { - source: hash, - root: self.get_root_hash(), - assist: assist.try_into().unwrap(), - index, - }, - )) - } } impl MongoMerkle { @@ -539,7 +511,7 @@ mod tests { // Init checking results const DEPTH: usize = 32; const TEST_ADDR: [u8; 32] = [2; 32]; - const INDEX1: u64 = 2_u64.pow(DEPTH as u32) - 1; + const INDEX1: u64 = 2_u64.pow(DEPTH as u32);// - 1; const LEAF1_DATA: [u8; 32] = [ 0, 16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, @@ -582,16 +554,16 @@ mod tests { } #[test] - /* Like the above test but use 13 depth - * 1. Update index=2_u32.pow(13) - 1 (first leaf) leave value. Check root. - * 2. Check index=2_u32.pow(13) - 1 leave value updated. + /* Like the above test but use 16 depth + * 1. Update index=2_u32.pow(16) - 1 (first leaf) leave value. Check root. + * 2. Check index=2_u32.pow(16) - 1 leave value updated. * 3. Load m tree from DB, check root and leave value. */ - fn test_mongo_merkle_single_leaf_update_13_depth() { + fn test_mongo_merkle_single_leaf_update_16_depth() { // Init checking results - const DEPTH: usize = 13; + const DEPTH: usize = 16; const TEST_ADDR: [u8; 32] = [4; 32]; - const INDEX1: u64 = 2_u64.pow(DEPTH as u32) - 1; + const INDEX1: u64 = 2_u64.pow(DEPTH as u32); // - 1; const LEAF1_DATA: [u8; 32] = [ 0, 16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, @@ -633,6 +605,73 @@ mod tests { assert_eq!(mt.verify_proof(&proof).unwrap(), true); } + + fn test_mongo_merkle_multi_leaves_update_16_depth(offset: u64) { + // Init checking results + const DEPTH: usize = 16; + const TEST_ADDR: [u8; 32] = [4; 32]; + let magic_number = offset * offset; + + let mongodb = Rc::new(RefCell::new(MongoDB::new(TEST_ADDR, None))); + + // 1 + let mut mt = MongoMerkle::::construct( + TEST_ADDR, + DEFAULT_HASH_VEC[DEPTH].clone(), + Some(mongodb.clone()), + ); + let cname = get_collection_name(MONGODB_DATA_NAME_PREFIX.to_string(), TEST_ADDR); + let collection = get_collection::( + mongodb.borrow().get_database_client().unwrap(), + MONGODB_DATABASE.to_string(), + cname, + ) + .unwrap(); + let _ = collection.delete_many(doc! {}, None); + + for i in offset..offset + 128 { + let full_index: u64 = 2_u64.pow(DEPTH as u32) - 1 + (i % 2_u64.pow(DEPTH as u32)); + let data_index = (full_index % 32) as usize; + let mut data = [0; 32]; + data[data_index] = (full_index % 256) as u8 + 1; + let (mut leaf, _) = mt.get_leaf_with_proof(full_index).unwrap(); + + leaf.set(&data.to_vec()); + mt.set_leaf_with_proof(&leaf).unwrap(); + + let (leaf, _) = mt.get_leaf_with_proof(full_index).unwrap(); + assert_eq!(leaf.index, full_index); + assert_eq!(leaf.data.unwrap(), data); + } + + // recheck + for i in offset..offset + 128 { + let full_index: u64 = 2_u64.pow(DEPTH as u32) - 1_u64 + (i % 2_u64.pow(DEPTH as u32)); + let data_index = (full_index % 32) as usize; + let mut data = [0; 32]; + data[data_index] = (full_index % 256) as u8 + 1; + let (leaf, proof) = mt.get_leaf_with_proof(full_index).unwrap(); + assert_eq!(leaf.index, full_index); + assert_eq!(leaf.data.unwrap(), data); + let verified = mt.verify_proof(&proof).unwrap(); + assert_eq!(true, verified); + } + } + + #[test] + /* Like the above test but use 16 depth + * 1. Update index=2_u32.pow(16) - 1 (first leaf) leave value. Check root. + * 2. Check index=2_u32.pow(16) - 1 leave value updated. + * 3. Load m tree from DB, check root and leave value. + */ + fn test_mongo_merkle_multi_leaves_16_depth() { + for i in 0..32 { + test_mongo_merkle_multi_leaves_update_16_depth(i+1); + } + } + + + /* Tests for 32 height m tree with updating multple leaves * 1. Clear m tree collection. Create default empty m tree. Check root (default one, A). * 2. Update index=2_u32.pow(32) - 1 (first leaf) leave value. Check root (1 leave updated, B). Check index=2_u32.pow(20) - 1 leave value updated. @@ -754,14 +793,24 @@ mod tests { 0, 0, 0, ]; - // 1 - let mut mt = - MongoMerkle::::construct(TEST_ADDR, DEFAULT_HASH_VEC[DEPTH].clone(), None); + let mongodb = Rc::new(RefCell::new(MongoDB::new(TEST_ADDR, None))); + let mut mt = MongoMerkle::::construct( + TEST_ADDR, + DEFAULT_HASH_VEC[DEPTH], + Some(mongodb.clone()), + ); + let cname = get_collection_name(MONGODB_DATA_NAME_PREFIX.to_string(), TEST_ADDR); - let collection = - get_collection::(MONGODB_DATABASE.to_string(), cname).unwrap(); + let collection = get_collection::( + mongodb.borrow().get_database_client().unwrap(), + MONGODB_DATABASE.to_string(), + cname, + ) + .unwrap(); let _ = collection.delete_many(doc! {}, None); + + // 1 let (mut leaf, _) = mt.get_leaf_with_proof(INDEX1).unwrap(); leaf.set(&LEAF1_DATA.to_vec()); mt.set_leaf_with_proof(&leaf).unwrap();