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(true) if the proof is accepted, Ok(false) if not. Err()
218    ///   if the proof is not well structured, E.g. not for this merkle tree.
219    fn verify(
220        commitment: impl Borrow<Self::Commitment>,
221        pos: impl Borrow<Self::Index>,
222        element: impl Borrow<Self::Element>,
223        proof: impl Borrow<Self::MembershipProof>,
224    ) -> Result<VerificationResult, MerkleTreeError>;
225
226    // fn batch_lookup(&self, pos: impl Iterator<Item = usize>) -> LookupResult<(),
227    // Self::BatchProof>; fn batch_verify(
228    //     &self,
229    //     pos: impl Iterator<Item = usize>,
230    //     proof: impl Borrow<Self::BatchProof>,
231    // ) -> Result<(), MerkleTreeError>;
232
233    /// Return an iterator that iterates through all element that are not
234    /// forgotten
235    fn iter(&'_ self) -> MerkleTreeIter<'_, Self::Element, Self::Index, Self::NodeValue>;
236}
237
238/// Merkle tree that allows insertion at back. Abstracted as a commitment for
239/// append-only vector.
240pub trait AppendableMerkleTreeScheme: MerkleTreeScheme<Index = u64> {
241    /// Insert a new value at the leftmost available slot
242    /// * `elem` - element to insert in the tree
243    /// * `returns` - Ok(()) if successful
244    fn push(&mut self, elem: impl Borrow<Self::Element>) -> Result<(), MerkleTreeError>;
245
246    /// Insert a list of new values at the leftmost available slots
247    /// * `elems` - elements to insert
248    /// * `returns` - Ok(()) if successful. If there are too many elements,
249    ///   insertions will be performed until the merkle tree is full, and will
250    ///   return an Err().
251    fn extend(
252        &mut self,
253        elems: impl IntoIterator<Item = impl Borrow<Self::Element>>,
254    ) -> Result<(), MerkleTreeError> {
255        for elem in elems {
256            self.push(elem)?;
257        }
258        Ok(())
259    }
260}
261
262/// A universal merkle tree is abstracted as a random-access array or a
263/// key-value map. It allows manipulation at any given position, and has ability
264/// to generate/verify a non-membership proof.
265pub trait UniversalMerkleTreeScheme: MerkleTreeScheme {
266    /// Non membership proof for a given index
267    type NonMembershipProof;
268    /// Batch non membership proof
269    type BatchNonMembershipProof;
270
271    /// Update the leaf value at a given position
272    /// * `pos` - zero-based index of the leaf in the tree
273    /// * `elem` - newly updated element
274    /// * `returns` - Err() if any error occurs internally or the given leaf is
275    ///   not in memory. Ok(result) if the update is success.
276    fn update(
277        &mut self,
278        pos: impl Borrow<Self::Index>,
279        elem: impl Borrow<Self::Element>,
280    ) -> Result<LookupResult<Self::Element, (), ()>, MerkleTreeError> {
281        self.update_with(pos, |_| Some(elem.borrow().clone()))
282    }
283
284    /// Remove a leaf at the given position
285    /// * `pos` - zero-based index of the leaf in the tree
286    /// * `returns` - Err() if any error occurs internally. Ok(result) if the
287    ///   update is success or the given leaf is not in memory.
288    fn remove(
289        &mut self,
290        pos: impl Borrow<Self::Index>,
291    ) -> Result<LookupResult<Self::Element, (), ()>, MerkleTreeError> {
292        self.update_with(pos, |_| None)
293    }
294
295    /// Apply an update function `f` at a given position
296    /// * `pos` - zero-based index of the leaf in the tree
297    /// * `f` - the update function, `None` means the given leaf doesn't exist
298    ///   or should be removed.
299    /// * `returns` - Err() if any error occurs internally. Ok(result) if the
300    ///   update is success or the given leaf is not in memory.
301    ///
302    /// # Example
303    /// ```ignore
304    /// merkle_tree.update_with(account, |balance| {
305    ///     Some(balance.cloned().unwrap_or_default().add(amount))
306    /// });
307    /// ```
308    fn update_with<F>(
309        &mut self,
310        pos: impl Borrow<Self::Index>,
311        f: F,
312    ) -> Result<LookupResult<Self::Element, (), ()>, MerkleTreeError>
313    where
314        F: FnOnce(Option<&Self::Element>) -> Option<Self::Element>;
315
316    /// Returns the leaf value given a position
317    /// * `pos` - zero-based index of the leaf in the tree
318    /// * `returns` - Leaf value at the position along with a proof.
319    ///   LookupResult::EmptyLeaf(p) if the leaf position is empty along with a
320    ///   proof p. LookupResult::NotInMemory if the leaf position has been
321    ///   forgotten.
322    fn universal_lookup(
323        &self,
324        pos: impl Borrow<Self::Index>,
325    ) -> LookupResult<&Self::Element, Self::MembershipProof, Self::NonMembershipProof>;
326
327    /// Verify an index is not in this merkle tree
328    /// * `pos` - zero-based index of the leaf in the tree
329    /// * `proof` - a merkle tree proof
330    /// * `returns` - Ok(true) if the proof is accepted, Ok(false) if not. Err()
331    ///   if the proof is not well structured, E.g. not for this merkle tree.
332    fn non_membership_verify(
333        commitment: impl Borrow<Self::Commitment>,
334        pos: impl Borrow<Self::Index>,
335        proof: impl Borrow<Self::NonMembershipProof>,
336    ) -> Result<VerificationResult, MerkleTreeError>;
337    // TODO(Chengyu): non-membership proof interfaces
338}
339
340/// Merkle tree that allows forget/remember elements from the memory
341pub trait ForgetableMerkleTreeScheme: MerkleTreeScheme {
342    /// Trim the leaf at position `i` from memory, if present.
343    /// Should not trim if position `i` is the last inserted leaf position.
344    /// Return is identical to result if `get_leaf(pos)` were called before this
345    /// call.
346    fn forget(
347        &mut self,
348        pos: impl Borrow<Self::Index>,
349    ) -> LookupResult<Self::Element, Self::MembershipProof, ()>;
350
351    /// "Re-insert" a leaf into the tree using its proof.
352    /// Returns Ok(()) if insertion is successful, or Err(err) if the
353    /// proof disagrees with the merkle tree
354    fn remember(
355        &mut self,
356        pos: impl Borrow<Self::Index>,
357        element: impl Borrow<Self::Element>,
358        proof: impl Borrow<Self::MembershipProof>,
359    ) -> Result<(), MerkleTreeError>;
360
361    /// Rebuild a merkle tree from a commitment.
362    /// Return a tree which is entirely forgotten.
363    fn from_commitment(
364        commitment: impl Borrow<Self::Commitment>,
365        height: usize,
366        num_leaves: u64,
367    ) -> Self;
368}
369
370/// Universal Merkle tree that allows forget/remember elements from the memory
371pub trait ForgetableUniversalMerkleTreeScheme:
372    ForgetableMerkleTreeScheme + UniversalMerkleTreeScheme
373{
374    /// Trim the leaf at position `pos` from memory.
375    ///
376    /// This is similar to [forget](ForgetableMerkleTreeScheme::forget), but it
377    /// may prune even an empty sub-tree at `pos` and will return a
378    /// non-membership proof for the pruned position if it does so. Note
379    /// that an implementation may choose _not_ to prune an empty sub-tree, as
380    /// it may be more efficient to represent an empty sub-tree than a
381    /// forgotten one. In this case,
382    /// [universal_lookup](UniversalMerkleTreeScheme::universal_lookup) may
383    /// return _either_ [LookupResult::NotInMemory] or
384    /// [LookupResult::NotFound] after a successful call to
385    /// [universal_forget](Self::universal_forget). In any case, if this
386    /// function is called for a `pos` which is in memory but not in the
387    /// tree, it will return a non-membership proof.
388    ///
389    /// The return value is the same as if
390    /// [universal_lookup](UniversalMerkleTreeScheme::universal_lookup) were
391    /// called before this call.
392    fn universal_forget(
393        &mut self,
394        pos: Self::Index,
395    ) -> LookupResult<Self::Element, Self::MembershipProof, Self::NonMembershipProof>;
396
397    /// "Re-insert" an empty leaf into the tree using its proof.
398    ///
399    /// Returns `Ok(())` if insertion is successful, or `Err(err)` if the proof
400    /// disagrees with the merkle tree
401    fn non_membership_remember(
402        &mut self,
403        pos: Self::Index,
404        proof: impl Borrow<Self::NonMembershipProof>,
405    ) -> Result<(), MerkleTreeError>;
406}
407
408/// A universal merkle tree that allows non destructive updates.
409/// A non destructive update doesn't directly modify the existing content, it
410/// creates a new copy about the update so that people could access both the old
411/// version and the new.
412pub trait PersistentUniversalMerkleTreeScheme: UniversalMerkleTreeScheme {
413    /// A non destructive update interface, check
414    /// [PersistentUniversalMerkleTreeScheme] and
415    /// [UniversalMerkleTreeScheme::update].
416    fn persistent_update(
417        &self,
418        pos: impl Borrow<Self::Index>,
419        elem: impl Borrow<Self::Element>,
420    ) -> Result<Self, MerkleTreeError> {
421        self.persistent_update_with(pos, |_| Some(elem.borrow().clone()))
422    }
423
424    /// A persistent remove interface, check
425    /// [PersistentUniversalMerkleTreeScheme] and
426    /// [UniversalMerkleTreeScheme::remove].
427    fn persistent_remove(&self, pos: Self::Index) -> Result<Self, MerkleTreeError> {
428        self.persistent_update_with(pos, |_| None)
429    }
430
431    /// A persistent update_with interface, check
432    /// [PersistentUniversalMerkleTreeScheme] and
433    /// [UniversalMerkleTreeScheme::update_with].
434    fn persistent_update_with<F>(
435        &self,
436        pos: impl Borrow<Self::Index>,
437        f: F,
438    ) -> Result<Self, MerkleTreeError>
439    where
440        F: FnOnce(Option<&Self::Element>) -> Option<Self::Element>;
441}
jf_merkle_tree/lib.rs

jf_merkle_tree/
lib.rs
