From e55b3ed2ce56d189f9cc6836348bf3c9967d1c51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philippe=20Laferri=C3=A8re?= Date: Tue, 6 Feb 2024 13:36:31 -0500 Subject: [PATCH] Introduce `SmtProof` (#270) * add conversion for `SmtLeaf` * introduce `SmtProof` scaffolding * implement `verify_membership()` * SmtLeaf: knows its index * `SmtLeaf::index` * `SmtLeaf::get_value()` returns an Option * fix `verify_membership()` * impl `SmtProof::get` * impl `into_parts()` * `SmtProof::compute_root` * use `SmtProof` in `Smt::open` * `SmtLeaf` constructors * Vec * impl `Error` for `SmtLeafError` * fix std Error * move Word/Digest conversions to LeafIndex * `SmtProof::new()` returns an error * `SparseMerkleTree::path_and_leaf_to_opening` * `SmtLeaf`: serializable/deserializable * `SmtProof`: serializable/deserializable * add tests for SmtLeaf serialization * move `SmtLeaf` to submodule * use constructors internally * fix docs * Add `Vec` * add `Vec` to tests * no_std use statements * fmt * `Errors`: make heading * use `SMT_DEPTH` * SmtLeaf single case: check leaf index * Multiple case: check consistency with leaf index * use `pub(super)` instead of `pub(crate)` * use `pub(super)` * `SmtLeaf`: add `num_entries()` accessor * Fix `SmtLeaf` serialization * improve leaf serialization tests --- src/merkle/error.rs | 10 + src/merkle/smt/full/error.rs | 86 ++++++++ src/merkle/smt/full/leaf.rs | 373 +++++++++++++++++++++++++++++++++++ src/merkle/smt/full/mod.rs | 224 +++------------------ src/merkle/smt/full/proof.rs | 104 ++++++++++ src/merkle/smt/full/tests.rs | 65 +++++- src/merkle/smt/mod.rs | 26 +-- src/merkle/smt/simple/mod.rs | 6 +- 8 files changed, 678 insertions(+), 216 deletions(-) create mode 100644 src/merkle/smt/full/error.rs create mode 100644 src/merkle/smt/full/leaf.rs create mode 100644 src/merkle/smt/full/proof.rs diff --git a/src/merkle/error.rs b/src/merkle/error.rs index b513212..2ecfeba 100644 --- a/src/merkle/error.rs +++ b/src/merkle/error.rs @@ -4,6 +4,8 @@ use crate::{ }; use core::fmt; +use super::smt::SmtLeafError; + #[derive(Clone, Debug, PartialEq, Eq)] pub enum MerkleError { ConflictingRoots(Vec), @@ -20,6 +22,7 @@ pub enum MerkleError { NodeNotInStore(RpoDigest, NodeIndex), NumLeavesNotPowerOfTwo(usize), RootNotInStore(RpoDigest), + SmtLeaf(SmtLeafError), } impl fmt::Display for MerkleError { @@ -50,9 +53,16 @@ impl fmt::Display for MerkleError { write!(f, "the leaves count {leaves} is not a power of 2") } RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root), + SmtLeaf(smt_leaf_error) => write!(f, "smt leaf error: {smt_leaf_error}"), } } } #[cfg(feature = "std")] impl std::error::Error for MerkleError {} + +impl From for MerkleError { + fn from(value: SmtLeafError) -> Self { + Self::SmtLeaf(value) + } +} diff --git a/src/merkle/smt/full/error.rs b/src/merkle/smt/full/error.rs new file mode 100644 index 0000000..b629043 --- /dev/null +++ b/src/merkle/smt/full/error.rs @@ -0,0 +1,86 @@ +use core::fmt; + +use crate::{ + hash::rpo::RpoDigest, + merkle::{LeafIndex, SMT_DEPTH}, + utils::collections::Vec, + Word, +}; + +// SMT LEAF ERROR +// ================================================================================================= + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SmtLeafError { + InconsistentKeys { + entries: Vec<(RpoDigest, Word)>, + key_1: RpoDigest, + key_2: RpoDigest, + }, + InvalidNumEntriesForMultiple(usize), + SingleKeyInconsistentWithLeafIndex { + key: RpoDigest, + leaf_index: LeafIndex, + }, + MultipleKeysInconsistentWithLeafIndex { + leaf_index_from_keys: LeafIndex, + leaf_index_supplied: LeafIndex, + }, +} + +#[cfg(feature = "std")] +impl std::error::Error for SmtLeafError {} + +impl fmt::Display for SmtLeafError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + use SmtLeafError::*; + match self { + InvalidNumEntriesForMultiple(num_entries) => { + write!(f, "Multiple leaf requires 2 or more entries. Got: {num_entries}") + } + InconsistentKeys { entries, key_1, key_2 } => { + write!(f, "Multiple leaf requires all keys to map to the same leaf index. Offending keys: {key_1} and {key_2}. Entries: {entries:?}.") + } + SingleKeyInconsistentWithLeafIndex { key, leaf_index } => { + write!( + f, + "Single key in leaf inconsistent with leaf index. Key: {key}, leaf index: {}", + leaf_index.value() + ) + } + MultipleKeysInconsistentWithLeafIndex { + leaf_index_from_keys, + leaf_index_supplied, + } => { + write!( + f, + "Keys in entries map to leaf index {}, but leaf index {} was supplied", + leaf_index_from_keys.value(), + leaf_index_supplied.value() + ) + } + } + } +} + +// SMT PROOF ERROR +// ================================================================================================= + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SmtProofError { + InvalidPathLength(usize), +} + +#[cfg(feature = "std")] +impl std::error::Error for SmtProofError {} + +impl fmt::Display for SmtProofError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + use SmtProofError::*; + match self { + InvalidPathLength(path_length) => { + write!(f, "Invalid Merkle path length. Expected {SMT_DEPTH}, got {path_length}") + } + } + } +} diff --git a/src/merkle/smt/full/leaf.rs b/src/merkle/smt/full/leaf.rs new file mode 100644 index 0000000..5d3e485 --- /dev/null +++ b/src/merkle/smt/full/leaf.rs @@ -0,0 +1,373 @@ +use core::cmp::Ordering; + +use crate::utils::{collections::Vec, string::ToString, vec}; +use winter_math::StarkField; +use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; + +use super::{Felt, LeafIndex, Rpo256, RpoDigest, SmtLeafError, Word, EMPTY_WORD, SMT_DEPTH}; + +#[derive(Clone, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] +pub enum SmtLeaf { + Empty(LeafIndex), + Single((RpoDigest, Word)), + Multiple(Vec<(RpoDigest, Word)>), +} + +impl SmtLeaf { + // CONSTRUCTORS + // --------------------------------------------------------------------------------------------- + + /// Returns a new leaf with the specified entries + /// + /// # Errors + /// - Returns an error if 2 keys in `entries` map to a different leaf index + /// - Returns an error if 1 or more keys in `entries` map to a leaf index + /// different from `leaf_index` + pub fn new( + entries: Vec<(RpoDigest, Word)>, + leaf_index: LeafIndex, + ) -> Result { + match entries.len() { + 0 => Ok(Self::new_empty(leaf_index)), + 1 => { + let (key, value) = entries[0]; + + if LeafIndex::::from(key) != leaf_index { + return Err(SmtLeafError::SingleKeyInconsistentWithLeafIndex { + key, + leaf_index, + }); + } + + Ok(Self::new_single(key, value)) + } + _ => { + let leaf = Self::new_multiple(entries)?; + + // `new_multiple()` checked that all keys map to the same leaf index. We still need + // to ensure that that leaf index is `leaf_index`. + if leaf.index() != leaf_index { + Err(SmtLeafError::MultipleKeysInconsistentWithLeafIndex { + leaf_index_from_keys: leaf.index(), + leaf_index_supplied: leaf_index, + }) + } else { + Ok(leaf) + } + } + } + } + + /// Returns a new empty leaf with the specified leaf index + pub fn new_empty(leaf_index: LeafIndex) -> Self { + Self::Empty(leaf_index) + } + + /// Returns a new single leaf with the specified entry. The leaf index is derived from the + /// entry's key. + pub fn new_single(key: RpoDigest, value: Word) -> Self { + Self::Single((key, value)) + } + + /// Returns a new single leaf with the specified entry. The leaf index is derived from the + /// entries' keys. + /// + /// # Errors + /// - Returns an error if 2 keys in `entries` map to a different leaf index + pub fn new_multiple(entries: Vec<(RpoDigest, Word)>) -> Result { + if entries.len() < 2 { + return Err(SmtLeafError::InvalidNumEntriesForMultiple(entries.len())); + } + + // Check that all keys map to the same leaf index + { + let mut keys = entries.iter().map(|(key, _)| key); + + let first_key = *keys.next().expect("ensured at least 2 entries"); + let first_leaf_index: LeafIndex = first_key.into(); + + for &next_key in keys { + let next_leaf_index: LeafIndex = next_key.into(); + + if next_leaf_index != first_leaf_index { + return Err(SmtLeafError::InconsistentKeys { + entries, + key_1: first_key, + key_2: next_key, + }); + } + } + } + + Ok(Self::Multiple(entries)) + } + + // PUBLIC ACCESSORS + // --------------------------------------------------------------------------------------------- + + /// Returns true if the leaf is empty + pub fn is_empty(&self) -> bool { + matches!(self, Self::Empty(_)) + } + + /// Returns the leaf's index in the [`super::Smt`] + pub fn index(&self) -> LeafIndex { + match self { + SmtLeaf::Empty(leaf_index) => *leaf_index, + SmtLeaf::Single((key, _)) => key.into(), + SmtLeaf::Multiple(entries) => { + // Note: All keys are guaranteed to have the same leaf index + let (first_key, _) = entries[0]; + first_key.into() + } + } + } + + /// Returns the number of entries stored in the leaf + pub fn num_entries(&self) -> u64 { + match self { + SmtLeaf::Empty(_) => 0, + SmtLeaf::Single(_) => 1, + SmtLeaf::Multiple(entries) => { + entries.len().try_into().expect("shouldn't have more than 2^64 entries") + } + } + } + + /// Computes the hash of the leaf + pub fn hash(&self) -> RpoDigest { + match self { + SmtLeaf::Empty(_) => EMPTY_WORD.into(), + SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]), + SmtLeaf::Multiple(kvs) => { + let elements: Vec = kvs.iter().copied().flat_map(kv_to_elements).collect(); + Rpo256::hash_elements(&elements) + } + } + } + + // ITERATORS + // --------------------------------------------------------------------------------------------- + + /// Returns the key-value pairs in the leaf + pub fn entries(&self) -> Vec<&(RpoDigest, Word)> { + match self { + SmtLeaf::Empty(_) => Vec::new(), + SmtLeaf::Single(kv_pair) => vec![kv_pair], + SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(), + } + } + + // CONVERSIONS + // --------------------------------------------------------------------------------------------- + + /// Converts a leaf to a list of field elements + pub fn to_elements(&self) -> Vec { + self.clone().into_elements() + } + + /// Converts a leaf to a list of field elements + pub fn into_elements(self) -> Vec { + self.into_entries().into_iter().flat_map(kv_to_elements).collect() + } + + /// Converts a leaf the key-value pairs in the leaf + pub fn into_entries(self) -> Vec<(RpoDigest, Word)> { + match self { + SmtLeaf::Empty(_) => Vec::new(), + SmtLeaf::Single(kv_pair) => vec![kv_pair], + SmtLeaf::Multiple(kv_pairs) => kv_pairs, + } + } + + // HELPERS + // --------------------------------------------------------------------------------------------- + + /// Returns the value associated with `key` in the leaf, or `None` if `key` maps to another leaf. + pub(super) fn get_value(&self, key: &RpoDigest) -> Option { + // Ensure that `key` maps to this leaf + if self.index() != key.into() { + return None; + } + + match self { + SmtLeaf::Empty(_) => Some(EMPTY_WORD), + SmtLeaf::Single((key_in_leaf, value_in_leaf)) => { + if key == key_in_leaf { + Some(*value_in_leaf) + } else { + Some(EMPTY_WORD) + } + } + SmtLeaf::Multiple(kv_pairs) => { + for (key_in_leaf, value_in_leaf) in kv_pairs { + if key == key_in_leaf { + return Some(*value_in_leaf); + } + } + + Some(EMPTY_WORD) + } + } + } + + /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if + /// any. + /// + /// The caller needs to ensure that `key` has the same leaf index as all other keys in the leaf + pub(super) fn insert(&mut self, key: RpoDigest, value: Word) -> Option { + match self { + SmtLeaf::Empty(_) => { + *self = SmtLeaf::new_single(key, value); + None + } + SmtLeaf::Single(kv_pair) => { + if kv_pair.0 == key { + // the key is already in this leaf. Update the value and return the previous + // value + let old_value = kv_pair.1; + kv_pair.1 = value; + Some(old_value) + } else { + // Another entry is present in this leaf. Transform the entry into a list + // entry, and make sure the key-value pairs are sorted by key + let mut pairs = vec![*kv_pair, (key, value)]; + pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2)); + + *self = SmtLeaf::Multiple(pairs); + + None + } + } + SmtLeaf::Multiple(kv_pairs) => { + match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { + Ok(pos) => { + let old_value = kv_pairs[pos].1; + kv_pairs[pos].1 = value; + + Some(old_value) + } + Err(pos) => { + kv_pairs.insert(pos, (key, value)); + + None + } + } + } + } + } + + /// Removes key-value pair from the leaf stored at key; returns the previous value associated + /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became + /// empty, and must be removed from the data structure it is contained in. + pub(super) fn remove(&mut self, key: RpoDigest) -> (Option, bool) { + match self { + SmtLeaf::Empty(_) => (None, false), + SmtLeaf::Single((key_at_leaf, value_at_leaf)) => { + if *key_at_leaf == key { + // our key was indeed stored in the leaf, so we return the value that was stored + // in it, and indicate that the leaf should be removed + let old_value = *value_at_leaf; + + // Note: this is not strictly needed, since the caller is expected to drop this + // `SmtLeaf` object. + *self = SmtLeaf::new_empty(key.into()); + + (Some(old_value), true) + } else { + // another key is stored at leaf; nothing to update + (None, false) + } + } + SmtLeaf::Multiple(kv_pairs) => { + match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { + Ok(pos) => { + let old_value = kv_pairs[pos].1; + + kv_pairs.remove(pos); + debug_assert!(!kv_pairs.is_empty()); + + if kv_pairs.len() == 1 { + // convert the leaf into `Single` + *self = SmtLeaf::Single(kv_pairs[0]); + } + + (Some(old_value), false) + } + Err(_) => { + // other keys are stored at leaf; nothing to update + (None, false) + } + } + } + } + } +} + +impl Serializable for SmtLeaf { + fn write_into(&self, target: &mut W) { + // Write: num entries + self.num_entries().write_into(target); + + // Write: leaf index + let leaf_index: u64 = self.index().value(); + leaf_index.write_into(target); + + // Write: entries + for (key, value) in self.entries() { + key.write_into(target); + value.write_into(target); + } + } +} + +impl Deserializable for SmtLeaf { + fn read_from(source: &mut R) -> Result { + // Read: num entries + let num_entries = source.read_u64()?; + + // Read: leaf index + let leaf_index: LeafIndex = { + let value = source.read_u64()?; + LeafIndex::new_max_depth(value) + }; + + // Read: entries + let mut entries: Vec<(RpoDigest, Word)> = Vec::new(); + for _ in 0..num_entries { + let key: RpoDigest = source.read()?; + let value: Word = source.read()?; + + entries.push((key, value)); + } + + Self::new(entries, leaf_index) + .map_err(|err| DeserializationError::InvalidValue(err.to_string())) + } +} + +// HELPER FUNCTIONS +// ================================================================================================ + +/// Converts a key-value tuple to an iterator of `Felt`s +fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { + let key_elements = key.into_iter(); + let value_elements = value.into_iter(); + + key_elements.chain(value_elements) +} + +/// Compares two keys, compared element-by-element using their integer representations starting with +/// the most significant element. +fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering { + for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() { + let v1 = v1.as_int(); + let v2 = v2.as_int(); + if v1 != v2 { + return v1.cmp(&v2); + } + } + + Ordering::Equal +} diff --git a/src/merkle/smt/full/mod.rs b/src/merkle/smt/full/mod.rs index 0dfd66c..deb8122 100644 --- a/src/merkle/smt/full/mod.rs +++ b/src/merkle/smt/full/mod.rs @@ -1,19 +1,23 @@ -use core::cmp::Ordering; - use winter_math::StarkField; use crate::hash::rpo::Rpo256; use crate::merkle::{EmptySubtreeRoots, InnerNodeInfo}; -use crate::utils::{ - collections::{BTreeMap, BTreeSet, Vec}, - vec, -}; +use crate::utils::collections::{BTreeMap, BTreeSet}; use crate::{Felt, EMPTY_WORD}; use super::{ InnerNode, LeafIndex, MerkleError, MerklePath, NodeIndex, RpoDigest, SparseMerkleTree, Word, }; +mod error; +pub use error::{SmtLeafError, SmtProofError}; + +mod leaf; +pub use leaf::SmtLeaf; + +mod proof; +pub use proof::SmtProof; + #[cfg(test)] mod tests; @@ -119,14 +123,14 @@ impl Smt { let leaf_pos = LeafIndex::::from(*key).value(); match self.leaves.get(&leaf_pos) { - Some(leaf) => leaf.get_value(key), + Some(leaf) => leaf.get_value(key).unwrap_or_default(), None => EMPTY_WORD, } } /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle /// path to the leaf, as well as the leaf itself. - pub fn open(&self, key: &RpoDigest) -> (MerklePath, SmtLeaf) { + pub fn open(&self, key: &RpoDigest) -> SmtProof { >::open(self, key) } @@ -208,7 +212,7 @@ impl SparseMerkleTree for Smt { type Key = RpoDigest; type Value = Word; type Leaf = SmtLeaf; - type Opening = (MerklePath, SmtLeaf); + type Opening = SmtProof; const EMPTY_VALUE: Self::Value = EMPTY_WORD; @@ -250,7 +254,7 @@ impl SparseMerkleTree for Smt { match self.leaves.get(&leaf_pos) { Some(leaf) => leaf.clone(), - None => SmtLeaf::Empty, + None => SmtLeaf::new_empty(key.into()), } } @@ -262,6 +266,10 @@ impl SparseMerkleTree for Smt { let most_significant_felt = key[3]; LeafIndex::new_max_depth(most_significant_felt.as_int()) } + + fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof { + SmtProof::new_unchecked(path, leaf) + } } impl Default for Smt { @@ -270,196 +278,24 @@ impl Default for Smt { } } -// LEAF +// CONVERSIONS // ================================================================================================ -#[derive(Clone, Debug, PartialEq, Eq)] -#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] -pub enum SmtLeaf { - Empty, - Single((RpoDigest, Word)), - Multiple(Vec<(RpoDigest, Word)>), -} - -impl SmtLeaf { - /// Converts a leaf to a list of field elements - pub fn to_elements(&self) -> Vec { - self.clone().into_elements() - } - - /// Converts a leaf to a list of field elements - pub fn into_elements(self) -> Vec { - self.into_entries().into_iter().flat_map(kv_to_elements).collect() - } - - /// Returns the key-value pairs in the leaf - pub fn entries(&self) -> Vec<&(RpoDigest, Word)> { - match self { - SmtLeaf::Empty => Vec::new(), - SmtLeaf::Single(kv_pair) => vec![kv_pair], - SmtLeaf::Multiple(kv_pairs) => kv_pairs.iter().collect(), - } - } - - /// Converts a leaf the key-value pairs in the leaf - pub fn into_entries(self) -> Vec<(RpoDigest, Word)> { - match self { - SmtLeaf::Empty => Vec::new(), - SmtLeaf::Single(kv_pair) => vec![kv_pair], - SmtLeaf::Multiple(kv_pairs) => kv_pairs, - } - } - - /// Computes the hash of the leaf - pub fn hash(&self) -> RpoDigest { - match self { - SmtLeaf::Empty => EMPTY_WORD.into(), - SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]), - SmtLeaf::Multiple(kvs) => { - let elements: Vec = kvs.iter().copied().flat_map(kv_to_elements).collect(); - Rpo256::hash_elements(&elements) - } - } - } - - // HELPERS - // --------------------------------------------------------------------------------------------- - - /// Returns the value associated with `key` in the leaf - fn get_value(&self, key: &RpoDigest) -> Word { - match self { - SmtLeaf::Empty => EMPTY_WORD, - SmtLeaf::Single((key_in_leaf, value_in_leaf)) => { - if key == key_in_leaf { - *value_in_leaf - } else { - EMPTY_WORD - } - } - SmtLeaf::Multiple(kv_pairs) => { - for (key_in_leaf, value_in_leaf) in kv_pairs { - if key == key_in_leaf { - return *value_in_leaf; - } - } - - EMPTY_WORD - } - } - } - - /// Inserts key-value pair into the leaf; returns the previous value associated with `key`, if - /// any. - fn insert(&mut self, key: RpoDigest, value: Word) -> Option { - match self { - SmtLeaf::Empty => { - *self = SmtLeaf::Single((key, value)); - None - } - SmtLeaf::Single(kv_pair) => { - if kv_pair.0 == key { - // the key is already in this leaf. Update the value and return the previous - // value - let old_value = kv_pair.1; - kv_pair.1 = value; - Some(old_value) - } else { - // Another entry is present in this leaf. Transform the entry into a list - // entry, and make sure the key-value pairs are sorted by key - let mut pairs = vec![*kv_pair, (key, value)]; - pairs.sort_by(|(key_1, _), (key_2, _)| cmp_keys(*key_1, *key_2)); - - *self = SmtLeaf::Multiple(pairs); - - None - } - } - SmtLeaf::Multiple(kv_pairs) => { - match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { - Ok(pos) => { - let old_value = kv_pairs[pos].1; - kv_pairs[pos].1 = value; - - Some(old_value) - } - Err(pos) => { - kv_pairs.insert(pos, (key, value)); - - None - } - } - } - } - } - - /// Removes key-value pair from the leaf stored at key; returns the previous value associated - /// with `key`, if any. Also returns an `is_empty` flag, indicating whether the leaf became - /// empty, and must be removed from the data structure it is contained in. - fn remove(&mut self, key: RpoDigest) -> (Option, bool) { - match self { - SmtLeaf::Empty => (None, false), - SmtLeaf::Single((key_at_leaf, value_at_leaf)) => { - if *key_at_leaf == key { - // our key was indeed stored in the leaf, so we return the value that was stored - // in it, and indicate that the leaf should be removed - let old_value = *value_at_leaf; - - // Note: this is not strictly needed, since the caller is expected to drop this - // `SmtLeaf` object. - *self = SmtLeaf::Empty; - - (Some(old_value), true) - } else { - // another key is stored at leaf; nothing to update - (None, false) - } - } - SmtLeaf::Multiple(kv_pairs) => { - match kv_pairs.binary_search_by(|kv_pair| cmp_keys(kv_pair.0, key)) { - Ok(pos) => { - let old_value = kv_pairs[pos].1; - - kv_pairs.remove(pos); - debug_assert!(!kv_pairs.is_empty()); - - if kv_pairs.len() == 1 { - // convert the leaf into `Single` - *self = SmtLeaf::Single(kv_pairs[0]); - } - - (Some(old_value), false) - } - Err(_) => { - // other keys are stored at leaf; nothing to update - (None, false) - } - } - } - } +impl From for LeafIndex { + fn from(value: Word) -> Self { + // We use the most significant `Felt` of a `Word` as the leaf index. + Self::new_max_depth(value[3].as_int()) } } -// HELPER FUNCTIONS -// ================================================================================================ - -/// Converts a key-value tuple to an iterator of `Felt`s -fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator { - let key_elements = key.into_iter(); - let value_elements = value.into_iter(); - - key_elements.chain(value_elements) +impl From for LeafIndex { + fn from(value: RpoDigest) -> Self { + Word::from(value).into() + } } -/// Compares two keys, compared element-by-element using their integer representations starting with -/// the most significant element. -fn cmp_keys(key_1: RpoDigest, key_2: RpoDigest) -> Ordering { - for (v1, v2) in key_1.iter().zip(key_2.iter()).rev() { - let v1 = v1.as_int(); - let v2 = v2.as_int(); - if v1 != v2 { - return v1.cmp(&v2); - } +impl From<&RpoDigest> for LeafIndex { + fn from(value: &RpoDigest) -> Self { + Word::from(value).into() } - - Ordering::Equal } diff --git a/src/merkle/smt/full/proof.rs b/src/merkle/smt/full/proof.rs new file mode 100644 index 0000000..ef29679 --- /dev/null +++ b/src/merkle/smt/full/proof.rs @@ -0,0 +1,104 @@ +use crate::utils::string::ToString; +use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; + +use super::{MerklePath, RpoDigest, SmtLeaf, SmtProofError, Word, SMT_DEPTH}; + +/// A proof which can be used to assert membership (or non-membership) of key-value pairs in a +/// [`super::Smt`]. +/// +/// The proof consists of a Merkle path and leaf which describes the node located at the base of the +/// path. +#[derive(PartialEq, Eq, Debug)] +pub struct SmtProof { + path: MerklePath, + leaf: SmtLeaf, +} + +impl SmtProof { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + + /// Returns a new instance of [`SmtProof`] instantiated from the specified path and leaf. + /// + /// # Errors + /// Returns an error if the path length is not [`SMT_DEPTH`]. + pub fn new(path: MerklePath, leaf: SmtLeaf) -> Result { + if path.len() != SMT_DEPTH.into() { + return Err(SmtProofError::InvalidPathLength(path.len())); + } + + Ok(Self { path, leaf }) + } + + /// Returns a new instance of [`SmtProof`] instantiated from the specified path and leaf. + /// + /// The length of the path is not checked. Reserved for internal use. + pub(super) fn new_unchecked(path: MerklePath, leaf: SmtLeaf) -> Self { + Self { path, leaf } + } + + // PROOF VERIFIER + // -------------------------------------------------------------------------------------------- + + /// Returns true if a [`super::Smt`] with the specified root contains the provided + /// key-value pair. + /// + /// Note: this method cannot be used to assert non-membership. That is, if false is returned, + /// it does not mean that the provided key-value pair is not in the tree. + pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool { + let maybe_value_in_leaf = self.leaf.get_value(key); + + match maybe_value_in_leaf { + Some(value_in_leaf) => { + // The value must match for the proof to be valid + if value_in_leaf != *value { + return false; + } + + // make sure the Merkle path resolves to the correct root + self.compute_root() == *root + } + // If the key maps to a different leaf, the proof cannot verify membership of `value` + None => false, + } + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the value associated with the specific key according to this proof, or None if + /// this proof does not contain a value for the specified key. + /// + /// A key-value pair generated by using this method should pass the `verify_membership()` check. + pub fn get(&self, key: &RpoDigest) -> Option { + self.leaf.get_value(key) + } + + /// Computes the root of a [`super::Smt`] to which this proof resolves. + pub fn compute_root(&self) -> RpoDigest { + self.path + .compute_root(self.leaf.index().value(), self.leaf.hash()) + .expect("failed to compute Merkle path root") + } + + /// Consume the proof and returns its parts. + pub fn into_parts(self) -> (MerklePath, SmtLeaf) { + (self.path, self.leaf) + } +} + +impl Serializable for SmtProof { + fn write_into(&self, target: &mut W) { + self.path.write_into(target); + self.leaf.write_into(target); + } +} + +impl Deserializable for SmtProof { + fn read_from(source: &mut R) -> Result { + let path = MerklePath::read_from(source)?; + let leaf = SmtLeaf::read_from(source)?; + + Self::new(path, leaf).map_err(|err| DeserializationError::InvalidValue(err.to_string())) + } +} diff --git a/src/merkle/smt/full/tests.rs b/src/merkle/smt/full/tests.rs index 2567236..08f4183 100644 --- a/src/merkle/smt/full/tests.rs +++ b/src/merkle/smt/full/tests.rs @@ -1,9 +1,15 @@ +use winter_utils::{Deserializable, Serializable}; + use super::*; use crate::{ merkle::{EmptySubtreeRoots, MerkleStore}, + utils::collections::Vec, ONE, WORD_SIZE, }; +// SMT +// -------------------------------------------------------------------------------------------- + /// This test checks that inserting twice at the same key functions as expected. The test covers /// only the case where the key is alone in its leaf #[test] @@ -129,7 +135,7 @@ fn test_smt_insert_and_remove_multiple_values() { assert_eq!(smt.root(), tree_root); let expected_path = store.get_path(tree_root, key_index).unwrap(); - assert_eq!(smt.open(&key).0, expected_path.path); + assert_eq!(smt.open(&key).into_parts().0, expected_path.path); } } let mut smt = Smt::default(); @@ -248,7 +254,7 @@ fn test_smt_removal() { let old_value_1 = smt.insert(key_1, EMPTY_WORD); assert_eq!(old_value_1, value_1); - assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Empty); + assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into())); } } @@ -323,12 +329,65 @@ fn test_smt_entries() { assert_eq!(&(key_2, value_2), entries.next().unwrap()); assert!(entries.next().is_none()); } + +// SMT LEAF +// -------------------------------------------------------------------------------------------- + +#[test] +fn test_empty_smt_leaf_serialization() { + let empty_leaf = SmtLeaf::new_empty(LeafIndex::new_max_depth(42)); + + let mut serialized = empty_leaf.to_bytes(); + // extend buffer with random bytes + serialized.extend([1, 2, 3, 4, 5]); + let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap(); + + assert_eq!(empty_leaf, deserialized); +} + +#[test] +fn test_single_smt_leaf_serialization() { + let single_leaf = SmtLeaf::new_single( + RpoDigest::from([10_u64.into(), 11_u64.into(), 12_u64.into(), 13_u64.into()]), + [1_u64.into(), 2_u64.into(), 3_u64.into(), 4_u64.into()], + ); + + let mut serialized = single_leaf.to_bytes(); + // extend buffer with random bytes + serialized.extend([1, 2, 3, 4, 5]); + let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap(); + + assert_eq!(single_leaf, deserialized); +} + +#[test] +fn test_multiple_smt_leaf_serialization_success() { + let multiple_leaf = SmtLeaf::new_multiple(vec![ + ( + RpoDigest::from([10_u64.into(), 11_u64.into(), 12_u64.into(), 13_u64.into()]), + [1_u64.into(), 2_u64.into(), 3_u64.into(), 4_u64.into()], + ), + ( + RpoDigest::from([100_u64.into(), 101_u64.into(), 102_u64.into(), 13_u64.into()]), + [11_u64.into(), 12_u64.into(), 13_u64.into(), 14_u64.into()], + ), + ]) + .unwrap(); + + let mut serialized = multiple_leaf.to_bytes(); + // extend buffer with random bytes + serialized.extend([1, 2, 3, 4, 5]); + let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap(); + + assert_eq!(multiple_leaf, deserialized); +} + // HELPERS // -------------------------------------------------------------------------------------------- fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest { if value == EMPTY_WORD { - SmtLeaf::Empty.hash() + SmtLeaf::new_empty(key.into()).hash() } else { SmtLeaf::Single((key, value)).hash() } diff --git a/src/merkle/smt/mod.rs b/src/merkle/smt/mod.rs index c8f2ba1..b4f65ac 100644 --- a/src/merkle/smt/mod.rs +++ b/src/merkle/smt/mod.rs @@ -1,5 +1,3 @@ -use winter_math::StarkField; - use crate::{ hash::rpo::{Rpo256, RpoDigest}, Word, @@ -8,7 +6,7 @@ use crate::{ use super::{EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Vec}; mod full; -pub use full::{Smt, SmtLeaf, SMT_DEPTH}; +pub use full::{Smt, SmtLeaf, SmtLeafError, SMT_DEPTH}; mod simple; pub use simple::SimpleSmt; @@ -52,7 +50,7 @@ pub(crate) trait SparseMerkleTree { /// The type for a leaf type Leaf; /// The type for an opening (i.e. a "proof") of a leaf - type Opening: From<(MerklePath, Self::Leaf)>; + type Opening; /// The default value used to compute the hash of empty leaves const EMPTY_VALUE: Self::Value; @@ -83,7 +81,7 @@ pub(crate) trait SparseMerkleTree { MerklePath::new(path) }; - (merkle_path, leaf).into() + Self::path_and_leaf_to_opening(merkle_path, leaf) } /// Inserts a value at the specified key, returning the previous value associated with that key. @@ -170,6 +168,11 @@ pub(crate) trait SparseMerkleTree { /// Maps a key to a leaf index fn key_to_leaf_index(key: &Self::Key) -> LeafIndex; + + /// Maps a (MerklePath, Self::Leaf) to an opening. + /// + /// The length `path` is guaranteed to be equal to `DEPTH` + fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening; } // INNER NODE @@ -240,16 +243,3 @@ impl TryFrom for LeafIndex { Self::new(node_index.value()) } } - -impl From for LeafIndex { - fn from(value: Word) -> Self { - // We use the most significant `Felt` of a `Word` as the leaf index. - Self::new_max_depth(value[3].as_int()) - } -} - -impl From for LeafIndex { - fn from(value: RpoDigest) -> Self { - Word::from(value).into() - } -} diff --git a/src/merkle/smt/simple/mod.rs b/src/merkle/smt/simple/mod.rs index 9ccc0f6..b0a4ab9 100644 --- a/src/merkle/smt/simple/mod.rs +++ b/src/merkle/smt/simple/mod.rs @@ -1,5 +1,5 @@ use crate::{ - merkle::{EmptySubtreeRoots, InnerNodeInfo, ValuePath}, + merkle::{EmptySubtreeRoots, InnerNodeInfo, MerklePath, ValuePath}, EMPTY_WORD, }; @@ -302,4 +302,8 @@ impl SparseMerkleTree for SimpleSmt { fn key_to_leaf_index(key: &LeafIndex) -> LeafIndex { *key } + + fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath { + (path, leaf).into() + } }