Merge pull request #118 from ralexstokes/multiproofs-take-2

implement multiproofs

ssz-rs/src/lib.rs

@@ -60,7 +60,7 @@ mod lib {
         pub use std::*;
     }
 
-    pub use self::core::{any, cmp, fmt, iter, slice};
+    pub use self::core::{any, cmp, fmt, iter, mem, slice};
 
     pub use self::{
         cmp::Ordering,
@@ -80,6 +80,12 @@ mod lib {
     #[cfg(feature = "std")]
     pub use std::vec::Vec;
 
+    #[cfg(not(feature = "std"))]
+    pub use alloc::collections::{BTreeMap as HashMap, BTreeSet as HashSet};
+
+    #[cfg(feature = "std")]
+    pub use std::collections::{HashMap, HashSet};
+
     #[cfg(feature = "serde")]
     pub use self::core::marker::PhantomData;
 }
@@ -109,7 +115,9 @@ mod exports {
         de::{Deserialize, DeserializeError},
         error::{Error as SimpleSerializeError, InstanceError, TypeError},
         list::List,
-        merkleization::{is_valid_merkle_branch, MerkleizationError, Merkleized, Node},
+        merkleization::{
+            is_valid_merkle_branch, multiproofs, MerkleizationError, Merkleized, Node,
+        },
         ser::{Serialize, SerializeError},
         uint::U256,
         utils::{deserialize, serialize},

ssz-rs/src/merkleization/mod.rs

@@ -1,3 +1,4 @@
+pub mod multiproofs;
 mod node;
 mod proofs;
 
@@ -33,6 +34,8 @@ pub enum MerkleizationError {
     InputExceedsLimit(usize),
     /// Proof verification failed
     InvalidProof,
+    /// Signals an invalid generalized index (e.g. `0`) was presented.
+    InvalidGeneralizedIndex,
 }
 
 impl From<SerializeError> for MerkleizationError {
@@ -49,6 +52,7 @@ impl Display for MerkleizationError {
             }
             Self::InputExceedsLimit(size) => write!(f, "data exceeds the declared limit {size}"),
             Self::InvalidProof => write!(f, "merkle proof verification failed"),
+            Self::InvalidGeneralizedIndex => write!(f, "invalid generalized index"),
         }
     }
 }

ssz-rs/src/merkleization/multiproofs/generalized_index.rs

@@ -0,0 +1,53 @@
+use crate::{lib::mem, merkleization::MerkleizationError as Error};
+
+const BITS_PER_BYTE: usize = crate::BITS_PER_BYTE as usize;
+
+// From: https://users.rust-lang.org/t/logarithm-of-integers/8506/5
+const fn num_bits<T>() -> usize {
+    mem::size_of::<T>() * BITS_PER_BYTE
+}
+
+// Return base 2 logarithm of `x`.
+// `None` is returned if `x` is `0` as this logarithm is undefined.
+fn log_2(x: usize) -> Option<u32> {
+    if x == 0 {
+        None
+    } else {
+        Some(num_bits::<usize>() as u32 - x.leading_zeros() - 1)
+    }
+}
+
+pub fn get_power_of_two_ceil(x: usize) -> usize {
+    match x {
+        x if x <= 1 => 1,
+        2 => 2,
+        x => 2 * get_power_of_two_ceil((x + 1) / 2),
+    }
+}
+
+pub type GeneralizedIndex = usize;
+
+pub fn get_path_length(index: GeneralizedIndex) -> Result<usize, Error> {
+    let length = log_2(index).ok_or(Error::InvalidGeneralizedIndex)?;
+    Ok(length as usize)
+}
+
+pub fn get_bit(index: GeneralizedIndex, position: usize) -> bool {
+    index & (1 << position) > 0
+}
+
+pub fn sibling(index: GeneralizedIndex) -> GeneralizedIndex {
+    index ^ 1
+}
+
+pub fn child_left(index: GeneralizedIndex) -> GeneralizedIndex {
+    index * 2
+}
+
+pub fn child_right(index: GeneralizedIndex) -> GeneralizedIndex {
+    index * 2 + 1
+}
+
+pub fn parent(index: GeneralizedIndex) -> GeneralizedIndex {
+    index / 2
+}

ssz-rs/src/merkleization/multiproofs/mod.rs

@@ -0,0 +1,149 @@
+mod generalized_index;
+
+pub use generalized_index::*;
+
+use crate::{
+    lib::*,
+    merkleization::{MerkleizationError as Error, Node},
+};
+use sha2::{Digest, Sha256};
+
+fn get_branch_indices(tree_index: GeneralizedIndex) -> Vec<GeneralizedIndex> {
+    let mut focus = sibling(tree_index);
+    let mut result = vec![focus];
+    while focus > 1 {
+        focus = sibling(parent(focus));
+        result.push(focus);
+    }
+    result.truncate(result.len() - 1);
+    result
+}
+
+fn get_path_indices(tree_index: GeneralizedIndex) -> Vec<GeneralizedIndex> {
+    let mut focus = tree_index;
+    let mut result = vec![focus];
+    while focus > 1 {
+        focus = parent(focus);
+        result.push(focus);
+    }
+    result.truncate(result.len() - 1);
+    result
+}
+
+fn get_helper_indices(indices: &[GeneralizedIndex]) -> Vec<GeneralizedIndex> {
+    let mut all_helper_indices = HashSet::new();
+    let mut all_path_indices = HashSet::new();
+
+    for index in indices {
+        all_helper_indices.extend(get_branch_indices(*index).iter());
+        all_path_indices.extend(get_path_indices(*index).iter());
+    }
+
+    let mut all_branch_indices =
+        all_helper_indices.difference(&all_path_indices).cloned().collect::<Vec<_>>();
+    all_branch_indices.sort_by(|a: &GeneralizedIndex, b: &GeneralizedIndex| b.cmp(a));
+    all_branch_indices
+}
+
+pub fn calculate_merkle_root(
+    leaf: Node,
+    proof: &[Node],
+    index: GeneralizedIndex,
+) -> Result<Node, Error> {
+    let path_length = get_path_length(index)?;
+    if path_length != proof.len() {
+        return Err(Error::InvalidProof)
+    }
+    let mut result = leaf;
+
+    let mut hasher = Sha256::new();
+    for (i, next) in proof.iter().enumerate() {
+        if get_bit(index, i) {
+            hasher.update(next.as_ref());
+            hasher.update(result.as_ref());
+        } else {
+            hasher.update(result.as_ref());
+            hasher.update(next.as_ref());
+        }
+        result.as_mut().copy_from_slice(&hasher.finalize_reset());
+    }
+    Ok(result)
+}
+
+pub fn verify_merkle_proof(
+    leaf: Node,
+    proof: &[Node],
+    index: GeneralizedIndex,
+    root: Node,
+) -> Result<(), Error> {
+    if calculate_merkle_root(leaf, proof, index)? == root {
+        Ok(())
+    } else {
+        Err(Error::InvalidProof)
+    }
+}
+
+pub fn calculate_multi_merkle_root(
+    leaves: &[Node],
+    proof: &[Node],
+    indices: &[GeneralizedIndex],
+) -> Result<Node, Error> {
+    if leaves.len() != indices.len() {
+        return Err(Error::InvalidProof)
+    }
+    let helper_indices = get_helper_indices(indices);
+    if proof.len() != helper_indices.len() {
+        return Err(Error::InvalidProof)
+    }
+
+    let mut objects = HashMap::new();
+    for (index, node) in indices.iter().zip(leaves.iter()) {
+        objects.insert(*index, *node);
+    }
+    for (index, node) in helper_indices.iter().zip(proof.iter()) {
+        objects.insert(*index, *node);
+    }
+
+    let mut keys = objects.keys().cloned().collect::<Vec<_>>();
+    keys.sort_by(|a, b| b.cmp(a));
+
+    let mut hasher = Sha256::new();
+    let mut pos = 0;
+    while pos < keys.len() {
+        let key = keys.get(pos).unwrap();
+        let key_present = objects.contains_key(key);
+        let sibling_present = objects.contains_key(&sibling(*key));
+        let parent_index = parent(*key);
+        let parent_missing = !objects.contains_key(&parent_index);
+        let should_compute = key_present && sibling_present && parent_missing;
+        if should_compute {
+            let right_index = key | 1;
+            let left_index = sibling(right_index);
+            let left_input = objects.get(&left_index).expect("contains index");
+            let right_input = objects.get(&right_index).expect("contains index");
+            hasher.update(left_input.as_ref());
+            hasher.update(right_input.as_ref());
+
+            let parent = objects.entry(parent_index).or_default();
+            parent.as_mut().copy_from_slice(&hasher.finalize_reset());
+            keys.push(parent_index);
+        }
+        pos += 1;
+    }
+
+    let root = *objects.get(&1).expect("contains index");
+    Ok(root)
+}
+
+pub fn verify_merkle_multiproof(
+    leaves: &[Node],
+    proof: &[Node],
+    indices: &[GeneralizedIndex],
+    root: Node,
+) -> Result<(), Error> {
+    if calculate_multi_merkle_root(leaves, proof, indices)? == root {
+        Ok(())
+    } else {
+        Err(Error::InvalidProof)
+    }
+}