Skip to content

Commit

Permalink
Introduce key_to_leaf_index
Browse files Browse the repository at this point in the history
  • Loading branch information
@plafer
plafer committed Jan 16, 2024
1 parent 81e5d2a commit c7149d5
Show file tree
Hide file tree
Showing 5 changed files with 64 additions and 79 deletions.
2 changes: 1 addition & 1 deletion src/merkle/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ pub(crate) mod sparse_merkle_tree;
pub use sparse_merkle_tree::{LeafIndex, SMT_MAX_DEPTH, SMT_MIN_DEPTH};

mod smt;
pub use smt::{Smt, SmtKey, SmtLeaf, SMT_DEPTH};
pub use smt::{Smt, SmtLeaf, SMT_DEPTH};

mod tiered_smt;
pub use tiered_smt::{TieredSmt, TieredSmtProof, TieredSmtProofError};
Expand Down
4 changes: 4 additions & 0 deletions src/merkle/simple_smt/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,10 @@ impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
// `SimpleSmt` takes the leaf value itself as the hash
leaf.into()
}

fn key_to_leaf_index(leaf: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
*leaf
}
}

// TRY APPLY DIFF
Expand Down
75 changes: 24 additions & 51 deletions src/merkle/smt/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ impl Smt {
/// - The number of entries exceeds 2^63 entries.
/// - The provided entries contain multiple values for the same key.
pub fn with_entries(
entries: impl IntoIterator<Item = (SmtKey, Word)>,
entries: impl IntoIterator<Item = (RpoDigest, Word)>,
) -> Result<Self, MerkleError> {
// create an empty tree
let mut tree = Self::new();
Expand Down Expand Up @@ -91,7 +91,7 @@ impl Smt {
}

/// Returns the leaf at the specified index.
pub fn get_leaf(&self, key: &SmtKey) -> SmtLeaf {
pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
<Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
}

Expand All @@ -103,7 +103,7 @@ impl Smt {
/// Returns a Merkle path from the leaf node specified by the key to the root.
///
/// The node itself is not included in the path.
pub fn open(&self, key: SmtKey) -> (MerklePath, SmtLeaf) {
pub fn open(&self, key: RpoDigest) -> (MerklePath, SmtLeaf) {
<Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
}

Expand All @@ -116,7 +116,7 @@ impl Smt {
///
/// This also recomputes all hashes between the leaf (associated with the key) and the root,
/// updating the root itself.
pub fn insert(&mut self, key: SmtKey, value: Word) -> Word {
pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
<Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
}

Expand All @@ -125,10 +125,10 @@ impl Smt {

/// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
/// value, such that this is indeed an insertion.
fn perform_insert(&mut self, key: SmtKey, value: Word) -> Option<Word> {
fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
debug_assert_ne!(value, Self::EMPTY_VALUE);

let leaf_index: LeafIndex<SMT_DEPTH> = key.into();
let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);

match self.leaves.get_mut(&leaf_index.value()) {
Some(leaf) => leaf.insert(key, value),
Expand All @@ -141,8 +141,8 @@ impl Smt {
}

/// Removes key-value pair at leaf index pointed to by `key` if it exists.
fn perform_remove(&mut self, key: SmtKey) -> Option<Word> {
let leaf_index: LeafIndex<SMT_DEPTH> = key.into();
fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);

if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
let (old_value, is_empty) = leaf.remove(key);
Expand All @@ -158,7 +158,7 @@ impl Smt {
}

impl SparseMerkleTree<SMT_DEPTH> for Smt {
type Key = SmtKey;
type Key = RpoDigest;
type Value = Word;
type Leaf = SmtLeaf;
type Opening = (MerklePath, SmtLeaf);
Expand Down Expand Up @@ -194,7 +194,7 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
}
}

fn get_leaf(&self, key: &Self::Key) -> Self::Leaf {
fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();

match self.leaves.get(&leaf_pos) {
Expand All @@ -206,6 +206,11 @@ impl SparseMerkleTree<SMT_DEPTH> for Smt {
fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
leaf.hash()
}

fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
let most_significant_felt = key[3];
LeafIndex::new_max_depth(most_significant_felt.as_int())
}
}

impl Default for Smt {
Expand All @@ -214,46 +219,14 @@ impl Default for Smt {
}
}

// SMT KEY
// ================================================================================================

/// Represents a key (256 bits) for the Smt.
///
/// The most significant `u64` determines the corresponding leaf index when inserting values into
/// the Smt.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct SmtKey {
word: RpoDigest,
}

impl From<SmtKey> for LeafIndex<SMT_DEPTH> {
fn from(key: SmtKey) -> Self {
let most_significant_felt = key.word[3];
Self::new_max_depth(most_significant_felt.as_int())
}
}

impl From<RpoDigest> for SmtKey {
fn from(digest: RpoDigest) -> Self {
Self { word: digest }
}
}

impl From<Word> for SmtKey {
fn from(word: Word) -> Self {
Self { word: word.into() }
}
}

// LEAF
// ================================================================================================

#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub enum SmtLeaf {
Single((SmtKey, Word)),
Multiple(Vec<(SmtKey, Word)>),
Single((RpoDigest, Word)),
Multiple(Vec<(RpoDigest, Word)>),
}

impl SmtLeaf {
Expand All @@ -273,7 +246,7 @@ impl SmtLeaf {
/// Compute the hash of the leaf
pub fn hash(&self) -> RpoDigest {
match self {
SmtLeaf::Single((key, value)) => Rpo256::merge(&[key.word, value.into()]),
SmtLeaf::Single((key, value)) => Rpo256::merge(&[*key, value.into()]),
SmtLeaf::Multiple(kvs) => {
let elements: Vec<Felt> = kvs.iter().copied().flat_map(kv_to_elements).collect();
Rpo256::hash_elements(&elements)
Expand All @@ -286,7 +259,7 @@ impl SmtLeaf {

/// Insert key-value pair into the leaf; return the previous value associated with `key`, if
/// any.
fn insert(&mut self, key: SmtKey, value: Word) -> Option<Word> {
fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
match self {
SmtLeaf::Single(kv_pair) => {
if kv_pair.0 == key {
Expand Down Expand Up @@ -327,7 +300,7 @@ impl SmtLeaf {
/// Remove key-value pair into the leaf stored at key; return the previous value associated with
/// `key`, if any. We also return an `is_empty` flag which indicates whether the leaf became
/// empty, and must be removed from the data structure it is contained in.
fn remove(&mut self, key: SmtKey) -> (Option<Word>, bool) {
fn remove(&mut self, key: RpoDigest) -> (Option<Word>, bool) {
match self {
SmtLeaf::Single((key_at_leaf, value_at_leaf)) => {
if *key_at_leaf == key {
Expand Down Expand Up @@ -370,17 +343,17 @@ impl SmtLeaf {
// ================================================================================================

/// Converts a key-value tuple to an iterator of `Felt`s
fn kv_to_elements((key, value): (SmtKey, Word)) -> impl Iterator<Item = Felt> {
let key_elements = key.word.into_iter();
fn kv_to_elements((key, value): (RpoDigest, Word)) -> impl Iterator<Item = Felt> {
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: SmtKey, key_2: SmtKey) -> Ordering {
for (v1, v2) in key_1.word.iter().zip(key_2.word.iter()).rev() {
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 {
Expand Down
46 changes: 23 additions & 23 deletions src/merkle/smt/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@ fn test_smt_insert_at_same_key() {

assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));

let key_1: SmtKey = {
let key_1: RpoDigest = {
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;

RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into()
RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
};
let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();

Expand Down Expand Up @@ -48,10 +48,10 @@ fn test_smt_insert_at_same_key() {
/// only the case where the leaf type is `SmtLeaf::Multiple`
#[test]
fn test_smt_insert_at_same_key_2() {
let key_already_present: SmtKey = {
let key_already_present: RpoDigest = {
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;

RpoDigest::from([ONE + ONE, ONE + ONE, ONE + ONE, Felt::new(raw)]).into()
RpoDigest::from([ONE + ONE, ONE + ONE, ONE + ONE, Felt::new(raw)])
};
let key_already_present_index: NodeIndex =
LeafIndex::<SMT_DEPTH>::from(key_already_present).into();
Expand All @@ -69,10 +69,10 @@ fn test_smt_insert_at_same_key_2() {
store
};

let key_1: SmtKey = {
let key_1: RpoDigest = {
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;

RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into()
RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
};
let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();

Expand Down Expand Up @@ -120,22 +120,22 @@ fn test_smt_insert_multiple_values() {

assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));

let key_1: SmtKey = {
let key_1: RpoDigest = {
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;

RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into()
RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
};

let key_2: SmtKey = {
let key_2: RpoDigest = {
let raw = 0b_11111111_11111111_11111111_11111111_11111111_11111111_11111111_11111111_u64;

RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into()
RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
};

let key_3: SmtKey = {
let key_3: RpoDigest = {
let raw = 0b_00000000_00000000_00000000_00000000_00000000_00000000_00000000_00000000_u64;

RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into()
RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
};

let value_1 = [ONE; WORD_SIZE];
Expand Down Expand Up @@ -169,11 +169,11 @@ fn test_smt_removal() {

let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;

let key_1: SmtKey = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into();
let key_2: SmtKey =
RpoDigest::from([2_u64.into(), 2_u64.into(), 2_u64.into(), Felt::new(raw)]).into();
let key_3: SmtKey =
RpoDigest::from([3_u64.into(), 3_u64.into(), 3_u64.into(), Felt::new(raw)]).into();
let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
let key_2: RpoDigest =
RpoDigest::from([2_u64.into(), 2_u64.into(), 2_u64.into(), Felt::new(raw)]);
let key_3: RpoDigest =
RpoDigest::from([3_u64.into(), 3_u64.into(), 3_u64.into(), Felt::new(raw)]);

let value_1 = [ONE; WORD_SIZE];
let value_2 = [2_u64.into(); WORD_SIZE];
Expand Down Expand Up @@ -242,9 +242,9 @@ fn test_smt_removal() {
fn test_smt_path_to_keys_in_same_leaf_are_equal() {
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;

let key_1: SmtKey = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]).into();
let key_2: SmtKey =
RpoDigest::from([2_u64.into(), 2_u64.into(), 2_u64.into(), Felt::new(raw)]).into();
let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
let key_2: RpoDigest =
RpoDigest::from([2_u64.into(), 2_u64.into(), 2_u64.into(), Felt::new(raw)]);

let value_1 = [ONE; WORD_SIZE];
let value_2 = [2_u64.into(); WORD_SIZE];
Expand All @@ -257,15 +257,15 @@ fn test_smt_path_to_keys_in_same_leaf_are_equal() {
// HELPERS
// --------------------------------------------------------------------------------------------

fn build_single_leaf_node(key: SmtKey, value: Word) -> RpoDigest {
fn build_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest {
SmtLeaf::Single((key, value)).hash()
}

fn build_multiple_leaf_node(kv_pairs: &[(SmtKey, Word)]) -> RpoDigest {
fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest {
let elements: Vec<Felt> = kv_pairs
.iter()
.flat_map(|(key, value)| {
let key_elements = key.word.into_iter();
let key_elements = key.into_iter();
let value_elements = (*value).into_iter();

key_elements.chain(value_elements)
Expand Down
16 changes: 12 additions & 4 deletions src/merkle/sparse_merkle_tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,8 @@ pub const SMT_MAX_DEPTH: u8 = 64;
///
/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
/// The type for a key, which must be convertible into a `u64` infaillibly
type Key: Into<LeafIndex<DEPTH>> + Clone;
/// The type for a key
type Key: Clone;
/// The type for a value
type Value: Clone + PartialEq;
/// The type for a leaf
Expand All @@ -61,7 +61,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
let leaf = self.get_leaf(&key);

let mut index: NodeIndex = {
let leaf_index: LeafIndex<DEPTH> = key.into();
let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
leaf_index.into()
};

Expand Down Expand Up @@ -97,7 +97,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {

let leaf = self.get_leaf(&key);
let node_index = {
let leaf_index: LeafIndex<DEPTH> = key.into();
let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
leaf_index.into()
};

Expand Down Expand Up @@ -148,6 +148,8 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {

/// Returns the hash of a leaf
fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;

fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
}

// INNER NODE
Expand Down Expand Up @@ -225,3 +227,9 @@ impl From<Word> for LeafIndex<SMT_MAX_DEPTH> {
Self::new_max_depth(value[3].as_int())
}
}

impl From<RpoDigest> for LeafIndex<SMT_MAX_DEPTH> {
fn from(value: RpoDigest) -> Self {
Word::from(value).into()
}
}

0 comments on commit c7149d5

Please sign in to comment.