This commit moves the previous implementation of `SparseMerkleTree` from miden-core to this crate. It also include a couple of new tests, a bench suite, and a couple of minor fixes. The original API was preserved to maintain compatibility with `AdviceTape`. closes #21al-gkr-basic-workflow
@ -0,0 +1,84 @@ | 
				
			|||
use core::mem::swap;
 | 
				
			|||
use criterion::{black_box, criterion_group, criterion_main, Criterion};
 | 
				
			|||
use miden_crypto::{merkle::SimpleSmt, Felt, Word};
 | 
				
			|||
use rand_utils::prng_array;
 | 
				
			|||
 | 
				
			|||
fn smt_rpo(c: &mut Criterion) {
 | 
				
			|||
    // setup trees
 | 
				
			|||
 | 
				
			|||
    let mut seed = [0u8; 32];
 | 
				
			|||
    let mut trees = vec![];
 | 
				
			|||
 | 
				
			|||
    for depth in 14..=20 {
 | 
				
			|||
        let leaves = ((1 << depth) - 1) as u64;
 | 
				
			|||
        for count in [1, leaves / 2, leaves] {
 | 
				
			|||
            let entries: Vec<_> = (0..count)
 | 
				
			|||
                .map(|i| {
 | 
				
			|||
                    let word = generate_word(&mut seed);
 | 
				
			|||
                    (i, word)
 | 
				
			|||
                })
 | 
				
			|||
                .collect();
 | 
				
			|||
            let tree = SimpleSmt::new(entries, depth).unwrap();
 | 
				
			|||
            trees.push(tree);
 | 
				
			|||
        }
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    let leaf = generate_word(&mut seed);
 | 
				
			|||
 | 
				
			|||
    // benchmarks
 | 
				
			|||
 | 
				
			|||
    let mut insert = c.benchmark_group(format!("smt update_leaf"));
 | 
				
			|||
 | 
				
			|||
    for tree in trees.iter_mut() {
 | 
				
			|||
        let depth = tree.depth();
 | 
				
			|||
        let count = tree.leaves_count() as u64;
 | 
				
			|||
        let key = count >> 2;
 | 
				
			|||
        insert.bench_with_input(
 | 
				
			|||
            format!("simple smt(depth:{depth},count:{count})"),
 | 
				
			|||
            &(key, leaf),
 | 
				
			|||
            |b, (key, leaf)| {
 | 
				
			|||
                b.iter(|| {
 | 
				
			|||
                    tree.update_leaf(black_box(*key), black_box(*leaf)).unwrap();
 | 
				
			|||
                });
 | 
				
			|||
            },
 | 
				
			|||
        );
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    insert.finish();
 | 
				
			|||
 | 
				
			|||
    let mut path = c.benchmark_group(format!("smt get_leaf_path"));
 | 
				
			|||
 | 
				
			|||
    for tree in trees.iter_mut() {
 | 
				
			|||
        let depth = tree.depth();
 | 
				
			|||
        let count = tree.leaves_count() as u64;
 | 
				
			|||
        let key = count >> 2;
 | 
				
			|||
        path.bench_with_input(
 | 
				
			|||
            format!("simple smt(depth:{depth},count:{count})"),
 | 
				
			|||
            &key,
 | 
				
			|||
            |b, key| {
 | 
				
			|||
                b.iter(|| {
 | 
				
			|||
                    tree.get_leaf_path(black_box(*key)).unwrap();
 | 
				
			|||
                });
 | 
				
			|||
            },
 | 
				
			|||
        );
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    path.finish();
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
criterion_group!(smt_group, smt_rpo);
 | 
				
			|||
criterion_main!(smt_group);
 | 
				
			|||
 | 
				
			|||
// HELPER FUNCTIONS
 | 
				
			|||
// --------------------------------------------------------------------------------------------
 | 
				
			|||
 | 
				
			|||
fn generate_word(seed: &mut [u8; 32]) -> Word {
 | 
				
			|||
    swap(seed, &mut prng_array(*seed));
 | 
				
			|||
    let nums: [u64; 4] = prng_array(*seed);
 | 
				
			|||
    [
 | 
				
			|||
        Felt::new(nums[0]),
 | 
				
			|||
        Felt::new(nums[1]),
 | 
				
			|||
        Felt::new(nums[2]),
 | 
				
			|||
        Felt::new(nums[3]),
 | 
				
			|||
    ]
 | 
				
			|||
}
 | 
				
			|||
@ -0,0 +1,269 @@ | 
				
			|||
use super::{BTreeMap, MerkleError, Rpo256, RpoDigest, Vec, Word};
 | 
				
			|||
 | 
				
			|||
#[cfg(test)]
 | 
				
			|||
mod tests;
 | 
				
			|||
 | 
				
			|||
// SPARSE MERKLE TREE
 | 
				
			|||
// ================================================================================================
 | 
				
			|||
 | 
				
			|||
/// A sparse Merkle tree with 63-bit keys and 4-element leaf values, without compaction.
 | 
				
			|||
/// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`.
 | 
				
			|||
/// The root of the tree is recomputed on each new leaf update.
 | 
				
			|||
#[derive(Clone, Debug)]
 | 
				
			|||
pub struct SimpleSmt {
 | 
				
			|||
    root: Word,
 | 
				
			|||
    depth: u32,
 | 
				
			|||
    store: Store,
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
impl SimpleSmt {
 | 
				
			|||
    // CONSTANTS
 | 
				
			|||
    // --------------------------------------------------------------------------------------------
 | 
				
			|||
 | 
				
			|||
    /// Minimum supported depth.
 | 
				
			|||
    pub const MIN_DEPTH: u32 = 1;
 | 
				
			|||
 | 
				
			|||
    /// Maximum supported depth.
 | 
				
			|||
    pub const MAX_DEPTH: u32 = 63;
 | 
				
			|||
 | 
				
			|||
    // CONSTRUCTORS
 | 
				
			|||
    // --------------------------------------------------------------------------------------------
 | 
				
			|||
 | 
				
			|||
    /// Creates a new simple SMT.
 | 
				
			|||
    ///
 | 
				
			|||
    /// The provided entries will be tuples of the leaves and their corresponding keys.
 | 
				
			|||
    ///
 | 
				
			|||
    /// # Errors
 | 
				
			|||
    ///
 | 
				
			|||
    /// The function will fail if the provided entries count exceed the maximum tree capacity, that
 | 
				
			|||
    /// is `2^{depth}`.
 | 
				
			|||
    pub fn new<R, I>(entries: R, depth: u32) -> Result<Self, MerkleError>
 | 
				
			|||
    where
 | 
				
			|||
        R: IntoIterator<IntoIter = I>,
 | 
				
			|||
        I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
 | 
				
			|||
    {
 | 
				
			|||
        let mut entries = entries.into_iter();
 | 
				
			|||
 | 
				
			|||
        // validate the range of the depth.
 | 
				
			|||
        let max = 1 << depth;
 | 
				
			|||
        if depth < Self::MIN_DEPTH {
 | 
				
			|||
            return Err(MerkleError::DepthTooSmall(depth));
 | 
				
			|||
        } else if Self::MAX_DEPTH < depth {
 | 
				
			|||
            return Err(MerkleError::DepthTooBig(depth));
 | 
				
			|||
        } else if entries.len() > max {
 | 
				
			|||
            return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
 | 
				
			|||
        }
 | 
				
			|||
 | 
				
			|||
        let (store, root) = Store::new(depth);
 | 
				
			|||
        let mut tree = Self { root, depth, store };
 | 
				
			|||
        entries.try_for_each(|(key, leaf)| tree.insert_leaf(key, leaf))?;
 | 
				
			|||
 | 
				
			|||
        Ok(tree)
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Returns the root of this Merkle tree.
 | 
				
			|||
    pub const fn root(&self) -> Word {
 | 
				
			|||
        self.root
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Returns the depth of this Merkle tree.
 | 
				
			|||
    pub const fn depth(&self) -> u32 {
 | 
				
			|||
        self.depth
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Returns the set count of the keys of the leaves.
 | 
				
			|||
    pub fn leaves_count(&self) -> usize {
 | 
				
			|||
        self.store.leaves_count()
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Returns a node at the specified key
 | 
				
			|||
    ///
 | 
				
			|||
    /// # Errors
 | 
				
			|||
    /// Returns an error if:
 | 
				
			|||
    /// * The specified depth is greater than the depth of the tree.
 | 
				
			|||
    /// * The specified key does not exist
 | 
				
			|||
    pub fn get_node(&self, depth: u32, key: u64) -> Result<Word, MerkleError> {
 | 
				
			|||
        if depth == 0 {
 | 
				
			|||
            Err(MerkleError::DepthTooSmall(depth))
 | 
				
			|||
        } else if depth > self.depth() {
 | 
				
			|||
            Err(MerkleError::DepthTooBig(depth))
 | 
				
			|||
        } else if depth == self.depth() {
 | 
				
			|||
            self.store.get_leaf_node(key)
 | 
				
			|||
        } else {
 | 
				
			|||
            let branch_node = self.store.get_branch_node(key, depth)?;
 | 
				
			|||
            Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
 | 
				
			|||
        }
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Returns a Merkle path from the node at the specified key to the root. The node itself is
 | 
				
			|||
    /// not included in the path.
 | 
				
			|||
    ///
 | 
				
			|||
    /// # Errors
 | 
				
			|||
    /// Returns an error if:
 | 
				
			|||
    /// * The specified key does not exist as a branch or leaf node
 | 
				
			|||
    /// * The specified depth is greater than the depth of the tree.
 | 
				
			|||
    pub fn get_path(&self, depth: u32, key: u64) -> Result<Vec<Word>, MerkleError> {
 | 
				
			|||
        if depth == 0 {
 | 
				
			|||
            return Err(MerkleError::DepthTooSmall(depth));
 | 
				
			|||
        } else if depth > self.depth() {
 | 
				
			|||
            return Err(MerkleError::DepthTooBig(depth));
 | 
				
			|||
        } else if depth == self.depth() && !self.store.check_leaf_node_exists(key) {
 | 
				
			|||
            return Err(MerkleError::InvalidIndex(self.depth(), key));
 | 
				
			|||
        }
 | 
				
			|||
 | 
				
			|||
        let mut path = Vec::with_capacity(depth as usize);
 | 
				
			|||
        let mut curr_key = key;
 | 
				
			|||
        for n in (0..depth).rev() {
 | 
				
			|||
            let parent_key = curr_key >> 1;
 | 
				
			|||
            let parent_node = self.store.get_branch_node(parent_key, n)?;
 | 
				
			|||
            let sibling_node = if curr_key & 1 == 1 {
 | 
				
			|||
                parent_node.left
 | 
				
			|||
            } else {
 | 
				
			|||
                parent_node.right
 | 
				
			|||
            };
 | 
				
			|||
            path.push(sibling_node.into());
 | 
				
			|||
            curr_key >>= 1;
 | 
				
			|||
        }
 | 
				
			|||
        Ok(path)
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
 | 
				
			|||
    /// included in the path.
 | 
				
			|||
    ///
 | 
				
			|||
    /// # Errors
 | 
				
			|||
    /// Returns an error if:
 | 
				
			|||
    /// * The specified key does not exist as a leaf node.
 | 
				
			|||
    pub fn get_leaf_path(&self, key: u64) -> Result<Vec<Word>, MerkleError> {
 | 
				
			|||
        self.get_path(self.depth(), key)
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
 | 
				
			|||
    ///
 | 
				
			|||
    /// # Errors
 | 
				
			|||
    /// Returns an error if the specified key is not a valid leaf index for this tree.
 | 
				
			|||
    pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
 | 
				
			|||
        if !self.store.check_leaf_node_exists(key) {
 | 
				
			|||
            return Err(MerkleError::InvalidIndex(self.depth(), key));
 | 
				
			|||
        }
 | 
				
			|||
        self.insert_leaf(key, value)?;
 | 
				
			|||
 | 
				
			|||
        Ok(())
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    /// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree
 | 
				
			|||
    pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
 | 
				
			|||
        self.store.insert_leaf_node(key, value);
 | 
				
			|||
 | 
				
			|||
        let depth = self.depth();
 | 
				
			|||
        let mut curr_key = key;
 | 
				
			|||
        let mut curr_node: RpoDigest = value.into();
 | 
				
			|||
        for n in (0..depth).rev() {
 | 
				
			|||
            let parent_key = curr_key >> 1;
 | 
				
			|||
            let parent_node = self
 | 
				
			|||
                .store
 | 
				
			|||
                .get_branch_node(parent_key, n)
 | 
				
			|||
                .unwrap_or_else(|_| self.store.get_empty_node((n + 1) as usize));
 | 
				
			|||
            let (left, right) = if curr_key & 1 == 1 {
 | 
				
			|||
                (parent_node.left, curr_node)
 | 
				
			|||
            } else {
 | 
				
			|||
                (curr_node, parent_node.right)
 | 
				
			|||
            };
 | 
				
			|||
 | 
				
			|||
            self.store.insert_branch_node(parent_key, n, left, right);
 | 
				
			|||
            curr_key = parent_key;
 | 
				
			|||
            curr_node = Rpo256::merge(&[left, right]);
 | 
				
			|||
        }
 | 
				
			|||
        self.root = curr_node.into();
 | 
				
			|||
 | 
				
			|||
        Ok(())
 | 
				
			|||
    }
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
// STORE
 | 
				
			|||
// ================================================================================================
 | 
				
			|||
 | 
				
			|||
/// A data store for sparse Merkle tree key-value pairs.
 | 
				
			|||
/// Leaves and branch nodes are stored separately in B-tree maps, indexed by key and (key, depth)
 | 
				
			|||
/// respectively. Hashes for blank subtrees at each layer are stored in `empty_hashes`, beginning
 | 
				
			|||
/// with the root hash of an empty tree, and ending with the zero value of a leaf node.
 | 
				
			|||
#[derive(Clone, Debug)]
 | 
				
			|||
struct Store {
 | 
				
			|||
    branches: BTreeMap<(u64, u32), BranchNode>,
 | 
				
			|||
    leaves: BTreeMap<u64, Word>,
 | 
				
			|||
    empty_hashes: Vec<RpoDigest>,
 | 
				
			|||
    depth: u32,
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[derive(Clone, Debug, Default)]
 | 
				
			|||
struct BranchNode {
 | 
				
			|||
    left: RpoDigest,
 | 
				
			|||
    right: RpoDigest,
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
impl Store {
 | 
				
			|||
    fn new(depth: u32) -> (Self, Word) {
 | 
				
			|||
        let branches = BTreeMap::new();
 | 
				
			|||
        let leaves = BTreeMap::new();
 | 
				
			|||
 | 
				
			|||
        // Construct empty node digests for each layer of the tree
 | 
				
			|||
        let empty_hashes: Vec<RpoDigest> = (0..depth + 1)
 | 
				
			|||
            .scan(Word::default().into(), |state, _| {
 | 
				
			|||
                let value = *state;
 | 
				
			|||
                *state = Rpo256::merge(&[value, value]);
 | 
				
			|||
                Some(value)
 | 
				
			|||
            })
 | 
				
			|||
            .collect::<Vec<_>>()
 | 
				
			|||
            .into_iter()
 | 
				
			|||
            .rev()
 | 
				
			|||
            .collect();
 | 
				
			|||
 | 
				
			|||
        let root = empty_hashes[0].into();
 | 
				
			|||
        let store = Self {
 | 
				
			|||
            branches,
 | 
				
			|||
            leaves,
 | 
				
			|||
            empty_hashes,
 | 
				
			|||
            depth,
 | 
				
			|||
        };
 | 
				
			|||
 | 
				
			|||
        (store, root)
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn get_empty_node(&self, depth: usize) -> BranchNode {
 | 
				
			|||
        let digest = self.empty_hashes[depth];
 | 
				
			|||
        BranchNode {
 | 
				
			|||
            left: digest,
 | 
				
			|||
            right: digest,
 | 
				
			|||
        }
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn check_leaf_node_exists(&self, key: u64) -> bool {
 | 
				
			|||
        self.leaves.contains_key(&key)
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn get_leaf_node(&self, key: u64) -> Result<Word, MerkleError> {
 | 
				
			|||
        self.leaves
 | 
				
			|||
            .get(&key)
 | 
				
			|||
            .cloned()
 | 
				
			|||
            .ok_or(MerkleError::InvalidIndex(self.depth, key))
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn insert_leaf_node(&mut self, key: u64, node: Word) {
 | 
				
			|||
        self.leaves.insert(key, node);
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn get_branch_node(&self, key: u64, depth: u32) -> Result<BranchNode, MerkleError> {
 | 
				
			|||
        self.branches
 | 
				
			|||
            .get(&(key, depth))
 | 
				
			|||
            .cloned()
 | 
				
			|||
            .ok_or(MerkleError::InvalidIndex(depth, key))
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn insert_branch_node(&mut self, key: u64, depth: u32, left: RpoDigest, right: RpoDigest) {
 | 
				
			|||
        let node = BranchNode { left, right };
 | 
				
			|||
        self.branches.insert((key, depth), node);
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    fn leaves_count(&self) -> usize {
 | 
				
			|||
        self.leaves.len()
 | 
				
			|||
    }
 | 
				
			|||
}
 | 
				
			|||
@ -0,0 +1,263 @@ | 
				
			|||
use super::{
 | 
				
			|||
    super::{MerkleTree, RpoDigest, SimpleSmt},
 | 
				
			|||
    Rpo256, Vec, Word,
 | 
				
			|||
};
 | 
				
			|||
use crate::{Felt, FieldElement};
 | 
				
			|||
use core::iter;
 | 
				
			|||
use proptest::prelude::*;
 | 
				
			|||
use rand_utils::prng_array;
 | 
				
			|||
 | 
				
			|||
const KEYS4: [u64; 4] = [0, 1, 2, 3];
 | 
				
			|||
const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
 | 
				
			|||
 | 
				
			|||
const VALUES4: [Word; 4] = [
 | 
				
			|||
    int_to_node(1),
 | 
				
			|||
    int_to_node(2),
 | 
				
			|||
    int_to_node(3),
 | 
				
			|||
    int_to_node(4),
 | 
				
			|||
];
 | 
				
			|||
 | 
				
			|||
const VALUES8: [Word; 8] = [
 | 
				
			|||
    int_to_node(1),
 | 
				
			|||
    int_to_node(2),
 | 
				
			|||
    int_to_node(3),
 | 
				
			|||
    int_to_node(4),
 | 
				
			|||
    int_to_node(5),
 | 
				
			|||
    int_to_node(6),
 | 
				
			|||
    int_to_node(7),
 | 
				
			|||
    int_to_node(8),
 | 
				
			|||
];
 | 
				
			|||
 | 
				
			|||
const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8];
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn build_empty_tree() {
 | 
				
			|||
    let smt = SimpleSmt::new(iter::empty(), 3).unwrap();
 | 
				
			|||
    let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap();
 | 
				
			|||
    assert_eq!(mt.root(), smt.root());
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn empty_digests_are_consistent() {
 | 
				
			|||
    let depth = 5;
 | 
				
			|||
    let root = SimpleSmt::new(iter::empty(), depth).unwrap().root();
 | 
				
			|||
    let computed: [RpoDigest; 2] = (0..depth).fold([Default::default(); 2], |state, _| {
 | 
				
			|||
        let digest = Rpo256::merge(&state);
 | 
				
			|||
        [digest; 2]
 | 
				
			|||
    });
 | 
				
			|||
 | 
				
			|||
    assert_eq!(Word::from(computed[0]), root);
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn build_sparse_tree() {
 | 
				
			|||
    let mut smt = SimpleSmt::new(iter::empty(), 3).unwrap();
 | 
				
			|||
    let mut values = ZERO_VALUES8.to_vec();
 | 
				
			|||
 | 
				
			|||
    // insert single value
 | 
				
			|||
    let key = 6;
 | 
				
			|||
    let new_node = int_to_node(7);
 | 
				
			|||
    values[key as usize] = new_node;
 | 
				
			|||
    smt.insert_leaf(key, new_node)
 | 
				
			|||
        .expect("Failed to insert leaf");
 | 
				
			|||
    let mt2 = MerkleTree::new(values.clone()).unwrap();
 | 
				
			|||
    assert_eq!(mt2.root(), smt.root());
 | 
				
			|||
    assert_eq!(mt2.get_path(3, 6).unwrap(), smt.get_path(3, 6).unwrap());
 | 
				
			|||
 | 
				
			|||
    // insert second value at distinct leaf branch
 | 
				
			|||
    let key = 2;
 | 
				
			|||
    let new_node = int_to_node(3);
 | 
				
			|||
    values[key as usize] = new_node;
 | 
				
			|||
    smt.insert_leaf(key, new_node)
 | 
				
			|||
        .expect("Failed to insert leaf");
 | 
				
			|||
    let mt3 = MerkleTree::new(values).unwrap();
 | 
				
			|||
    assert_eq!(mt3.root(), smt.root());
 | 
				
			|||
    assert_eq!(mt3.get_path(3, 2).unwrap(), smt.get_path(3, 2).unwrap());
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn build_full_tree() {
 | 
				
			|||
    let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
 | 
				
			|||
 | 
				
			|||
    let (root, node2, node3) = compute_internal_nodes();
 | 
				
			|||
    assert_eq!(root, tree.root());
 | 
				
			|||
    assert_eq!(node2, tree.get_node(1, 0).unwrap());
 | 
				
			|||
    assert_eq!(node3, tree.get_node(1, 1).unwrap());
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn get_values() {
 | 
				
			|||
    let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
 | 
				
			|||
 | 
				
			|||
    // check depth 2
 | 
				
			|||
    assert_eq!(VALUES4[0], tree.get_node(2, 0).unwrap());
 | 
				
			|||
    assert_eq!(VALUES4[1], tree.get_node(2, 1).unwrap());
 | 
				
			|||
    assert_eq!(VALUES4[2], tree.get_node(2, 2).unwrap());
 | 
				
			|||
    assert_eq!(VALUES4[3], tree.get_node(2, 3).unwrap());
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn get_path() {
 | 
				
			|||
    let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
 | 
				
			|||
 | 
				
			|||
    let (_, node2, node3) = compute_internal_nodes();
 | 
				
			|||
 | 
				
			|||
    // check depth 2
 | 
				
			|||
    assert_eq!(vec![VALUES4[1], node3], tree.get_path(2, 0).unwrap());
 | 
				
			|||
    assert_eq!(vec![VALUES4[0], node3], tree.get_path(2, 1).unwrap());
 | 
				
			|||
    assert_eq!(vec![VALUES4[3], node2], tree.get_path(2, 2).unwrap());
 | 
				
			|||
    assert_eq!(vec![VALUES4[2], node2], tree.get_path(2, 3).unwrap());
 | 
				
			|||
 | 
				
			|||
    // check depth 1
 | 
				
			|||
    assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
 | 
				
			|||
    assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn update_leaf() {
 | 
				
			|||
    let mut tree = SimpleSmt::new(KEYS8.into_iter().zip(VALUES8.into_iter()), 3).unwrap();
 | 
				
			|||
 | 
				
			|||
    // update one value
 | 
				
			|||
    let key = 3;
 | 
				
			|||
    let new_node = int_to_node(9);
 | 
				
			|||
    let mut expected_values = VALUES8.to_vec();
 | 
				
			|||
    expected_values[key] = new_node;
 | 
				
			|||
    let expected_tree = SimpleSmt::new(
 | 
				
			|||
        KEYS8.into_iter().zip(expected_values.clone().into_iter()),
 | 
				
			|||
        3,
 | 
				
			|||
    )
 | 
				
			|||
    .unwrap();
 | 
				
			|||
 | 
				
			|||
    tree.update_leaf(key as u64, new_node).unwrap();
 | 
				
			|||
    assert_eq!(expected_tree.root, tree.root);
 | 
				
			|||
 | 
				
			|||
    // update another value
 | 
				
			|||
    let key = 6;
 | 
				
			|||
    let new_node = int_to_node(10);
 | 
				
			|||
    expected_values[key] = new_node;
 | 
				
			|||
    let expected_tree =
 | 
				
			|||
        SimpleSmt::new(KEYS8.into_iter().zip(expected_values.into_iter()), 3).unwrap();
 | 
				
			|||
 | 
				
			|||
    tree.update_leaf(key as u64, new_node).unwrap();
 | 
				
			|||
    assert_eq!(expected_tree.root, tree.root);
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
#[test]
 | 
				
			|||
fn small_tree_opening_is_consistent() {
 | 
				
			|||
    //        ____k____
 | 
				
			|||
    //       /         \
 | 
				
			|||
    //     _i_         _j_
 | 
				
			|||
    //    /   \       /   \
 | 
				
			|||
    //   e     f     g     h
 | 
				
			|||
    //  / \   / \   / \   / \
 | 
				
			|||
    // a   b 0   0 c   0 0   d
 | 
				
			|||
 | 
				
			|||
    let z = Word::from(RpoDigest::default());
 | 
				
			|||
 | 
				
			|||
    let a = Word::from(Rpo256::merge(&[z.into(); 2]));
 | 
				
			|||
    let b = Word::from(Rpo256::merge(&[a.into(); 2]));
 | 
				
			|||
    let c = Word::from(Rpo256::merge(&[b.into(); 2]));
 | 
				
			|||
    let d = Word::from(Rpo256::merge(&[c.into(); 2]));
 | 
				
			|||
 | 
				
			|||
    let e = Word::from(Rpo256::merge(&[a.into(), b.into()]));
 | 
				
			|||
    let f = Word::from(Rpo256::merge(&[z.into(), z.into()]));
 | 
				
			|||
    let g = Word::from(Rpo256::merge(&[c.into(), z.into()]));
 | 
				
			|||
    let h = Word::from(Rpo256::merge(&[z.into(), d.into()]));
 | 
				
			|||
 | 
				
			|||
    let i = Word::from(Rpo256::merge(&[e.into(), f.into()]));
 | 
				
			|||
    let j = Word::from(Rpo256::merge(&[g.into(), h.into()]));
 | 
				
			|||
 | 
				
			|||
    let k = Word::from(Rpo256::merge(&[i.into(), j.into()]));
 | 
				
			|||
 | 
				
			|||
    let depth = 3;
 | 
				
			|||
    let entries = vec![(0, a), (1, b), (4, c), (7, d)];
 | 
				
			|||
    let tree = SimpleSmt::new(entries, depth).unwrap();
 | 
				
			|||
 | 
				
			|||
    assert_eq!(tree.root(), Word::from(k));
 | 
				
			|||
 | 
				
			|||
    let cases: Vec<(u32, u64, Vec<Word>)> = vec![
 | 
				
			|||
        (3, 0, vec![b, f, j]),
 | 
				
			|||
        (3, 1, vec![a, f, j]),
 | 
				
			|||
        (3, 4, vec![z, h, i]),
 | 
				
			|||
        (3, 7, vec![z, g, i]),
 | 
				
			|||
        (2, 0, vec![f, j]),
 | 
				
			|||
        (2, 1, vec![e, j]),
 | 
				
			|||
        (2, 2, vec![h, i]),
 | 
				
			|||
        (2, 3, vec![g, i]),
 | 
				
			|||
        (1, 0, vec![j]),
 | 
				
			|||
        (1, 1, vec![i]),
 | 
				
			|||
    ];
 | 
				
			|||
 | 
				
			|||
    for (depth, key, path) in cases {
 | 
				
			|||
        let opening = tree.get_path(depth, key).unwrap();
 | 
				
			|||
 | 
				
			|||
        assert_eq!(path, opening);
 | 
				
			|||
    }
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
proptest! {
 | 
				
			|||
    #[test]
 | 
				
			|||
    fn arbitrary_openings_single_leaf(
 | 
				
			|||
        depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
 | 
				
			|||
        key in prop::num::u64::ANY,
 | 
				
			|||
        leaf in prop::num::u64::ANY,
 | 
				
			|||
    ) {
 | 
				
			|||
        let mut tree = SimpleSmt::new(iter::empty(), depth).unwrap();
 | 
				
			|||
 | 
				
			|||
        let key = key % (1 << depth as u64);
 | 
				
			|||
        let leaf = int_to_node(leaf);
 | 
				
			|||
 | 
				
			|||
        tree.insert_leaf(key, leaf.into()).unwrap();
 | 
				
			|||
        tree.get_leaf_path(key).unwrap();
 | 
				
			|||
 | 
				
			|||
        // traverse to root, fetching all paths
 | 
				
			|||
        for d in 1..depth {
 | 
				
			|||
            let k = key >> (depth - d);
 | 
				
			|||
            tree.get_path(d, k).unwrap();
 | 
				
			|||
        }
 | 
				
			|||
    }
 | 
				
			|||
 | 
				
			|||
    #[test]
 | 
				
			|||
    fn arbitrary_openings_multiple_leaves(
 | 
				
			|||
        depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
 | 
				
			|||
        count in 2u8..10u8,
 | 
				
			|||
        ref seed in any::<[u8; 32]>()
 | 
				
			|||
    ) {
 | 
				
			|||
        let mut tree = SimpleSmt::new(iter::empty(), depth).unwrap();
 | 
				
			|||
        let mut seed = *seed;
 | 
				
			|||
        let leaves = (1 << depth) - 1;
 | 
				
			|||
 | 
				
			|||
        for _ in 0..count {
 | 
				
			|||
            seed = prng_array(seed);
 | 
				
			|||
 | 
				
			|||
            let mut key = [0u8; 8];
 | 
				
			|||
            let mut leaf = [0u8; 8];
 | 
				
			|||
 | 
				
			|||
            key.copy_from_slice(&seed[..8]);
 | 
				
			|||
            leaf.copy_from_slice(&seed[8..16]);
 | 
				
			|||
 | 
				
			|||
            let key = u64::from_le_bytes(key);
 | 
				
			|||
            let key = key % leaves;
 | 
				
			|||
            let leaf = u64::from_le_bytes(leaf);
 | 
				
			|||
            let leaf = int_to_node(leaf);
 | 
				
			|||
 | 
				
			|||
            tree.insert_leaf(key, leaf).unwrap();
 | 
				
			|||
            tree.get_leaf_path(key).unwrap();
 | 
				
			|||
        }
 | 
				
			|||
    }
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
// HELPER FUNCTIONS
 | 
				
			|||
// --------------------------------------------------------------------------------------------
 | 
				
			|||
 | 
				
			|||
fn compute_internal_nodes() -> (Word, Word, Word) {
 | 
				
			|||
    let node2 = Rpo256::hash_elements(&[VALUES4[0], VALUES4[1]].concat());
 | 
				
			|||
    let node3 = Rpo256::hash_elements(&[VALUES4[2], VALUES4[3]].concat());
 | 
				
			|||
    let root = Rpo256::merge(&[node2, node3]);
 | 
				
			|||
 | 
				
			|||
    (root.into(), node2.into(), node3.into())
 | 
				
			|||
}
 | 
				
			|||
 | 
				
			|||
const fn int_to_node(value: u64) -> Word {
 | 
				
			|||
    [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
 | 
				
			|||
}
 | 
				
			|||