refactor: refactor to clean up and simplify things

This commit is contained in:
tohrnii
2023-06-13 10:28:15 +01:00
parent fe9aa8c28c
commit 0e0a3fda4f
6 changed files with 197 additions and 228 deletions

View File

@@ -4,7 +4,7 @@ use super::{
};
use crate::{
hash::rpo::Rpo256,
merkle::{int_to_leaf, int_to_node, MerklePathSet, MerkleTree, SimpleSmt},
merkle::{digests_to_words, int_to_leaf, int_to_node, MerklePathSet, MerkleTree, SimpleSmt},
Felt, Word, WORD_SIZE,
};
@@ -15,17 +15,17 @@ use std::error::Error;
// ================================================================================================
const KEYS4: [u64; 4] = [0, 1, 2, 3];
const VALUES4: [Word; 4] = [int_to_leaf(1), int_to_leaf(2), int_to_leaf(3), int_to_leaf(4)];
const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];
const VALUES8: [Word; 8] = [
int_to_leaf(1),
int_to_leaf(2),
int_to_leaf(3),
int_to_leaf(4),
int_to_leaf(5),
int_to_leaf(6),
int_to_leaf(7),
int_to_leaf(8),
const VALUES8: [RpoDigest; 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),
];
// TESTS
@@ -33,16 +33,16 @@ const VALUES8: [Word; 8] = [
#[test]
fn test_root_not_in_store() -> Result<(), MerkleError> {
let mtree = MerkleTree::new(VALUES4.to_vec())?;
let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
let store = MerkleStore::from(&mtree);
assert_eq!(
store.get_node(VALUES4[0].into(), NodeIndex::make(mtree.depth(), 0)),
Err(MerkleError::RootNotInStore(VALUES4[0].into())),
store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
Err(MerkleError::RootNotInStore(VALUES4[0])),
"Leaf 0 is not a root"
);
assert_eq!(
store.get_path(VALUES4[0].into(), NodeIndex::make(mtree.depth(), 0)),
Err(MerkleError::RootNotInStore(VALUES4[0].into())),
store.get_path(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
Err(MerkleError::RootNotInStore(VALUES4[0])),
"Leaf 0 is not a root"
);
@@ -51,29 +51,29 @@ fn test_root_not_in_store() -> Result<(), MerkleError> {
#[test]
fn test_merkle_tree() -> Result<(), MerkleError> {
let mtree = MerkleTree::new(VALUES4.to_vec())?;
let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
let store = MerkleStore::from(&mtree);
// STORE LEAVES ARE CORRECT -------------------------------------------------------------------
// checks the leaves in the store corresponds to the expected values
assert_eq!(
store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
Ok(VALUES4[0].into()),
Ok(VALUES4[0]),
"node 0 must be in the tree"
);
assert_eq!(
store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
Ok(VALUES4[1].into()),
Ok(VALUES4[1]),
"node 1 must be in the tree"
);
assert_eq!(
store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
Ok(VALUES4[2].into()),
Ok(VALUES4[2]),
"node 2 must be in the tree"
);
assert_eq!(
store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
Ok(VALUES4[3].into()),
Ok(VALUES4[3]),
"node 3 must be in the tree"
);
@@ -104,7 +104,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
// assert the merkle path returned by the store is the same as the one in the tree
let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap();
assert_eq!(
VALUES4[0], *result.value,
VALUES4[0], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -115,7 +115,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap();
assert_eq!(
VALUES4[1], *result.value,
VALUES4[1], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -126,7 +126,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap();
assert_eq!(
VALUES4[2], *result.value,
VALUES4[2], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -137,7 +137,7 @@ fn test_merkle_tree() -> Result<(), MerkleError> {
let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap();
assert_eq!(
VALUES4[3], *result.value,
VALUES4[3], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -157,7 +157,7 @@ fn test_empty_roots() {
for depth in 0..255 {
root = Rpo256::merge(&[root; 2]);
assert!(
store.get_node(root.into(), NodeIndex::make(0, 0)).is_ok(),
store.get_node(root, NodeIndex::make(0, 0)).is_ok(),
"The root of the empty tree of depth {depth} must be registered"
);
}
@@ -197,7 +197,8 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> {
#[test]
fn test_get_invalid_node() {
let mtree = MerkleTree::new(VALUES4.to_vec()).expect("creating a merkle tree must work");
let mtree =
MerkleTree::new(digests_to_words(&VALUES4)).expect("creating a merkle tree must work");
let store = MerkleStore::from(&mtree);
let _ = store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3));
}
@@ -211,7 +212,7 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
let idx = NodeIndex::make(1, 0);
assert_eq!(smt.get_node(idx).unwrap(), leaves2[0].into());
assert_eq!(store.get_node(smt.root().into(), idx).unwrap(), smt.get_node(idx).unwrap());
assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap());
let idx = NodeIndex::make(1, 1);
assert_eq!(smt.get_node(idx).unwrap(), leaves2[1].into());
@@ -222,36 +223,38 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
#[test]
fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
let smt =
SimpleSmt::with_leaves(SimpleSmt::MAX_DEPTH, KEYS4.into_iter().zip(VALUES4.into_iter()))
.unwrap();
let smt = SimpleSmt::with_leaves(
SimpleSmt::MAX_DEPTH,
KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()),
)
.unwrap();
let store = MerkleStore::from(&smt);
// STORE LEAVES ARE CORRECT ==============================================================
// checks the leaves in the store corresponds to the expected values
assert_eq!(
store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 0)),
Ok(VALUES4[0].into()),
store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)),
Ok(VALUES4[0]),
"node 0 must be in the tree"
);
assert_eq!(
store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 1)),
Ok(VALUES4[1].into()),
store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)),
Ok(VALUES4[1]),
"node 1 must be in the tree"
);
assert_eq!(
store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 2)),
Ok(VALUES4[2].into()),
store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)),
Ok(VALUES4[2]),
"node 2 must be in the tree"
);
assert_eq!(
store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 3)),
Ok(VALUES4[3].into()),
store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)),
Ok(VALUES4[3]),
"node 3 must be in the tree"
);
assert_eq!(
store.get_node(smt.root().into(), NodeIndex::make(smt.depth(), 4)),
store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)),
Ok(RpoDigest::default()),
"unmodified node 4 must be ZERO"
);
@@ -288,7 +291,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
// assert the merkle path returned by the store is the same as the one in the tree
let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 0)).unwrap();
assert_eq!(
VALUES4[0], *result.value,
VALUES4[0], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -299,7 +302,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 1)).unwrap();
assert_eq!(
VALUES4[1], *result.value,
VALUES4[1], result.value,
"Value for merkle path at index 1 must match leaf value"
);
assert_eq!(
@@ -310,7 +313,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 2)).unwrap();
assert_eq!(
VALUES4[2], *result.value,
VALUES4[2], result.value,
"Value for merkle path at index 2 must match leaf value"
);
assert_eq!(
@@ -321,7 +324,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 3)).unwrap();
assert_eq!(
VALUES4[3], *result.value,
VALUES4[3], result.value,
"Value for merkle path at index 3 must match leaf value"
);
assert_eq!(
@@ -347,7 +350,7 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
#[test]
fn test_add_merkle_paths() -> Result<(), MerkleError> {
let mtree = MerkleTree::new(VALUES4.to_vec())?;
let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
let i0 = 0;
let p0 = mtree.get_path(NodeIndex::make(2, i0)).unwrap();
@@ -362,10 +365,10 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
let p3 = mtree.get_path(NodeIndex::make(2, i3)).unwrap();
let paths = [
(i0, VALUES4[i0 as usize].into(), p0),
(i1, VALUES4[i1 as usize].into(), p1),
(i2, VALUES4[i2 as usize].into(), p2),
(i3, VALUES4[i3 as usize].into(), p3),
(i0, VALUES4[i0 as usize], p0),
(i1, VALUES4[i1 as usize], p1),
(i2, VALUES4[i2 as usize], p2),
(i3, VALUES4[i3 as usize], p3),
];
let mut store = MerkleStore::default();
@@ -378,22 +381,22 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
// checks the leaves in the store corresponds to the expected values
assert_eq!(
store.get_node(set.root(), NodeIndex::make(set.depth(), 0)),
Ok(VALUES4[0].into()),
Ok(VALUES4[0]),
"node 0 must be in the set"
);
assert_eq!(
store.get_node(set.root(), NodeIndex::make(set.depth(), 1)),
Ok(VALUES4[1].into()),
Ok(VALUES4[1]),
"node 1 must be in the set"
);
assert_eq!(
store.get_node(set.root(), NodeIndex::make(set.depth(), 2)),
Ok(VALUES4[2].into()),
Ok(VALUES4[2]),
"node 2 must be in the set"
);
assert_eq!(
store.get_node(set.root(), NodeIndex::make(set.depth(), 3)),
Ok(VALUES4[3].into()),
Ok(VALUES4[3]),
"node 3 must be in the set"
);
@@ -424,7 +427,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
// assert the merkle path returned by the store is the same as the one in the set
let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 0)).unwrap();
assert_eq!(
VALUES4[0], *result.value,
VALUES4[0], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -435,7 +438,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 1)).unwrap();
assert_eq!(
VALUES4[1], *result.value,
VALUES4[1], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -446,7 +449,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 2)).unwrap();
assert_eq!(
VALUES4[2], *result.value,
VALUES4[2], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -457,7 +460,7 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> {
let result = store.get_path(set.root(), NodeIndex::make(set.depth(), 3)).unwrap();
assert_eq!(
VALUES4[3], *result.value,
VALUES4[3], result.value,
"Value for merkle path at index 0 must match leaf value"
);
assert_eq!(
@@ -517,7 +520,7 @@ fn store_path_opens_from_leaf() {
let mtree = MerkleTree::new(vec![a, b, c, d, e, f, g, h]).unwrap();
let store = MerkleStore::from(&mtree);
let path = store.get_path(root.into(), NodeIndex::make(3, 1)).unwrap().path;
let path = store.get_path(root, NodeIndex::make(3, 1)).unwrap().path;
let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec());
assert_eq!(path, expected);
@@ -525,7 +528,7 @@ fn store_path_opens_from_leaf() {
#[test]
fn test_set_node() -> Result<(), MerkleError> {
let mtree = MerkleTree::new(VALUES4.to_vec())?;
let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
let mut store = MerkleStore::from(&mtree);
let value = int_to_node(42);
let index = NodeIndex::make(mtree.depth(), 0);
@@ -537,7 +540,7 @@ fn test_set_node() -> Result<(), MerkleError> {
#[test]
fn test_constructors() -> Result<(), MerkleError> {
let mtree = MerkleTree::new(VALUES4.to_vec())?;
let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
let store = MerkleStore::from(&mtree);
let depth = mtree.depth();
@@ -549,7 +552,11 @@ fn test_constructors() -> Result<(), MerkleError> {
}
let depth = 32;
let smt = SimpleSmt::with_leaves(depth, KEYS4.into_iter().zip(VALUES4.into_iter())).unwrap();
let smt = SimpleSmt::with_leaves(
depth,
KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()),
)
.unwrap();
let store = MerkleStore::from(&smt);
let depth = smt.depth();
@@ -561,20 +568,20 @@ fn test_constructors() -> Result<(), MerkleError> {
let d = 2;
let paths = [
(0, VALUES4[0].into(), mtree.get_path(NodeIndex::make(d, 0)).unwrap()),
(1, VALUES4[1].into(), mtree.get_path(NodeIndex::make(d, 1)).unwrap()),
(2, VALUES4[2].into(), mtree.get_path(NodeIndex::make(d, 2)).unwrap()),
(3, VALUES4[3].into(), mtree.get_path(NodeIndex::make(d, 3)).unwrap()),
(0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0)).unwrap()),
(1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1)).unwrap()),
(2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2)).unwrap()),
(3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()),
];
let mut store1 = MerkleStore::default();
store1.add_merkle_paths(paths.clone())?;
let mut store2 = MerkleStore::default();
store2.add_merkle_path(0, VALUES4[0].into(), mtree.get_path(NodeIndex::make(d, 0))?)?;
store2.add_merkle_path(1, VALUES4[1].into(), mtree.get_path(NodeIndex::make(d, 1))?)?;
store2.add_merkle_path(2, VALUES4[2].into(), mtree.get_path(NodeIndex::make(d, 2))?)?;
store2.add_merkle_path(3, VALUES4[3].into(), mtree.get_path(NodeIndex::make(d, 3))?)?;
store2.add_merkle_path(0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0))?)?;
store2.add_merkle_path(1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1))?)?;
store2.add_merkle_path(2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2))?)?;
store2.add_merkle_path(3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?;
let set = MerklePathSet::new(d).with_paths(paths).unwrap();
for key in [0, 1, 2, 3] {
@@ -744,16 +751,16 @@ fn get_leaf_depth_works_with_depth_8() {
#[test]
fn mstore_subset() {
// add a Merkle tree of depth 3 to the store
let mtree = MerkleTree::new(VALUES8.to_vec()).unwrap();
let mtree = MerkleTree::new(digests_to_words(&VALUES8)).unwrap();
let mut store = MerkleStore::default();
let empty_store_num_nodes = store.nodes.len();
store.extend(mtree.inner_nodes());
// build 3 subtrees contained within the above Merkle tree; note that subtree2 is a subset
// of subtree1
let subtree1 = MerkleTree::new(VALUES8[..4].to_vec()).unwrap();
let subtree2 = MerkleTree::new(VALUES8[2..4].to_vec()).unwrap();
let subtree3 = MerkleTree::new(VALUES8[6..].to_vec()).unwrap();
let subtree1 = MerkleTree::new(digests_to_words(&VALUES8[..4])).unwrap();
let subtree2 = MerkleTree::new(digests_to_words(&VALUES8[2..4])).unwrap();
let subtree3 = MerkleTree::new(digests_to_words(&VALUES8[6..])).unwrap();
// --- extract all 3 subtrees ---------------------------------------------
@@ -798,7 +805,7 @@ fn check_mstore_subtree(store: &MerkleStore, subtree: &MerkleTree) {
#[cfg(feature = "std")]
#[test]
fn test_serialization() -> Result<(), Box<dyn Error>> {
let mtree = MerkleTree::new(VALUES4.to_vec())?;
let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
let store = MerkleStore::from(&mtree);
let decoded = MerkleStore::read_from_bytes(&store.to_bytes()).expect("deserialization failed");
assert_eq!(store, decoded);