jf_merkle_tree/
lib.rs

1// Copyright (c) 2022 Espresso Systems (espressosys.com)
2// This file is part of the Jellyfish library.
3
4// You should have received a copy of the MIT License
5// along with the Jellyfish library. If not, see <https://mit-license.org/>.
6
7//! Merkle Tree traits and implementations
8
9#![cfg_attr(not(feature = "std"), no_std)]
10#![deny(missing_docs)]
11#[cfg(test)]
12extern crate std;
13
14#[cfg(any(not(feature = "std"), target_has_atomic = "ptr"))]
15#[doc(hidden)]
16extern crate alloc;
17
18pub mod append_only;
19pub mod errors;
20pub mod examples;
21#[cfg(feature = "gadgets")]
22pub mod gadgets;
23pub mod hasher;
24pub mod light_weight;
25pub mod macros;
26pub mod universal_merkle_tree;
27
28pub(crate) mod internal;
29
30pub mod prelude;
31pub use crate::errors::MerkleTreeError;
32
33use self::internal::MerkleTreeIter;
34use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
35use ark_std::{borrow::Borrow, fmt::Debug, hash::Hash, vec, vec::Vec};
36use num_bigint::BigUint;
37use num_traits::ToPrimitive;
38use serde::{Deserialize, Serialize};
39
40/// Glorified bool type
41pub(crate) type VerificationResult = Result<(), ()>;
42/// Glorified true
43pub const SUCCESS: VerificationResult = Ok(());
44/// Glorified false
45pub const FAIL: VerificationResult = Err(());
46
47#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
48/// The result of querying at an index in the tree
49/// Typically, F for element type, P for membership proof type and N for
50/// non-membership proof type
51pub enum LookupResult<F, P, N> {
52    /// The value at the given index, and a proof of validity
53    Ok(F, P),
54    /// The index is valid but we do not have the leaf in memory
55    NotInMemory,
56    /// The index is outside the occupied range in the tree, and a
57    /// non-membership proof
58    NotFound(N),
59}
60
61impl<F, P, N> LookupResult<F, P, N> {
62    /// Assert the lookup result is Ok. Return a tuple of element and membership
63    /// proof.
64    pub fn expect_ok(self) -> Result<(F, P), MerkleTreeError> {
65        match self {
66            LookupResult::Ok(x, proof) => Ok((x, proof)),
67            LookupResult::NotInMemory => Err(MerkleTreeError::ForgottenLeaf),
68            LookupResult::NotFound(_) => Err(MerkleTreeError::NotFound),
69        }
70    }
71
72    /// Assert the lookup result is NotFound. Return a non-membership proof.
73    pub fn expect_not_found(self) -> Result<N, MerkleTreeError> {
74        match self {
75            LookupResult::NotFound(n) => Ok(n),
76            LookupResult::Ok(..) => Err(MerkleTreeError::ExistingLeaf),
77            LookupResult::NotInMemory => Err(MerkleTreeError::ForgottenLeaf),
78        }
79    }
80
81    /// Assert the lookup result is NotInMemory.
82    pub fn expect_not_in_memory(self) -> Result<(), MerkleTreeError> {
83        match self {
84            LookupResult::NotInMemory => Ok(()),
85            LookupResult::Ok(..) => Err(MerkleTreeError::ExistingLeaf),
86            LookupResult::NotFound(..) => Err(MerkleTreeError::NotFound),
87        }
88    }
89}
90
91/// An element of a Merkle tree.
92pub trait Element: Clone + Eq + PartialEq + Hash {}
93impl<T: Clone + Eq + PartialEq + Hash> Element for T {}
94
95/// An index type of a leaf in a Merkle tree.
96pub trait Index: Debug + Eq + PartialEq + Hash + Ord + PartialOrd + Clone {}
97impl<T: Debug + Eq + PartialEq + Hash + Ord + PartialOrd + Clone> Index for T {}
98
99/// An internal node value type in a Merkle tree.
100pub trait NodeValue:
101    Default + Eq + PartialEq + Hash + Copy + Clone + Debug + CanonicalSerialize + CanonicalDeserialize
102{
103}
104impl<T> NodeValue for T where
105    T: Default
106        + Eq
107        + PartialEq
108        + Hash
109        + Copy
110        + Clone
111        + Debug
112        + CanonicalSerialize
113        + CanonicalDeserialize
114{
115}
116
117/// Merkle tree hash function
118/// WARN: it's important to domain separate the two digest functions. Otherwise,
119/// you may suffer from the length extension attack.
120pub trait DigestAlgorithm<E, I, T>
121where
122    E: Element,
123    I: Index,
124    T: NodeValue,
125{
126    /// Digest a list of values
127    fn digest(data: &[T]) -> Result<T, MerkleTreeError>;
128
129    /// Digest an indexed element
130    fn digest_leaf(pos: &I, elem: &E) -> Result<T, MerkleTreeError>;
131}
132
133/// A trait for Merkle tree index type.
134pub trait ToTraversalPath<const ARITY: usize> {
135    /// Convert the given index to a vector of branch indices given tree height
136    /// and ARITY.
137    fn to_traversal_path(&self, height: usize) -> Vec<usize>;
138}
139
140impl_to_traversal_path_primitives!(usize);
141impl_to_traversal_path_primitives!(u8);
142impl_to_traversal_path_primitives!(u16);
143impl_to_traversal_path_primitives!(u32);
144impl_to_traversal_path_primitives!(u64);
145impl_to_traversal_path_biguint!(u128);
146impl_to_traversal_path_biguint!(BigUint);
147impl_to_traversal_path_biguint!(ark_bn254::Fr);
148impl_to_traversal_path_biguint!(ark_bls12_377::Fr);
149impl_to_traversal_path_biguint!(ark_bls12_381::Fr);
150impl_to_traversal_path_biguint!(ark_bn254::Fq);
151impl_to_traversal_path_biguint!(ark_bls12_377::Fq);
152impl_to_traversal_path_biguint!(ark_bls12_381::Fq);
153
154/// Trait for a Merkle proof
155pub trait MerkleProof<T: NodeValue>:
156    Eq
157    + PartialEq
158    + Hash
159    + Clone
160    + CanonicalSerialize
161    + CanonicalDeserialize
162    + Serialize
163    + for<'a> Deserialize<'a>
164{
165    /// Expected height of the Merkle tree.
166    fn height(&self) -> usize;
167
168    /// Return all values of siblings of this Merkle path
169    fn path_values(&self) -> &[Vec<T>];
170}
171
172/// Basic functionalities for a merkle tree implementation. Abstracted as an
173/// accumulator for fixed-length array. Supports generating membership proof at
174/// a given position and verify a membership proof.
175pub trait MerkleTreeScheme: Sized {
176    /// Merkle tree element type
177    type Element: Element;
178    /// Index type for this merkle tree
179    type Index: Index;
180    /// Internal and root node value
181    type NodeValue: NodeValue;
182    /// Merkle proof
183    type MembershipProof: MerkleProof<Self::NodeValue>;
184    /// Batch proof
185    type BatchMembershipProof: Clone;
186    /// Merkle tree commitment
187    type Commitment: NodeValue;
188
189    /// Tree ARITY
190    const ARITY: usize;
191
192    /// Return the height of this merkle tree
193    fn height(&self) -> usize;
194    /// Return the maximum allowed number leaves
195    fn capacity(&self) -> BigUint;
196    /// Return the current number of leaves
197    fn num_leaves(&self) -> u64;
198
199    /// Return a merkle commitment
200    fn commitment(&self) -> Self::Commitment;
201
202    /// Returns the leaf value given a position
203    /// * `pos` - zero-based index of the leaf in the tree
204    /// * `returns` - Leaf value at the position along with a proof.
205    ///   LookupResult::EmptyLeaf if the leaf position is empty or invalid,
206    ///   LookupResult::NotInMemory if the leaf position has been forgotten.
207    fn lookup(
208        &self,
209        pos: impl Borrow<Self::Index>,
210    ) -> LookupResult<&Self::Element, Self::MembershipProof, ()>;
211
212    /// Verify an element is a leaf of a Merkle tree given the proof
213    /// * `commitment` - a merkle tree commitment
214    /// * `pos` - zero-based index of the leaf in the tree
215    /// * `element` - the leaf value
216    /// * `proof` - a membership proof for `element` at given `pos`
217    /// * `returns` - Ok(SUCCESS) if the proof is accepted, Ok(FAIL) if not.
218    ///   Err() if the proof is not well structured, E.g. not for this merkle
219    ///   tree.
220    fn verify(
221        commitment: impl Borrow<Self::Commitment>,
222        pos: impl Borrow<Self::Index>,
223        element: impl Borrow<Self::Element>,
224        proof: impl Borrow<Self::MembershipProof>,
225    ) -> Result<VerificationResult, MerkleTreeError>;
226
227    // fn batch_lookup(&self, pos: impl Iterator<Item = usize>) -> LookupResult<(),
228    // Self::BatchProof>; fn batch_verify(
229    //     &self,
230    //     pos: impl Iterator<Item = usize>,
231    //     proof: impl Borrow<Self::BatchProof>,
232    // ) -> Result<(), MerkleTreeError>;
233
234    /// Return an iterator that iterates through all element that are not
235    /// forgotten
236    fn iter(&'_ self) -> MerkleTreeIter<'_, Self::Element, Self::Index, Self::NodeValue>;
237}
238
239/// Merkle tree that supports generating range proofs
240pub trait RangeProofMerkleTreeScheme: MerkleTreeScheme {
241    /// Range proof for a given range
242    type RangeMembershipProof;
243
244    /// Generate a range proof for leaves in [start, end], inclusive.
245    /// * `start` - zero-based start index of the range (inclusive)
246    /// * `end` - zero-based end index of the range (inclusive)
247    /// * `returns` - Indexed leaf values in the range along with a range
248    ///   membership proof.
249    ///
250    /// NOTE: this interface will return `NotInMemory` or `NotFound` if any
251    /// of the leaves in the range are not in memory or not found.
252    #[allow(clippy::type_complexity)]
253    fn range_lookup(
254        &self,
255        start: impl Borrow<Self::Index>,
256        end: impl Borrow<Self::Index>,
257    ) -> LookupResult<(Vec<Self::Index>, Vec<Self::Element>), Self::RangeMembershipProof, ()>;
258
259    /// Verify a range proof for leaves in [start, end], inclusive.
260    /// * `commitment` - a merkle tree commitment
261    /// * `indices` - zero-based indices of the leaves in the tree
262    /// * `elements` - the leaf values in the range
263    /// * `proof` - a range membership proof for the given range
264    /// * `returns` - Ok(SUCCESS) if the proof is accepted, Ok(FAIL) if not.
265    ///   Err() if the proof is not well structured, E.g. not for this merkle
266    ///   tree.
267    ///
268    /// NOTE: all indices are needed here because we cannot generate them from
269    /// [start, end] under the current Index trait bound.
270    fn verify_range_proof(
271        commitment: impl Borrow<Self::Commitment>,
272        indices: &[impl Borrow<Self::Index>],
273        elements: &[impl Borrow<Self::Element>],
274        proof: impl Borrow<Self::RangeMembershipProof>,
275    ) -> Result<VerificationResult, MerkleTreeError>;
276}
277
278/// Merkle tree that allows insertion at back. Abstracted as a commitment for
279/// append-only vector.
280pub trait AppendableMerkleTreeScheme: MerkleTreeScheme<Index = u64> {
281    /// Insert a new value at the leftmost available slot
282    /// * `elem` - element to insert in the tree
283    /// * `returns` - Ok(()) if successful
284    fn push(&mut self, elem: impl Borrow<Self::Element>) -> Result<(), MerkleTreeError>;
285
286    /// Insert a list of new values at the leftmost available slots
287    /// * `elems` - elements to insert
288    /// * `returns` - Ok(()) if successful. If there are too many elements,
289    ///   insertions will be performed until the merkle tree is full, and will
290    ///   return an Err().
291    fn extend(
292        &mut self,
293        elems: impl IntoIterator<Item = impl Borrow<Self::Element>>,
294    ) -> Result<(), MerkleTreeError> {
295        for elem in elems {
296            self.push(elem)?;
297        }
298        Ok(())
299    }
300}
301
302/// A universal merkle tree is abstracted as a random-access array or a
303/// key-value map. It allows manipulation at any given position, and has ability
304/// to generate/verify a non-membership proof.
305pub trait UniversalMerkleTreeScheme: MerkleTreeScheme {
306    /// Non membership proof for a given index
307    type NonMembershipProof;
308    /// Batch non membership proof
309    type BatchNonMembershipProof;
310
311    /// Update the leaf value at a given position
312    /// * `pos` - zero-based index of the leaf in the tree
313    /// * `elem` - newly updated element
314    /// * `returns` - Err() if any error occurs internally or the given leaf is
315    ///   not in memory. Ok(result) if the update is success.
316    fn update(
317        &mut self,
318        pos: impl Borrow<Self::Index>,
319        elem: impl Borrow<Self::Element>,
320    ) -> Result<LookupResult<Self::Element, (), ()>, MerkleTreeError> {
321        self.update_with(pos, |_| Some(elem.borrow().clone()))
322    }
323
324    /// Remove a leaf at the given position
325    /// * `pos` - zero-based index of the leaf in the tree
326    /// * `returns` - Err() if any error occurs internally. Ok(result) if the
327    ///   update is success or the given leaf is not in memory.
328    fn remove(
329        &mut self,
330        pos: impl Borrow<Self::Index>,
331    ) -> Result<LookupResult<Self::Element, (), ()>, MerkleTreeError> {
332        self.update_with(pos, |_| None)
333    }
334
335    /// Apply an update function `f` at a given position
336    /// * `pos` - zero-based index of the leaf in the tree
337    /// * `f` - the update function, `None` means the given leaf doesn't exist
338    ///   or should be removed.
339    /// * `returns` - Err() if any error occurs internally. Ok(result) if the
340    ///   update is success or the given leaf is not in memory.
341    ///
342    /// # Example
343    /// ```ignore
344    /// merkle_tree.update_with(account, |balance| {
345    ///     Some(balance.cloned().unwrap_or_default().add(amount))
346    /// });
347    /// ```
348    fn update_with<F>(
349        &mut self,
350        pos: impl Borrow<Self::Index>,
351        f: F,
352    ) -> Result<LookupResult<Self::Element, (), ()>, MerkleTreeError>
353    where
354        F: FnOnce(Option<&Self::Element>) -> Option<Self::Element>;
355
356    /// Returns the leaf value given a position
357    /// * `pos` - zero-based index of the leaf in the tree
358    /// * `returns` - Leaf value at the position along with a proof.
359    ///   LookupResult::EmptyLeaf(p) if the leaf position is empty along with a
360    ///   proof p. LookupResult::NotInMemory if the leaf position has been
361    ///   forgotten.
362    fn universal_lookup(
363        &self,
364        pos: impl Borrow<Self::Index>,
365    ) -> LookupResult<&Self::Element, Self::MembershipProof, Self::NonMembershipProof>;
366
367    /// Verify an index is not in this merkle tree
368    /// * `pos` - zero-based index of the leaf in the tree
369    /// * `proof` - a merkle tree proof
370    /// * `returns` - Ok(SUCCESS) if the proof is accepted, Ok(FAIL) if not.
371    ///   Err() if the proof is not well structured, E.g. not for this merkle
372    ///   tree.
373    fn non_membership_verify(
374        commitment: impl Borrow<Self::Commitment>,
375        pos: impl Borrow<Self::Index>,
376        proof: impl Borrow<Self::NonMembershipProof>,
377    ) -> Result<VerificationResult, MerkleTreeError>;
378    // TODO(Chengyu): non-membership proof interfaces
379}
380
381/// Merkle tree that allows forget/remember elements from the memory
382pub trait ForgetableMerkleTreeScheme: MerkleTreeScheme {
383    /// Trim the leaf at position `i` from memory, if present.
384    /// Should not trim if position `i` is the last inserted leaf position.
385    /// Return is identical to result if `get_leaf(pos)` were called before this
386    /// call.
387    fn forget(
388        &mut self,
389        pos: impl Borrow<Self::Index>,
390    ) -> LookupResult<Self::Element, Self::MembershipProof, ()>;
391
392    /// "Re-insert" a leaf into the tree using its proof.
393    /// Returns Ok(()) if insertion is successful, or Err(err) if the
394    /// proof disagrees with the merkle tree
395    fn remember(
396        &mut self,
397        pos: impl Borrow<Self::Index>,
398        element: impl Borrow<Self::Element>,
399        proof: impl Borrow<Self::MembershipProof>,
400    ) -> Result<(), MerkleTreeError>;
401
402    /// Rebuild a merkle tree from a commitment.
403    /// Return a tree which is entirely forgotten.
404    fn from_commitment(
405        commitment: impl Borrow<Self::Commitment>,
406        height: usize,
407        num_leaves: u64,
408    ) -> Self;
409}
410
411/// Universal Merkle tree that allows forget/remember elements from the memory
412pub trait ForgetableUniversalMerkleTreeScheme:
413    ForgetableMerkleTreeScheme + UniversalMerkleTreeScheme
414{
415    /// Trim the leaf at position `pos` from memory.
416    ///
417    /// This is similar to [forget](ForgetableMerkleTreeScheme::forget), but it
418    /// may prune even an empty sub-tree at `pos` and will return a
419    /// non-membership proof for the pruned position if it does so. Note
420    /// that an implementation may choose _not_ to prune an empty sub-tree, as
421    /// it may be more efficient to represent an empty sub-tree than a
422    /// forgotten one. In this case,
423    /// [universal_lookup](UniversalMerkleTreeScheme::universal_lookup) may
424    /// return _either_ [LookupResult::NotInMemory] or
425    /// [LookupResult::NotFound] after a successful call to
426    /// [universal_forget](Self::universal_forget). In any case, if this
427    /// function is called for a `pos` which is in memory but not in the
428    /// tree, it will return a non-membership proof.
429    ///
430    /// The return value is the same as if
431    /// [universal_lookup](UniversalMerkleTreeScheme::universal_lookup) were
432    /// called before this call.
433    fn universal_forget(
434        &mut self,
435        pos: Self::Index,
436    ) -> LookupResult<Self::Element, Self::MembershipProof, Self::NonMembershipProof>;
437
438    /// "Re-insert" an empty leaf into the tree using its proof.
439    ///
440    /// Returns `Ok(())` if insertion is successful, or `Err(err)` if the proof
441    /// disagrees with the merkle tree
442    fn non_membership_remember(
443        &mut self,
444        pos: Self::Index,
445        proof: impl Borrow<Self::NonMembershipProof>,
446    ) -> Result<(), MerkleTreeError>;
447}
448
449/// A universal merkle tree that allows non destructive updates.
450/// A non destructive update doesn't directly modify the existing content, it
451/// creates a new copy about the update so that people could access both the old
452/// version and the new.
453pub trait PersistentUniversalMerkleTreeScheme: UniversalMerkleTreeScheme {
454    /// A non destructive update interface, check
455    /// [PersistentUniversalMerkleTreeScheme] and
456    /// [UniversalMerkleTreeScheme::update].
457    fn persistent_update(
458        &self,
459        pos: impl Borrow<Self::Index>,
460        elem: impl Borrow<Self::Element>,
461    ) -> Result<Self, MerkleTreeError> {
462        self.persistent_update_with(pos, |_| Some(elem.borrow().clone()))
463    }
464
465    /// A persistent remove interface, check
466    /// [PersistentUniversalMerkleTreeScheme] and
467    /// [UniversalMerkleTreeScheme::remove].
468    fn persistent_remove(&self, pos: Self::Index) -> Result<Self, MerkleTreeError> {
469        self.persistent_update_with(pos, |_| None)
470    }
471
472    /// A persistent update_with interface, check
473    /// [PersistentUniversalMerkleTreeScheme] and
474    /// [UniversalMerkleTreeScheme::update_with].
475    fn persistent_update_with<F>(
476        &self,
477        pos: impl Borrow<Self::Index>,
478        f: F,
479    ) -> Result<Self, MerkleTreeError>
480    where
481        F: FnOnce(Option<&Self::Element>) -> Option<Self::Element>;
482}
jf_merkle_tree/lib.rs

jf_merkle_tree/
lib.rs
