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

jf_merkle_tree/
lib.rs
