Browse Source

Merge pull request #32 from 0xPolygonMiden/next

Tracking PR for next release
al-gkr-basic-workflow
Bobbin Threadbare 2 years ago
committed by GitHub
parent
commit
3c60484e21
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 794 additions and 24 deletions
  1. +5
    -1
      Cargo.toml
  2. +6
    -4
      README.md
  3. +49
    -0
      benches/README.md
  4. +61
    -2
      benches/hash.rs
  5. +84
    -0
      benches/smt.rs
  6. +12
    -8
      src/hash/blake/mod.rs
  7. +4
    -1
      src/lib.rs
  8. +7
    -6
      src/merkle/merkle_tree.rs
  9. +34
    -2
      src/merkle/mod.rs
  10. +269
    -0
      src/merkle/simple_smt/mod.rs
  11. +263
    -0
      src/merkle/simple_smt/tests.rs

+ 5
- 1
Cargo.toml

@ -14,6 +14,10 @@ edition = "2021"
name = "hash" name = "hash"
harness = false harness = false
[[bench]]
name = "smt"
harness = false
[features] [features]
default = ["blake3/default", "std", "winter_crypto/default", "winter_math/default", "winter_utils/default"] default = ["blake3/default", "std", "winter_crypto/default", "winter_math/default", "winter_utils/default"]
std = ["blake3/std", "winter_crypto/std", "winter_math/std", "winter_utils/std"] std = ["blake3/std", "winter_crypto/std", "winter_math/std", "winter_utils/std"]
@ -25,6 +29,6 @@ winter_math = { version = "0.4.1", package = "winter-math", default-features = f
winter_utils = { version = "0.4.1", package = "winter-utils", default-features = false } winter_utils = { version = "0.4.1", package = "winter-utils", default-features = false }
[dev-dependencies] [dev-dependencies]
criterion = "0.4"
criterion = { version = "0.4", features = ["html_reports"] }
proptest = "1.0.0" proptest = "1.0.0"
rand_utils = { version = "0.4", package = "winter-rand-utils" } rand_utils = { version = "0.4", package = "winter-rand-utils" }

+ 6
- 4
README.md

@ -2,24 +2,26 @@
This crate contains cryptographic primitives used in Polygon Miden. This crate contains cryptographic primitives used in Polygon Miden.
## Hash ## Hash
[Hash module](./src/hash) provides a set of cryptographic hash functions which are used by Miden VM and Miden Rollup. Currently, these functions are:
[Hash module](./src/hash) provides a set of cryptographic hash functions which are used by the Miden VM and the Miden rollup. Currently, these functions are:
* [BLAKE3](https://github.com/BLAKE3-team/BLAKE3) hash function with 256-bit, 192-bit, or 160-bit output. The 192-bit and 160-bit outputs are obtained by truncating the 256-bit output of the standard BLAKE3. * [BLAKE3](https://github.com/BLAKE3-team/BLAKE3) hash function with 256-bit, 192-bit, or 160-bit output. The 192-bit and 160-bit outputs are obtained by truncating the 256-bit output of the standard BLAKE3.
* [RPO](https://eprint.iacr.org/2022/1577) hash function with 256-bit output. This hash function is an algebraic hash function suitable for recursive STARKs. * [RPO](https://eprint.iacr.org/2022/1577) hash function with 256-bit output. This hash function is an algebraic hash function suitable for recursive STARKs.
For performance benchmarks of these hash functions and their comparison to other popular hash functions please see [here](./benches/).
## Merkle ## Merkle
[Merkle module](./src/merkle/) provides a set of data structures related to Merkle tree. All these data structures are implemented using RPO hash function described above. The data structure are:
[Merkle module](./src/merkle/) provides a set of data structures related to Merkle trees. All these data structures are implemented using the RPO hash function described above. The data structures are:
* `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64. * `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64.
* `MerklePathSet`: a collection of Merkle authentication paths all resolving to the same root. The length of the paths can be at most 64. * `MerklePathSet`: a collection of Merkle authentication paths all resolving to the same root. The length of the paths can be at most 64.
## Crate features ## Crate features
This carate can be compiled with the following features:
This crate can be compiled with the following features:
* `std` - enabled by default and relies on the Rust standard library. * `std` - enabled by default and relies on the Rust standard library.
* `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly. * `no_std` does not rely on the Rust standard library and enables compilation to WebAssembly.
Both of these features imply use of [alloc](https://doc.rust-lang.org/alloc/) to support heap-allocated collections.
Both of these features imply the use of [alloc](https://doc.rust-lang.org/alloc/) to support heap-allocated collections.
To compile with `no_std`, disable default features via `--no-default-features` flag. To compile with `no_std`, disable default features via `--no-default-features` flag.

+ 49
- 0
benches/README.md

@ -0,0 +1,49 @@
# Miden VM Hash Functions
In the Miden VM, we make use of different hash functions. Some of these are "traditional" hash functions, like `BLAKE3`, which are optimized for out-of-STARK performance, while others are algebraic hash functions, like `Rescue Prime`, and are more optimized for a better performance inside the STARK. In what follows, we benchmark several such hash functions and compare against other constructions that are used by other proving systems. More precisely, we benchmark:
* **BLAKE3** as specified [here](https://github.com/BLAKE3-team/BLAKE3-specs/blob/master/blake3.pdf) and implemented [here](https://github.com/BLAKE3-team/BLAKE3) (with a wrapper exposed via this crate).
* **SHA3** as specified [here](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.202.pdf) and implemented [here](https://github.com/novifinancial/winterfell/blob/46dce1adf0/crypto/src/hash/sha/mod.rs).
* **Poseidon** as specified [here](https://eprint.iacr.org/2019/458.pdf) and implemented [here](https://github.com/mir-protocol/plonky2/blob/806b88d7d6e69a30dc0b4775f7ba275c45e8b63b/plonky2/src/hash/poseidon_goldilocks.rs) (but in pure Rust, without vectorized instructions).
* **Rescue Prime (RP)** as specified [here](https://eprint.iacr.org/2020/1143) and implemented [here](https://github.com/novifinancial/winterfell/blob/46dce1adf0/crypto/src/hash/rescue/rp64_256/mod.rs).
* **Rescue Prime Optimized (RPO)** as specified [here](https://eprint.iacr.org/2022/1577) and implemented in this crate.
## Comparison and Instructions
### Comparison
We benchmark the above hash functions using two scenarios. The first is a 2-to-1 $(a,b)\mapsto h(a,b)$ hashing where both $a$, $b$ and $h(a,b)$ are the digests corresponding to each of the hash functions.
The second scenario is that of sequential hashing where we take a sequence of length $100$ field elements and hash these to produce a single digest. The digests are $4$ field elements in a prime field with modulus $2^{64} - 2^{32} + 1$ (i.e., 32 bytes) for Poseidon, Rescue Prime and RPO, and an array `[u8; 32]` for SHA3 and BLAKE3.
#### Scenario 1: 2-to-1 hashing `h(a,b)`
| Function | BLAKE3 | SHA3 | Poseidon | Rp64_256 | RPO_256 |
| ------------------- | ------ | --------| --------- | --------- | ------- |
| Apple M1 Pro | 80 ns | 245 ns | 1.5 us | 9.1 us | 5.4 us |
| Apple M2 | 76 ns | 233 ns | 1.3 us | 7.9 us | 5.0 us |
| Amazon Graviton 3 | 116 ns | | | | 8.8 us |
| AMD Ryzen 9 5950X | 64 ns | 273 ns | 1.2 us | 9.1 us | 5.5 us |
| Intel Core i5-8279U | 80 ns | | | | 8.7 us |
| Intel Xeon 8375C | 67 ns | | | | 8.2 us |
#### Scenario 2: Sequential hashing of 100 elements `h([a_0,...,a_99])`
| Function | BLAKE3 | SHA3 | Poseidon | Rp64_256 | RPO_256 |
| ------------------- | -------| ------- | --------- | --------- | ------- |
| Apple M1 Pro | 1.1 us | 1.5 us | 19.4 us | 118 us | 70 us |
| Apple M2 | 1.0 us | 1.5 us | 17.4 us | 103 us | 65 us |
| Amazon Graviton 3 | 1.4 us | | | | 114 us |
| AMD Ryzen 9 5950X | 0.8 us | 1.7 us | 15.7 us | 120 us | 72 us |
| Intel Core i5-8279U | 1.0 us | | | | 116 us |
| Intel Xeon 8375C | 0.8 ns | | | | 110 us |
### Instructions
Before you can run the benchmarks, you'll need to make sure you have Rust [installed](https://www.rust-lang.org/tools/install). After that, to run the benchmarks for RPO and BLAKE3, clone the current repository, and from the root directory of the repo run the following:
```
cargo bench hash
```
To run the benchmarks for Rescue Prime, Poseidon and SHA3, clone the following [repository](https://github.com/Dominik1999/winterfell.git) as above, then checkout the `hash-functions-benches` branch, and from the root directory of the repo run the following:
```
cargo bench hash
```

+ 61
- 2
benches/hash.rs

@ -1,9 +1,13 @@
use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion}; use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
use miden_crypto::{ use miden_crypto::{
hash::rpo::{Rpo256, RpoDigest},
hash::{
blake::Blake3_256,
rpo::{Rpo256, RpoDigest},
},
Felt, Felt,
}; };
use rand_utils::rand_value; use rand_utils::rand_value;
use winter_crypto::Hasher;
fn rpo256_2to1(c: &mut Criterion) { fn rpo256_2to1(c: &mut Criterion) {
let v: [RpoDigest; 2] = [Rpo256::hash(&[1_u8]), Rpo256::hash(&[2_u8])]; let v: [RpoDigest; 2] = [Rpo256::hash(&[1_u8]), Rpo256::hash(&[2_u8])];
@ -53,5 +57,60 @@ fn rpo256_sequential(c: &mut Criterion) {
}); });
} }
criterion_group!(hash_group, rpo256_sequential, rpo256_2to1);
fn blake3_2to1(c: &mut Criterion) {
let v: [<Blake3_256 as Hasher>::Digest; 2] =
[Blake3_256::hash(&[1_u8]), Blake3_256::hash(&[2_u8])];
c.bench_function("Blake3 2-to-1 hashing (cached)", |bench| {
bench.iter(|| Blake3_256::merge(black_box(&v)))
});
c.bench_function("Blake3 2-to-1 hashing (random)", |bench| {
bench.iter_batched(
|| {
[
Blake3_256::hash(&rand_value::<u64>().to_le_bytes()),
Blake3_256::hash(&rand_value::<u64>().to_le_bytes()),
]
},
|state| Blake3_256::merge(&state),
BatchSize::SmallInput,
)
});
}
fn blake3_sequential(c: &mut Criterion) {
let v: [Felt; 100] = (0..100)
.into_iter()
.map(Felt::new)
.collect::<Vec<Felt>>()
.try_into()
.expect("should not fail");
c.bench_function("Blake3 sequential hashing (cached)", |bench| {
bench.iter(|| Blake3_256::hash_elements(black_box(&v)))
});
c.bench_function("Blake3 sequential hashing (random)", |bench| {
bench.iter_batched(
|| {
let v: [Felt; 100] = (0..100)
.into_iter()
.map(|_| Felt::new(rand_value()))
.collect::<Vec<Felt>>()
.try_into()
.expect("should not fail");
v
},
|state| Blake3_256::hash_elements(&state),
BatchSize::SmallInput,
)
});
}
criterion_group!(
hash_group,
rpo256_2to1,
rpo256_sequential,
blake3_2to1,
blake3_sequential
);
criterion_main!(hash_group); criterion_main!(hash_group);

+ 84
- 0
benches/smt.rs

@ -0,0 +1,84 @@
use core::mem::swap;
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use miden_crypto::{merkle::SimpleSmt, Felt, Word};
use rand_utils::prng_array;
fn smt_rpo(c: &mut Criterion) {
// setup trees
let mut seed = [0u8; 32];
let mut trees = vec![];
for depth in 14..=20 {
let leaves = ((1 << depth) - 1) as u64;
for count in [1, leaves / 2, leaves] {
let entries: Vec<_> = (0..count)
.map(|i| {
let word = generate_word(&mut seed);
(i, word)
})
.collect();
let tree = SimpleSmt::new(entries, depth).unwrap();
trees.push(tree);
}
}
let leaf = generate_word(&mut seed);
// benchmarks
let mut insert = c.benchmark_group(format!("smt update_leaf"));
for tree in trees.iter_mut() {
let depth = tree.depth();
let count = tree.leaves_count() as u64;
let key = count >> 2;
insert.bench_with_input(
format!("simple smt(depth:{depth},count:{count})"),
&(key, leaf),
|b, (key, leaf)| {
b.iter(|| {
tree.update_leaf(black_box(*key), black_box(*leaf)).unwrap();
});
},
);
}
insert.finish();
let mut path = c.benchmark_group(format!("smt get_leaf_path"));
for tree in trees.iter_mut() {
let depth = tree.depth();
let count = tree.leaves_count() as u64;
let key = count >> 2;
path.bench_with_input(
format!("simple smt(depth:{depth},count:{count})"),
&key,
|b, key| {
b.iter(|| {
tree.get_leaf_path(black_box(*key)).unwrap();
});
},
);
}
path.finish();
}
criterion_group!(smt_group, smt_rpo);
criterion_main!(smt_group);
// HELPER FUNCTIONS
// --------------------------------------------------------------------------------------------
fn generate_word(seed: &mut [u8; 32]) -> Word {
swap(seed, &mut prng_array(*seed));
let nums: [u64; 4] = prng_array(*seed);
[
Felt::new(nums[0]),
Felt::new(nums[1]),
Felt::new(nums[2]),
Felt::new(nums[3]),
]
}

+ 12
- 8
src/hash/blake/mod.rs

@ -1,5 +1,7 @@
use super::{Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField}; use super::{Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField};
use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
use crate::utils::{
uninit_vector, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
};
use core::{ use core::{
mem::{size_of, transmute, transmute_copy}, mem::{size_of, transmute, transmute_copy},
ops::Deref, ops::Deref,
@ -276,13 +278,15 @@ where
let digest = if Felt::IS_CANONICAL { let digest = if Felt::IS_CANONICAL {
blake3::hash(E::elements_as_bytes(elements)) blake3::hash(E::elements_as_bytes(elements))
} else { } else {
E::as_base_elements(elements)
.iter()
.fold(blake3::Hasher::new(), |mut hasher, felt| {
hasher.update(&felt.as_int().to_le_bytes());
hasher
})
.finalize()
let base_elements = E::as_base_elements(elements);
let blen = base_elements.len() << 3;
let mut bytes = unsafe { uninit_vector(blen) };
for (idx, element) in base_elements.iter().enumerate() {
bytes[idx * 8..(idx + 1) * 8].copy_from_slice(&element.as_int().to_le_bytes());
}
blake3::hash(&bytes)
}; };
*shrink_bytes(&digest.into()) *shrink_bytes(&digest.into())
} }

+ 4
- 1
src/lib.rs

@ -23,11 +23,14 @@ pub mod utils {
// ================================================================================================ // ================================================================================================
/// A group of four field elements in the Miden base field. /// A group of four field elements in the Miden base field.
pub type Word = [Felt; 4];
pub type Word = [Felt; WORD_SIZE];
// CONSTANTS // CONSTANTS
// ================================================================================================ // ================================================================================================
/// Number of field elements in a word.
pub const WORD_SIZE: usize = 4;
/// Field element representing ZERO in the Miden base filed. /// Field element representing ZERO in the Miden base filed.
pub const ZERO: Felt = Felt::ZERO; pub const ZERO: Felt = Felt::ZERO;

+ 7
- 6
src/merkle/merkle_tree.rs

@ -1,4 +1,4 @@
use super::{Digest, Felt, MerkleError, Rpo256, Vec, Word};
use super::{Felt, MerkleError, Rpo256, RpoDigest, Vec, Word};
use crate::{utils::uninit_vector, FieldElement}; use crate::{utils::uninit_vector, FieldElement};
use core::slice; use core::slice;
use winter_math::log2; use winter_math::log2;
@ -22,7 +22,7 @@ impl MerkleTree {
pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> { pub fn new(leaves: Vec<Word>) -> Result<Self, MerkleError> {
let n = leaves.len(); let n = leaves.len();
if n <= 1 { if n <= 1 {
return Err(MerkleError::DepthTooSmall);
return Err(MerkleError::DepthTooSmall(n as u32));
} else if !n.is_power_of_two() { } else if !n.is_power_of_two() {
return Err(MerkleError::NumLeavesNotPowerOfTwo(n)); return Err(MerkleError::NumLeavesNotPowerOfTwo(n));
} }
@ -35,7 +35,8 @@ impl MerkleTree {
nodes[n..].copy_from_slice(&leaves); nodes[n..].copy_from_slice(&leaves);
// re-interpret nodes as an array of two nodes fused together // re-interpret nodes as an array of two nodes fused together
let two_nodes = unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [Digest; 2], n) };
let two_nodes =
unsafe { slice::from_raw_parts(nodes.as_ptr() as *const [RpoDigest; 2], n) };
// calculate all internal tree nodes // calculate all internal tree nodes
for i in (1..n).rev() { for i in (1..n).rev() {
@ -68,7 +69,7 @@ impl MerkleTree {
/// * The specified index not valid for the specified depth. /// * The specified index not valid for the specified depth.
pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> { pub fn get_node(&self, depth: u32, index: u64) -> Result<Word, MerkleError> {
if depth == 0 { if depth == 0 {
return Err(MerkleError::DepthTooSmall);
return Err(MerkleError::DepthTooSmall(depth));
} else if depth > self.depth() { } else if depth > self.depth() {
return Err(MerkleError::DepthTooBig(depth)); return Err(MerkleError::DepthTooBig(depth));
} }
@ -89,7 +90,7 @@ impl MerkleTree {
/// * The specified index not valid for the specified depth. /// * The specified index not valid for the specified depth.
pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> { pub fn get_path(&self, depth: u32, index: u64) -> Result<Vec<Word>, MerkleError> {
if depth == 0 { if depth == 0 {
return Err(MerkleError::DepthTooSmall);
return Err(MerkleError::DepthTooSmall(depth));
} else if depth > self.depth() { } else if depth > self.depth() {
return Err(MerkleError::DepthTooBig(depth)); return Err(MerkleError::DepthTooBig(depth));
} }
@ -123,7 +124,7 @@ impl MerkleTree {
let n = self.nodes.len() / 2; let n = self.nodes.len() / 2;
let two_nodes = let two_nodes =
unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [Digest; 2], n) };
unsafe { slice::from_raw_parts(self.nodes.as_ptr() as *const [RpoDigest; 2], n) };
for _ in 0..depth { for _ in 0..depth {
index /= 2; index /= 2;

+ 34
- 2
src/merkle/mod.rs

@ -1,8 +1,9 @@
use super::{ use super::{
hash::rpo::{Rpo256, RpoDigest as Digest},
hash::rpo::{Rpo256, RpoDigest},
utils::collections::{BTreeMap, Vec}, utils::collections::{BTreeMap, Vec},
Felt, Word, ZERO, Felt, Word, ZERO,
}; };
use core::fmt;
mod merkle_tree; mod merkle_tree;
pub use merkle_tree::MerkleTree; pub use merkle_tree::MerkleTree;
@ -10,20 +11,51 @@ pub use merkle_tree::MerkleTree;
mod merkle_path_set; mod merkle_path_set;
pub use merkle_path_set::MerklePathSet; pub use merkle_path_set::MerklePathSet;
mod simple_smt;
pub use simple_smt::SimpleSmt;
// ERRORS // ERRORS
// ================================================================================================ // ================================================================================================
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub enum MerkleError { pub enum MerkleError {
DepthTooSmall,
DepthTooSmall(u32),
DepthTooBig(u32), DepthTooBig(u32),
NumLeavesNotPowerOfTwo(usize), NumLeavesNotPowerOfTwo(usize),
InvalidIndex(u32, u64), InvalidIndex(u32, u64),
InvalidDepth(u32, u32), InvalidDepth(u32, u32),
InvalidPath(Vec<Word>), InvalidPath(Vec<Word>),
InvalidEntriesCount(usize, usize),
NodeNotInSet(u64), NodeNotInSet(u64),
} }
impl fmt::Display for MerkleError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
use MerkleError::*;
match self {
DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"),
DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"),
NumLeavesNotPowerOfTwo(leaves) => {
write!(f, "the leaves count {leaves} is not a power of 2")
}
InvalidIndex(depth, index) => write!(
f,
"the leaf index {index} is not valid for the depth {depth}"
),
InvalidDepth(expected, provided) => write!(
f,
"the provided depth {provided} is not valid for {expected}"
),
InvalidPath(_path) => write!(f, "the provided path is not valid"),
InvalidEntriesCount(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"),
}
}
}
#[cfg(feature = "std")]
impl std::error::Error for MerkleError {}
// HELPER FUNCTIONS // HELPER FUNCTIONS
// ================================================================================================ // ================================================================================================

+ 269
- 0
src/merkle/simple_smt/mod.rs

@ -0,0 +1,269 @@
use super::{BTreeMap, MerkleError, Rpo256, RpoDigest, Vec, Word};
#[cfg(test)]
mod tests;
// SPARSE MERKLE TREE
// ================================================================================================
/// A sparse Merkle tree with 63-bit keys and 4-element leaf values, without compaction.
/// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`.
/// The root of the tree is recomputed on each new leaf update.
#[derive(Clone, Debug)]
pub struct SimpleSmt {
root: Word,
depth: u32,
store: Store,
}
impl SimpleSmt {
// CONSTANTS
// --------------------------------------------------------------------------------------------
/// Minimum supported depth.
pub const MIN_DEPTH: u32 = 1;
/// Maximum supported depth.
pub const MAX_DEPTH: u32 = 63;
// CONSTRUCTORS
// --------------------------------------------------------------------------------------------
/// Creates a new simple SMT.
///
/// The provided entries will be tuples of the leaves and their corresponding keys.
///
/// # Errors
///
/// The function will fail if the provided entries count exceed the maximum tree capacity, that
/// is `2^{depth}`.
pub fn new<R, I>(entries: R, depth: u32) -> Result<Self, MerkleError>
where
R: IntoIterator<IntoIter = I>,
I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
{
let mut entries = entries.into_iter();
// validate the range of the depth.
let max = 1 << depth;
if depth < Self::MIN_DEPTH {
return Err(MerkleError::DepthTooSmall(depth));
} else if Self::MAX_DEPTH < depth {
return Err(MerkleError::DepthTooBig(depth));
} else if entries.len() > max {
return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
}
let (store, root) = Store::new(depth);
let mut tree = Self { root, depth, store };
entries.try_for_each(|(key, leaf)| tree.insert_leaf(key, leaf))?;
Ok(tree)
}
/// Returns the root of this Merkle tree.
pub const fn root(&self) -> Word {
self.root
}
/// Returns the depth of this Merkle tree.
pub const fn depth(&self) -> u32 {
self.depth
}
/// Returns the set count of the keys of the leaves.
pub fn leaves_count(&self) -> usize {
self.store.leaves_count()
}
/// Returns a node at the specified key
///
/// # Errors
/// Returns an error if:
/// * The specified depth is greater than the depth of the tree.
/// * The specified key does not exist
pub fn get_node(&self, depth: u32, key: u64) -> Result<Word, MerkleError> {
if depth == 0 {
Err(MerkleError::DepthTooSmall(depth))
} else if depth > self.depth() {
Err(MerkleError::DepthTooBig(depth))
} else if depth == self.depth() {
self.store.get_leaf_node(key)
} else {
let branch_node = self.store.get_branch_node(key, depth)?;
Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into())
}
}
/// Returns a Merkle path from the node at the specified key to the root. The node itself is
/// not included in the path.
///
/// # Errors
/// Returns an error if:
/// * The specified key does not exist as a branch or leaf node
/// * The specified depth is greater than the depth of the tree.
pub fn get_path(&self, depth: u32, key: u64) -> Result<Vec<Word>, MerkleError> {
if depth == 0 {
return Err(MerkleError::DepthTooSmall(depth));
} else if depth > self.depth() {
return Err(MerkleError::DepthTooBig(depth));
} else if depth == self.depth() && !self.store.check_leaf_node_exists(key) {
return Err(MerkleError::InvalidIndex(self.depth(), key));
}
let mut path = Vec::with_capacity(depth as usize);
let mut curr_key = key;
for n in (0..depth).rev() {
let parent_key = curr_key >> 1;
let parent_node = self.store.get_branch_node(parent_key, n)?;
let sibling_node = if curr_key & 1 == 1 {
parent_node.left
} else {
parent_node.right
};
path.push(sibling_node.into());
curr_key >>= 1;
}
Ok(path)
}
/// Return a Merkle path from the leaf at the specified key to the root. The leaf itself is not
/// included in the path.
///
/// # Errors
/// Returns an error if:
/// * The specified key does not exist as a leaf node.
pub fn get_leaf_path(&self, key: u64) -> Result<Vec<Word>, MerkleError> {
self.get_path(self.depth(), key)
}
/// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree
///
/// # Errors
/// Returns an error if the specified key is not a valid leaf index for this tree.
pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
if !self.store.check_leaf_node_exists(key) {
return Err(MerkleError::InvalidIndex(self.depth(), key));
}
self.insert_leaf(key, value)?;
Ok(())
}
/// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree
pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> {
self.store.insert_leaf_node(key, value);
let depth = self.depth();
let mut curr_key = key;
let mut curr_node: RpoDigest = value.into();
for n in (0..depth).rev() {
let parent_key = curr_key >> 1;
let parent_node = self
.store
.get_branch_node(parent_key, n)
.unwrap_or_else(|_| self.store.get_empty_node((n + 1) as usize));
let (left, right) = if curr_key & 1 == 1 {
(parent_node.left, curr_node)
} else {
(curr_node, parent_node.right)
};
self.store.insert_branch_node(parent_key, n, left, right);
curr_key = parent_key;
curr_node = Rpo256::merge(&[left, right]);
}
self.root = curr_node.into();
Ok(())
}
}
// STORE
// ================================================================================================
/// A data store for sparse Merkle tree key-value pairs.
/// Leaves and branch nodes are stored separately in B-tree maps, indexed by key and (key, depth)
/// respectively. Hashes for blank subtrees at each layer are stored in `empty_hashes`, beginning
/// with the root hash of an empty tree, and ending with the zero value of a leaf node.
#[derive(Clone, Debug)]
struct Store {
branches: BTreeMap<(u64, u32), BranchNode>,
leaves: BTreeMap<u64, Word>,
empty_hashes: Vec<RpoDigest>,
depth: u32,
}
#[derive(Clone, Debug, Default)]
struct BranchNode {
left: RpoDigest,
right: RpoDigest,
}
impl Store {
fn new(depth: u32) -> (Self, Word) {
let branches = BTreeMap::new();
let leaves = BTreeMap::new();
// Construct empty node digests for each layer of the tree
let empty_hashes: Vec<RpoDigest> = (0..depth + 1)
.scan(Word::default().into(), |state, _| {
let value = *state;
*state = Rpo256::merge(&[value, value]);
Some(value)
})
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
let root = empty_hashes[0].into();
let store = Self {
branches,
leaves,
empty_hashes,
depth,
};
(store, root)
}
fn get_empty_node(&self, depth: usize) -> BranchNode {
let digest = self.empty_hashes[depth];
BranchNode {
left: digest,
right: digest,
}
}
fn check_leaf_node_exists(&self, key: u64) -> bool {
self.leaves.contains_key(&key)
}
fn get_leaf_node(&self, key: u64) -> Result<Word, MerkleError> {
self.leaves
.get(&key)
.cloned()
.ok_or(MerkleError::InvalidIndex(self.depth, key))
}
fn insert_leaf_node(&mut self, key: u64, node: Word) {
self.leaves.insert(key, node);
}
fn get_branch_node(&self, key: u64, depth: u32) -> Result<BranchNode, MerkleError> {
self.branches
.get(&(key, depth))
.cloned()
.ok_or(MerkleError::InvalidIndex(depth, key))
}
fn insert_branch_node(&mut self, key: u64, depth: u32, left: RpoDigest, right: RpoDigest) {
let node = BranchNode { left, right };
self.branches.insert((key, depth), node);
}
fn leaves_count(&self) -> usize {
self.leaves.len()
}
}

+ 263
- 0
src/merkle/simple_smt/tests.rs

@ -0,0 +1,263 @@
use super::{
super::{MerkleTree, RpoDigest, SimpleSmt},
Rpo256, Vec, Word,
};
use crate::{Felt, FieldElement};
use core::iter;
use proptest::prelude::*;
use rand_utils::prng_array;
const KEYS4: [u64; 4] = [0, 1, 2, 3];
const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
const VALUES4: [Word; 4] = [
int_to_node(1),
int_to_node(2),
int_to_node(3),
int_to_node(4),
];
const VALUES8: [Word; 8] = [
int_to_node(1),
int_to_node(2),
int_to_node(3),
int_to_node(4),
int_to_node(5),
int_to_node(6),
int_to_node(7),
int_to_node(8),
];
const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8];
#[test]
fn build_empty_tree() {
let smt = SimpleSmt::new(iter::empty(), 3).unwrap();
let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap();
assert_eq!(mt.root(), smt.root());
}
#[test]
fn empty_digests_are_consistent() {
let depth = 5;
let root = SimpleSmt::new(iter::empty(), depth).unwrap().root();
let computed: [RpoDigest; 2] = (0..depth).fold([Default::default(); 2], |state, _| {
let digest = Rpo256::merge(&state);
[digest; 2]
});
assert_eq!(Word::from(computed[0]), root);
}
#[test]
fn build_sparse_tree() {
let mut smt = SimpleSmt::new(iter::empty(), 3).unwrap();
let mut values = ZERO_VALUES8.to_vec();
// insert single value
let key = 6;
let new_node = int_to_node(7);
values[key as usize] = new_node;
smt.insert_leaf(key, new_node)
.expect("Failed to insert leaf");
let mt2 = MerkleTree::new(values.clone()).unwrap();
assert_eq!(mt2.root(), smt.root());
assert_eq!(mt2.get_path(3, 6).unwrap(), smt.get_path(3, 6).unwrap());
// insert second value at distinct leaf branch
let key = 2;
let new_node = int_to_node(3);
values[key as usize] = new_node;
smt.insert_leaf(key, new_node)
.expect("Failed to insert leaf");
let mt3 = MerkleTree::new(values).unwrap();
assert_eq!(mt3.root(), smt.root());
assert_eq!(mt3.get_path(3, 2).unwrap(), smt.get_path(3, 2).unwrap());
}
#[test]
fn build_full_tree() {
let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
let (root, node2, node3) = compute_internal_nodes();
assert_eq!(root, tree.root());
assert_eq!(node2, tree.get_node(1, 0).unwrap());
assert_eq!(node3, tree.get_node(1, 1).unwrap());
}
#[test]
fn get_values() {
let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
// check depth 2
assert_eq!(VALUES4[0], tree.get_node(2, 0).unwrap());
assert_eq!(VALUES4[1], tree.get_node(2, 1).unwrap());
assert_eq!(VALUES4[2], tree.get_node(2, 2).unwrap());
assert_eq!(VALUES4[3], tree.get_node(2, 3).unwrap());
}
#[test]
fn get_path() {
let tree = SimpleSmt::new(KEYS4.into_iter().zip(VALUES4.into_iter()), 2).unwrap();
let (_, node2, node3) = compute_internal_nodes();
// check depth 2
assert_eq!(vec![VALUES4[1], node3], tree.get_path(2, 0).unwrap());
assert_eq!(vec![VALUES4[0], node3], tree.get_path(2, 1).unwrap());
assert_eq!(vec![VALUES4[3], node2], tree.get_path(2, 2).unwrap());
assert_eq!(vec![VALUES4[2], node2], tree.get_path(2, 3).unwrap());
// check depth 1
assert_eq!(vec![node3], tree.get_path(1, 0).unwrap());
assert_eq!(vec![node2], tree.get_path(1, 1).unwrap());
}
#[test]
fn update_leaf() {
let mut tree = SimpleSmt::new(KEYS8.into_iter().zip(VALUES8.into_iter()), 3).unwrap();
// update one value
let key = 3;
let new_node = int_to_node(9);
let mut expected_values = VALUES8.to_vec();
expected_values[key] = new_node;
let expected_tree = SimpleSmt::new(
KEYS8.into_iter().zip(expected_values.clone().into_iter()),
3,
)
.unwrap();
tree.update_leaf(key as u64, new_node).unwrap();
assert_eq!(expected_tree.root, tree.root);
// update another value
let key = 6;
let new_node = int_to_node(10);
expected_values[key] = new_node;
let expected_tree =
SimpleSmt::new(KEYS8.into_iter().zip(expected_values.into_iter()), 3).unwrap();
tree.update_leaf(key as u64, new_node).unwrap();
assert_eq!(expected_tree.root, tree.root);
}
#[test]
fn small_tree_opening_is_consistent() {
// ____k____
// / \
// _i_ _j_
// / \ / \
// e f g h
// / \ / \ / \ / \
// a b 0 0 c 0 0 d
let z = Word::from(RpoDigest::default());
let a = Word::from(Rpo256::merge(&[z.into(); 2]));
let b = Word::from(Rpo256::merge(&[a.into(); 2]));
let c = Word::from(Rpo256::merge(&[b.into(); 2]));
let d = Word::from(Rpo256::merge(&[c.into(); 2]));
let e = Word::from(Rpo256::merge(&[a.into(), b.into()]));
let f = Word::from(Rpo256::merge(&[z.into(), z.into()]));
let g = Word::from(Rpo256::merge(&[c.into(), z.into()]));
let h = Word::from(Rpo256::merge(&[z.into(), d.into()]));
let i = Word::from(Rpo256::merge(&[e.into(), f.into()]));
let j = Word::from(Rpo256::merge(&[g.into(), h.into()]));
let k = Word::from(Rpo256::merge(&[i.into(), j.into()]));
let depth = 3;
let entries = vec![(0, a), (1, b), (4, c), (7, d)];
let tree = SimpleSmt::new(entries, depth).unwrap();
assert_eq!(tree.root(), Word::from(k));
let cases: Vec<(u32, u64, Vec<Word>)> = vec![
(3, 0, vec![b, f, j]),
(3, 1, vec![a, f, j]),
(3, 4, vec![z, h, i]),
(3, 7, vec![z, g, i]),
(2, 0, vec![f, j]),
(2, 1, vec![e, j]),
(2, 2, vec![h, i]),
(2, 3, vec![g, i]),
(1, 0, vec![j]),
(1, 1, vec![i]),
];
for (depth, key, path) in cases {
let opening = tree.get_path(depth, key).unwrap();
assert_eq!(path, opening);
}
}
proptest! {
#[test]
fn arbitrary_openings_single_leaf(
depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
key in prop::num::u64::ANY,
leaf in prop::num::u64::ANY,
) {
let mut tree = SimpleSmt::new(iter::empty(), depth).unwrap();
let key = key % (1 << depth as u64);
let leaf = int_to_node(leaf);
tree.insert_leaf(key, leaf.into()).unwrap();
tree.get_leaf_path(key).unwrap();
// traverse to root, fetching all paths
for d in 1..depth {
let k = key >> (depth - d);
tree.get_path(d, k).unwrap();
}
}
#[test]
fn arbitrary_openings_multiple_leaves(
depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
count in 2u8..10u8,
ref seed in any::<[u8; 32]>()
) {
let mut tree = SimpleSmt::new(iter::empty(), depth).unwrap();
let mut seed = *seed;
let leaves = (1 << depth) - 1;
for _ in 0..count {
seed = prng_array(seed);
let mut key = [0u8; 8];
let mut leaf = [0u8; 8];
key.copy_from_slice(&seed[..8]);
leaf.copy_from_slice(&seed[8..16]);
let key = u64::from_le_bytes(key);
let key = key % leaves;
let leaf = u64::from_le_bytes(leaf);
let leaf = int_to_node(leaf);
tree.insert_leaf(key, leaf).unwrap();
tree.get_leaf_path(key).unwrap();
}
}
}
// HELPER FUNCTIONS
// --------------------------------------------------------------------------------------------
fn compute_internal_nodes() -> (Word, Word, Word) {
let node2 = Rpo256::hash_elements(&[VALUES4[0], VALUES4[1]].concat());
let node3 = Rpo256::hash_elements(&[VALUES4[2], VALUES4[3]].concat());
let root = Rpo256::merge(&[node2, node3]);
(root.into(), node2.into(), node3.into())
}
const fn int_to_node(value: u64) -> Word {
[Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
}

Loading…
Cancel
Save