Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix(restore-root): fix logic to restore root #74

Merged
merged 6 commits into from
Apr 12, 2024
Merged

fix(restore-root): fix logic to restore root #74

Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d5812a8
recompute root, only the dense prefix is stored in storage
0xKitsune Apr 12, 2024
abbaf0a
update logic to calculate root when restoring
0xKitsune Apr 12, 2024
eff6e2f
fmt
0xKitsune Apr 12, 2024
85a0d6f
update test
0xKitsune Apr 12, 2024
e7c24e1
use recompute root
0xKitsune Apr 12, 2024
2a1aa64
fmt
0xKitsune Apr 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
195 changes: 117 additions & 78 deletions src/cascading_merkle_tree.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use std::ops::{Deref, DerefMut};

Check warning on line 1 in src/cascading_merkle_tree.rs

View workflow job for this annotation

GitHub Actions / Test 

unused imports: `DerefMut`, `Deref`

use color_eyre::eyre::{bail, ensure, Result};
use itertools::Itertools;

Check warning on line 4 in src/cascading_merkle_tree.rs

View workflow job for this annotation

GitHub Actions / Test 

unused import: `itertools::Itertools`
use rayon::prelude::*;

Check warning on line 5 in src/cascading_merkle_tree.rs

View workflow job for this annotation

GitHub Actions / Test 

unused import: `rayon::prelude`

use crate::merkle_tree::{Branch, Hasher, Proof};

Expand Down Expand Up @@ -37,12 +37,12 @@
where
H: Hasher,
{
depth: usize,
root: H::Hash,
empty_value: H::Hash,
depth: usize,
root: H::Hash,
empty_value: H::Hash,
sparse_column: Vec<H::Hash>,
storage: S,
_marker: std::marker::PhantomData<H>,
storage: S,
_marker: std::marker::PhantomData<H>,
}

impl<H, S> CascadingMerkleTree<H, S>
Expand Down Expand Up @@ -85,13 +85,15 @@

let mut tree = CascadingMerkleTree {
depth,
root: storage.storage_root(),
root: *empty_value,
empty_value: *empty_value,
sparse_column,
storage,
_marker: std::marker::PhantomData,
};

tree.recompute_root();

Ok(tree)
}

Expand Down Expand Up @@ -382,10 +384,15 @@
#[cfg(test)]
mod tests {

use rand::Rng;
use serial_test::serial;

use super::*;
use crate::generic_storage::{GenericStorage, MmapVec};
use crate::{
generic_storage::{GenericStorage, MmapVec},
poseidon_tree::PoseidonHash,
Field,
};

#[derive(Debug, Clone, PartialEq, Eq)]
struct TestHasher;
Expand Down Expand Up @@ -585,12 +592,12 @@
let leaves = vec![1; num_leaves];
let tree = CascadingMerkleTree::<TestHasher>::new_with_leaves(vec![], 10, &0, &leaves);
let expected = CascadingMerkleTree::<TestHasher> {
depth: 10,
root: 5,
empty_value: 0,
depth: 10,
root: 5,
empty_value: 0,
sparse_column: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
storage: vec![5, 1, 2, 1, 4, 2, 1, 1, 5, 1, 1, 0, 1, 0, 0, 0],
_marker: std::marker::PhantomData,
storage: vec![5, 1, 2, 1, 4, 2, 1, 1, 5, 1, 1, 0, 1, 0, 0, 0],
_marker: std::marker::PhantomData,
};
debug_tree(&tree);
tree.validate().unwrap();
Expand All @@ -604,12 +611,12 @@
let empty = 0;
let tree = CascadingMerkleTree::<TestHasher>::new_with_leaves(vec![], 10, &empty, &leaves);
let expected = CascadingMerkleTree::<TestHasher> {
depth: 10,
root: 8,
empty_value: 0,
depth: 10,
root: 8,
empty_value: 0,
sparse_column: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
storage: vec![8, 1, 2, 1, 4, 2, 1, 1, 8, 4, 2, 2, 1, 1, 1, 1],
_marker: std::marker::PhantomData,
storage: vec![8, 1, 2, 1, 4, 2, 1, 1, 8, 4, 2, 2, 1, 1, 1, 1],
_marker: std::marker::PhantomData,
};
debug_tree(&tree);
tree.validate().unwrap();
Expand All @@ -622,12 +629,12 @@
let empty = 0;
let tree = CascadingMerkleTree::<TestHasher>::new_with_leaves(vec![], 10, &empty, &leaves);
let expected = CascadingMerkleTree::<TestHasher> {
depth: 10,
root: 0,
empty_value: 0,
depth: 10,
root: 0,
empty_value: 0,
sparse_column: vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
storage: vec![0, 0],
_marker: std::marker::PhantomData,
storage: vec![0, 0],
_marker: std::marker::PhantomData,
};
debug_tree(&tree);
tree.validate().unwrap();
Expand All @@ -640,12 +647,12 @@
let empty = 1;
let tree = CascadingMerkleTree::<TestHasher>::new_with_leaves(vec![], 10, &empty, &leaves);
let expected = CascadingMerkleTree::<TestHasher> {
depth: 10,
root: 1024,
empty_value: 1,
depth: 10,
root: 1024,
empty_value: 1,
sparse_column: vec![1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024],
storage: vec![0, 1],
_marker: std::marker::PhantomData,
storage: vec![0, 1],
_marker: std::marker::PhantomData,
};
debug_tree(&tree);
tree.validate().unwrap();
Expand All @@ -659,12 +666,12 @@
let empty = 1;
let tree = CascadingMerkleTree::<TestHasher>::new_with_leaves(vec![], 4, &empty, &leaves);
let expected = CascadingMerkleTree::<TestHasher> {
depth: 4,
root: 8,
empty_value: 1,
depth: 4,
root: 8,
empty_value: 1,
sparse_column: vec![1, 2, 4, 8, 16],
storage: vec![8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
_marker: std::marker::PhantomData,
storage: vec![8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
_marker: std::marker::PhantomData,
};
debug_tree(&tree);
tree.validate().unwrap();
Expand Down Expand Up @@ -727,9 +734,12 @@
tree.validate().unwrap();
debug_tree(&tree);
let expected = vec![
(0, vec![
1usize, 3, 6, 7, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31,
]),
(
0,
vec![
1usize, 3, 6, 7, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31,
],
),
(1, vec![2, 5, 10, 11, 20, 21, 22, 23]),
(2, vec![4, 9, 18, 19]),
(3, vec![8, 17]),
Expand Down Expand Up @@ -777,42 +787,60 @@
debug_tree(&tree);
tree.validate().unwrap();
let expected = vec![
(1, vec![
Branch::Left(2),
Branch::Left(7),
Branch::Left(13),
Branch::Left(8),
]),
(2, vec![
Branch::Right(1),
Branch::Left(7),
Branch::Left(13),
Branch::Left(8),
]),
(3, vec![
Branch::Left(4),
Branch::Right(3),
Branch::Left(13),
Branch::Left(8),
]),
(4, vec![
Branch::Right(3),
Branch::Right(3),
Branch::Left(13),
Branch::Left(8),
]),
(5, vec![
Branch::Left(6),
Branch::Left(2),
Branch::Right(10),
Branch::Left(8),
]),
(6, vec![
Branch::Right(5),
Branch::Left(2),
Branch::Right(10),
Branch::Left(8),
]),
(
1,
vec![
Branch::Left(2),
Branch::Left(7),
Branch::Left(13),
Branch::Left(8),
],
),
(
2,
vec![
Branch::Right(1),
Branch::Left(7),
Branch::Left(13),
Branch::Left(8),
],
),
(
3,
vec![
Branch::Left(4),
Branch::Right(3),
Branch::Left(13),
Branch::Left(8),
],
),
(
4,
vec![
Branch::Right(3),
Branch::Right(3),
Branch::Left(13),
Branch::Left(8),
],
),
(
5,
vec![
Branch::Left(6),
Branch::Left(2),
Branch::Right(10),
Branch::Left(8),
],
),
(
6,
vec![
Branch::Right(5),
Branch::Left(2),
Branch::Right(10),
Branch::Left(8),
],
),
];
for (leaf, expected_proof) in expected {
let proof = tree.proof_from_hash(leaf).unwrap();
Expand Down Expand Up @@ -881,31 +909,42 @@
#[test]
#[serial]
fn test_restore_from_cache() -> color_eyre::Result<()> {
let empty = 0;
let leaves = vec![1; 1 << 20];
let mut rng = rand::thread_rng();

let leaves: Vec<Field> = (0..1 << 2)
.map(|_| {
let val = rng.gen::<usize>();

Field::from(val)
})
.collect::<Vec<Field>>();

// Create a new tmp file for mmap storage
let tempfile = tempfile::NamedTempFile::new()?;
let file_path = tempfile.path().to_owned();

// Initialize the expected tree
let mmap_vec: MmapVec<_> = unsafe { MmapVec::new(tempfile.reopen()?).unwrap() };
let expected_tree = CascadingMerkleTree::<TestHasher, MmapVec<_>>::new_with_leaves(
mmap_vec, 30, &empty, &leaves,
let expected_tree = CascadingMerkleTree::<PoseidonHash, MmapVec<_>>::new_with_leaves(
mmap_vec,
3,
&Field::ZERO,
&leaves,
);

let expected_root = expected_tree.root();
let expected_leaves = expected_tree.leaves().collect::<Vec<usize>>();
let expected_leaves = expected_tree.leaves().collect::<Vec<Field>>();

drop(expected_tree);

// Restore the tree
let mmap_vec: MmapVec<_> = unsafe { MmapVec::restore(file_path).unwrap() };
let tree = CascadingMerkleTree::<TestHasher, MmapVec<_>>::restore(mmap_vec, 30, &empty)?;
let tree =
CascadingMerkleTree::<PoseidonHash, MmapVec<_>>::restore(mmap_vec, 3, &Field::ZERO)?;

// Assert that the root and the leaves are as expected
assert_eq!(tree.root(), expected_root);
assert_eq!(tree.leaves().collect::<Vec<usize>>(), expected_leaves);
assert_eq!(tree.leaves().collect::<Vec<Field>>(), expected_leaves);

Ok(())
}
Expand Down
Loading