incrementalmerkletree/src/sample.rs

/// Sample implementation of the Tree interface.
use crate::{TreeHasher, Recording, Tree};

#[derive(Clone)]
pub struct CompleteTree<H: TreeHasher> {
    leaves: Vec<H::Digest>,
    current_position: usize,
    witnesses: Vec<(usize, H::Digest)>,
    checkpoints: Vec<usize>,
    depth: usize,
}

impl<H: TreeHasher> CompleteTree<H> {
    /// Creates a new, empty binary tree of specified depth.
    ///
    /// # Panics
    ///
    /// Panics if the specified depth is zero.
    pub fn new(depth: usize) -> Self {
        if depth == 0 {
            panic!("invalid depth for incremental merkle tree");
        }

        CompleteTree {
            leaves: vec![H::empty_leaf(); 1 << depth],
            current_position: 0,
            witnesses: vec![],
            checkpoints: vec![],
            depth,
        }
    }
}

impl<H: TreeHasher> Tree<H> for CompleteTree<H> {
    type Recording = CompleteRecording<H>;

    /// Appends a new value to the tree at the next available slot. Returns true
    /// if successful and false if the tree is full.
    fn append(&mut self, value: &H::Digest) -> bool {
        if self.current_position == (1 << self.depth) {
            false
        } else {
            self.leaves[self.current_position] = value.clone();
            self.current_position += 1;
            true
        }
    }

    /// Obtains the current root of this Merkle tree.
    fn root(&self) -> H::Digest {
        lazy_root::<H>(self.leaves.clone())
    }

    /// Marks the current tree state leaf as a value that we're interested in
    /// witnessing. Returns true if successful and false if the tree is empty.
    fn witness(&mut self) -> bool {
        if self.current_position == 0 {
            return false;
        } else {
            let value = self.leaves[self.current_position - 1].clone();
            self.witnesses.push((self.current_position - 1, value));
            true
        }
    }

    /// Obtains an authentication path to the value specified in the tree.
    /// Returns `None` if there is no available authentication path to the
    /// specified value.
    fn authentication_path(&self, value: &H::Digest) -> Option<(usize, Vec<H::Digest>)> {
        self.witnesses
            .iter()
            .find(|witness| witness.1 == *value)
            .map(|&(pos, _)| {
                let mut path = vec![];

                let mut index = pos;
                for bit in 0..self.depth {
                    index ^= 1 << bit;
                    path.push(lazy_root::<H>(self.leaves[index..][0..(1 << bit)].to_vec()));
                    index &= usize::MAX << (bit + 1);
                }

                (pos, path)
            })
    }

    /// Marks the specified tree state value as a value we're no longer
    /// interested in maintaining a witness for. Returns true if successful and
    /// false if the value is not a known witness.
    fn remove_witness(&mut self, value: &H::Digest) -> bool {
        if let Some((position, _)) = self
            .witnesses
            .iter()
            .enumerate()
            .find(|witness| (witness.1).1 == *value)
        {
            self.witnesses.remove(position);

            true
        } else {
            false
        }
    }

    /// Marks the current tree state as a checkpoint if it is not already a
    /// checkpoint.
    fn checkpoint(&mut self) {
        self.checkpoints.push(self.current_position);
    }

    /// Rewinds the tree state to the previous checkpoint. This function will
    /// fail and return false if there is no previous checkpoint or in the event
    /// witness data would be destroyed in the process.
    fn rewind(&mut self) -> bool {
        if let Some(checkpoint) = self.checkpoints.pop() {
            if self.witnesses.iter().any(|&(pos, _)| pos >= checkpoint) {
                self.checkpoints.push(checkpoint);
                return false;
            }

            self.current_position = checkpoint;
            if checkpoint != (1 << self.depth) {
                self.leaves[checkpoint..].fill(H::empty_leaf());
            }

            true
        } else {
            false
        }
    }

    /// Removes the oldest checkpoint. Returns true if successful and false if
    /// there are no checkpoints.
    fn pop_checkpoint(&mut self) -> bool {
        if self.checkpoints.is_empty() {
            false
        } else {
            self.checkpoints.remove(0);
            true
        }
    }

    /// Start a recording of append operations performed on a tree.
    fn recording(&self) -> CompleteRecording<H> {
        CompleteRecording {
            start_position: self.current_position,
            current_position: self.current_position,
            depth: self.depth,
            appends: vec![],
        }
    }

    /// Plays a recording of append operations back. Returns true if successful
    /// and false if the recording is incompatible with the current tree state.
    fn play(&mut self, recording: &CompleteRecording<H>) -> bool {
        if recording.start_position == self.current_position && self.depth == recording.depth {
            for val in recording.appends.iter() {
                self.append(val);
            }
            true
        } else {
            false
        }
    }
}

#[derive(Clone)]
pub struct CompleteRecording<H: TreeHasher> {
    start_position: usize,
    current_position: usize,
    depth: usize,
    appends: Vec<H::Digest>,
}

impl<H: TreeHasher> Recording<H> for CompleteRecording<H> {
    /// Appends a new value to the tree at the next available slot. Returns true
    /// if successful and false if the tree is full.
    fn append(&mut self, value: &H::Digest) -> bool {
        if self.current_position == (1 << self.depth) {
            false
        } else {
            self.appends.push(value.clone());
            self.current_position += 1;

            true
        }
    }

    /// Plays a recording of append operations back. Returns true if successful
    /// and false if the provided recording is incompatible with `Self`.
    fn play(&mut self, recording: &Self) -> bool {
        if self.current_position == recording.start_position && self.depth == recording.depth {
            self.appends.extend_from_slice(&recording.appends);
            self.current_position = recording.current_position;
            true
        } else {
            false
        }
    }
}

pub(crate) fn lazy_root<H: TreeHasher>(mut leaves: Vec<H::Digest>) -> H::Digest {
    while leaves.len() != 1 {
        leaves = leaves
            .iter()
            .enumerate()
            .filter(|(i, _)| (i % 2) == 0)
            .map(|(_, a)| a)
            .zip(
                leaves
                    .iter()
                    .enumerate()
                    .filter(|(i, _)| (i % 2) == 1)
                    .map(|(_, b)| b),
            )
            .map(|(a, b)| H::combine(a, b))
            .collect();
    }

    leaves[0].clone()
}


#[cfg(test)]
mod tests {
    #![allow(deprecated)]
    use crate::tests::{compute_root_from_auth_path};
    use crate::{Tree, TreeHasher};

    use super::CompleteTree;

    use std::hash::SipHasher as Hash;

    #[test]
    fn correct_empty_root() {
        const DEPTH: usize = 5;
        let mut expected = 0u64;
        for _ in 0..DEPTH {
            expected = Hash::combine(&expected, &expected);
        }

        let tree = CompleteTree::<Hash>::new(DEPTH);
        assert_eq!(tree.root(), expected);
    }

    #[test]
    fn correct_root() {
        const DEPTH: usize = 3;
        let values: Vec<u64> = (0..(1 << DEPTH)).collect();

        let mut tree = CompleteTree::<Hash>::new(DEPTH);
        for value in values.iter() {
            assert!(tree.append(value));
        }
        assert!(!tree.append(&0));

        let expected = Hash::combine(
            &Hash::combine(&Hash::combine(&0, &1), &Hash::combine(&2, &3)),
            &Hash::combine(&Hash::combine(&4, &5), &Hash::combine(&6, &7)),
        );

        assert_eq!(tree.root(), expected);
    }

    #[test]
    fn correct_auth_path() {
        const DEPTH: usize = 3;
        let values: Vec<u64> = (0..(1 << DEPTH)).collect();

        let mut tree = CompleteTree::<Hash>::new(DEPTH);
        for value in values.iter() {
            assert!(tree.append(value));
            tree.witness();
        }
        assert!(!tree.append(&0));

        let expected = Hash::combine(
            &Hash::combine(&Hash::combine(&0, &1), &Hash::combine(&2, &3)),
            &Hash::combine(&Hash::combine(&4, &5), &Hash::combine(&6, &7)),
        );

        assert_eq!(tree.root(), expected);

        for i in 0..(1 << DEPTH) {
            println!("value: {}", i);
            let (position, path) = tree.authentication_path(&i).unwrap();
            assert_eq!(
                compute_root_from_auth_path::<Hash>(i, position, &path),
                expected
            );
        }
    }
}