Skip to content

Commit

Permalink
refactor: Use a regular Vec inside BoundedVec, add more tests
Browse files Browse the repository at this point in the history
* There is no need to use manual allocations in `BoundedVec`. Instead,
  we can just embed a regular `Vec`, make it private and make sure that
  the methods we expose never reallocate it.
* Add more test cases covering all error variants.
  • Loading branch information
vadorovsky committed May 7, 2024
1 parent 459318e commit 17cf033
Show file tree
Hide file tree
Showing 9 changed files with 224 additions and 294 deletions.
326 changes: 163 additions & 163 deletions merkle-tree/bounded-vec/src/lib.rs

Large diffs are not rendered by default.

4 changes: 1 addition & 3 deletions merkle-tree/concurrent/src/changelog.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use light_bounded_vec::{BoundedVec, Pod};
use light_bounded_vec::BoundedVec;

use crate::errors::ConcurrentMerkleTreeError;

Expand All @@ -18,8 +18,6 @@ pub type ChangelogEntry26 = ChangelogEntry<26>;
pub type ChangelogEntry32 = ChangelogEntry<32>;
pub type ChangelogEntry40 = ChangelogEntry<40>;

unsafe impl<const HEIGHT: usize> Pod for ChangelogEntry<HEIGHT> {}

impl<const HEIGHT: usize> ChangelogEntry<HEIGHT> {
pub fn new(root: [u8; 32], path: [[u8; 32]; HEIGHT], index: usize) -> Self {
let index = index as u64;
Expand Down
20 changes: 10 additions & 10 deletions merkle-tree/concurrent/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ use crate::{
// const generic here is that removing it would require keeping a `BoundecVec`
// inside `CyclicBoundedVec`. Casting byte slices to such nested vector is not
// a trivial task, but we might eventually do it at some point.
pub struct ConcurrentMerkleTree<'a, H, const HEIGHT: usize>
pub struct ConcurrentMerkleTree<H, const HEIGHT: usize>
where
H: Hasher,
{
Expand All @@ -53,23 +53,23 @@ where
pub rightmost_leaf: [u8; 32],

/// Hashes of subtrees.
pub filled_subtrees: BoundedVec<'a, [u8; 32]>,
pub filled_subtrees: BoundedVec<[u8; 32]>,
/// History of Merkle proofs.
pub changelog: CyclicBoundedVec<'a, ChangelogEntry<HEIGHT>>,
pub changelog: CyclicBoundedVec<ChangelogEntry<HEIGHT>>,
/// History of roots.
pub roots: CyclicBoundedVec<'a, [u8; 32]>,
pub roots: CyclicBoundedVec<[u8; 32]>,
/// Cached upper nodes.
pub canopy: BoundedVec<'a, [u8; 32]>,
pub canopy: BoundedVec<[u8; 32]>,

pub _hasher: PhantomData<H>,
}

pub type ConcurrentMerkleTree22<'a, H> = ConcurrentMerkleTree<'a, H, 22>;
pub type ConcurrentMerkleTree26<'a, H> = ConcurrentMerkleTree<'a, H, 26>;
pub type ConcurrentMerkleTree32<'a, H> = ConcurrentMerkleTree<'a, H, 32>;
pub type ConcurrentMerkleTree40<'a, H> = ConcurrentMerkleTree<'a, H, 40>;
pub type ConcurrentMerkleTree22<H> = ConcurrentMerkleTree<H, 22>;
pub type ConcurrentMerkleTree26<H> = ConcurrentMerkleTree<H, 26>;
pub type ConcurrentMerkleTree32<H> = ConcurrentMerkleTree<H, 32>;
pub type ConcurrentMerkleTree40<H> = ConcurrentMerkleTree<H, 40>;

impl<'a, H, const HEIGHT: usize> ConcurrentMerkleTree<'a, H, HEIGHT>
impl<'a, H, const HEIGHT: usize> ConcurrentMerkleTree<H, HEIGHT>
where
H: Hasher,
{
Expand Down
42 changes: 12 additions & 30 deletions merkle-tree/indexed/src/copy.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::{marker::PhantomData, mem, slice};

use light_bounded_vec::{BoundedVec, CyclicBoundedVec, Pod};
use light_bounded_vec::{BoundedVec, CyclicBoundedVec};
use light_concurrent_merkle_tree::{
changelog::ChangelogEntry, errors::ConcurrentMerkleTreeError, ConcurrentMerkleTree,
};
Expand All @@ -10,39 +10,21 @@ use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};
use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree, RawIndexedElement};

#[derive(Debug)]
pub struct IndexedMerkleTreeCopy<'a, H, I, const HEIGHT: usize>(
pub IndexedMerkleTree<'a, H, I, HEIGHT>,
)
pub struct IndexedMerkleTreeCopy<H, I, const HEIGHT: usize>(pub IndexedMerkleTree<H, I, HEIGHT>)
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>;

pub type IndexedMerkleTreeCopy22<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 22>;
pub type IndexedMerkleTreeCopy26<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 26>;
pub type IndexedMerkleTreeCopy32<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 32>;
pub type IndexedMerkleTreeCopy40<'a, H, I> = IndexedMerkleTreeCopy<'a, H, I, 40>;
pub type IndexedMerkleTreeCopy22<H, I> = IndexedMerkleTreeCopy<H, I, 22>;
pub type IndexedMerkleTreeCopy26<H, I> = IndexedMerkleTreeCopy<H, I, 26>;
pub type IndexedMerkleTreeCopy32<H, I> = IndexedMerkleTreeCopy<H, I, 32>;
pub type IndexedMerkleTreeCopy40<H, I> = IndexedMerkleTreeCopy<H, I, 40>;

impl<'a, H, I, const HEIGHT: usize> IndexedMerkleTreeCopy<'a, H, I, HEIGHT>
impl<H, I, const HEIGHT: usize> IndexedMerkleTreeCopy<H, I, HEIGHT>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
/// Casts a byte slice into wrapped `IndexedMerkleTree` structure reference,
Expand All @@ -66,9 +48,9 @@ where
bytes_changelog: &[u8],
bytes_roots: &[u8],
bytes_canopy: &[u8],
bytes_indexed_changelog: &'a [u8],
bytes_indexed_changelog: &[u8],
) -> Result<Self, IndexedMerkleTreeError> {
let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<'a, H, I, HEIGHT>>();
let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<H, I, HEIGHT>>();
if bytes_struct.len() != expected_bytes_struct_size {
return Err(IndexedMerkleTreeError::ConcurrentMerkleTree(
ConcurrentMerkleTreeError::StructBufferSize(
Expand All @@ -77,7 +59,7 @@ where
),
));
}
let struct_ref: *mut IndexedMerkleTree<'a, H, I, HEIGHT> = bytes_struct.as_ptr() as _;
let struct_ref: *mut IndexedMerkleTree<H, I, HEIGHT> = bytes_struct.as_ptr() as _;

let mut merkle_tree = unsafe {
ConcurrentMerkleTree {
Expand Down
41 changes: 12 additions & 29 deletions merkle-tree/indexed/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::marker::PhantomData;

use array::{IndexedArray, IndexedElement};
use light_bounded_vec::{BoundedVec, CyclicBoundedVec, Pod};
use light_bounded_vec::{BoundedVec, CyclicBoundedVec};
use light_concurrent_merkle_tree::{
errors::ConcurrentMerkleTreeError, light_hasher::Hasher, ConcurrentMerkleTree,
};
Expand All @@ -23,54 +23,37 @@ pub const FIELD_SIZE_SUB_ONE: &str =
#[derive(Debug, Default, Clone, Copy)]
pub struct RawIndexedElement<I>
where
I: Clone + Pod,
I: Clone,
{
pub value: [u8; 32],
pub next_index: I,
pub next_value: [u8; 32],
pub index: I,
}
unsafe impl<I> Pod for RawIndexedElement<I> where I: Pod + Clone {}

#[derive(Debug)]
#[repr(C)]
pub struct IndexedMerkleTree<'a, H, I, const HEIGHT: usize>
pub struct IndexedMerkleTree<H, I, const HEIGHT: usize>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
pub merkle_tree: ConcurrentMerkleTree<'a, H, HEIGHT>,
pub changelog: CyclicBoundedVec<'a, RawIndexedElement<I>>,
pub merkle_tree: ConcurrentMerkleTree<H, HEIGHT>,
pub changelog: CyclicBoundedVec<RawIndexedElement<I>>,

_index: PhantomData<I>,
}

pub type IndexedMerkleTree22<'a, H, I> = IndexedMerkleTree<'a, H, I, 22>;
pub type IndexedMerkleTree26<'a, H, I> = IndexedMerkleTree<'a, H, I, 26>;
pub type IndexedMerkleTree32<'a, H, I> = IndexedMerkleTree<'a, H, I, 32>;
pub type IndexedMerkleTree40<'a, H, I> = IndexedMerkleTree<'a, H, I, 40>;
pub type IndexedMerkleTree22<H, I> = IndexedMerkleTree<H, I, 22>;
pub type IndexedMerkleTree26<H, I> = IndexedMerkleTree<H, I, 26>;
pub type IndexedMerkleTree32<H, I> = IndexedMerkleTree<H, I, 32>;
pub type IndexedMerkleTree40<H, I> = IndexedMerkleTree<H, I, 40>;

impl<'a, H, I, const HEIGHT: usize> IndexedMerkleTree<'a, H, I, HEIGHT>
impl<H, I, const HEIGHT: usize> IndexedMerkleTree<H, I, HEIGHT>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
pub fn new(
Expand Down
4 changes: 2 additions & 2 deletions merkle-tree/indexed/src/reference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -172,12 +172,12 @@ where
/// We prove non-inclusion by:
/// 1. Showing that value is greater than leaf_lower_range_value and less than leaf_higher_range_value
/// 2. Showing that the leaf_hash H(leaf_lower_range_value, leaf_next_index, leaf_higher_value) is included in the root (Merkle tree)
pub struct NonInclusionProof<'a> {
pub struct NonInclusionProof {
pub root: [u8; 32],
pub value: [u8; 32],
pub leaf_lower_range_value: [u8; 32],
pub leaf_higher_range_value: [u8; 32],
pub leaf_index: usize,
pub next_index: usize,
pub merkle_proof: BoundedVec<'a, [u8; 32]>,
pub merkle_proof: BoundedVec<[u8; 32]>,
}
73 changes: 20 additions & 53 deletions merkle-tree/indexed/src/zero_copy.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::mem;

use light_bounded_vec::{BoundedVec, CyclicBoundedVec, Pod};
use light_bounded_vec::{BoundedVec, CyclicBoundedVec};
use light_concurrent_merkle_tree::{
changelog::ChangelogEntry, errors::ConcurrentMerkleTreeError, ConcurrentMerkleTree,
};
Expand All @@ -13,18 +13,10 @@ use crate::{errors::IndexedMerkleTreeError, IndexedMerkleTree, RawIndexedElement
pub struct IndexedMerkleTreeZeroCopy<'a, H, I, const HEIGHT: usize>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
pub merkle_tree: &'a IndexedMerkleTree<'a, H, I, HEIGHT>,
pub merkle_tree: &'a IndexedMerkleTree<H, I, HEIGHT>,
}

pub type IndexedMerkleTreeZeroCopy22<'a, H, I> = IndexedMerkleTreeZeroCopy<'a, H, I, 22>;
Expand All @@ -35,15 +27,7 @@ pub type IndexedMerkleTreeZeroCopy40<'a, H, I> = IndexedMerkleTreeZeroCopy<'a, H
impl<'a, H, I, const HEIGHT: usize> IndexedMerkleTreeZeroCopy<'a, H, I, HEIGHT>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
/// Casts a byte slice into wrapped `IndexedMerkleTree` structure reference,
Expand All @@ -64,7 +48,7 @@ where
pub unsafe fn struct_from_bytes_zero_copy(
bytes_struct: &'a [u8],
) -> Result<Self, IndexedMerkleTreeError> {
let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<'a, H, I, HEIGHT>>();
let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<H, I, HEIGHT>>();
if bytes_struct.len() != expected_bytes_struct_size {
return Err(IndexedMerkleTreeError::ConcurrentMerkleTree(
ConcurrentMerkleTreeError::StructBufferSize(
Expand All @@ -73,7 +57,7 @@ where
),
));
}
let tree: *const IndexedMerkleTree<'a, H, I, HEIGHT> = bytes_struct.as_ptr() as _;
let tree: *const IndexedMerkleTree<H, I, HEIGHT> = bytes_struct.as_ptr() as _;

Ok(Self {
merkle_tree: &*tree,
Expand Down Expand Up @@ -162,15 +146,14 @@ where
),
));
}
tree.merkle_tree.merkle_tree.roots =
ConcurrentMerkleTree::<'a, H, HEIGHT>::roots_from_bytes(
bytes_roots,
tree.merkle_tree.merkle_tree.current_root_index + 1,
tree.merkle_tree.merkle_tree.roots_length,
tree.merkle_tree.merkle_tree.roots_capacity,
)?;

let canopy_size = ConcurrentMerkleTree::<'a, H, HEIGHT>::canopy_size(
tree.merkle_tree.merkle_tree.roots = ConcurrentMerkleTree::<H, HEIGHT>::roots_from_bytes(
bytes_roots,
tree.merkle_tree.merkle_tree.current_root_index + 1,
tree.merkle_tree.merkle_tree.roots_length,
tree.merkle_tree.merkle_tree.roots_capacity,
)?;

let canopy_size = ConcurrentMerkleTree::<H, HEIGHT>::canopy_size(
tree.merkle_tree.merkle_tree.canopy_depth,
);
let expected_canopy_size = mem::size_of::<[u8; 32]>() * canopy_size;
Expand Down Expand Up @@ -210,18 +193,10 @@ where
pub struct IndexedMerkleTreeZeroCopyMut<'a, H, I, const HEIGHT: usize>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
pub merkle_tree: &'a mut IndexedMerkleTree<'a, H, I, HEIGHT>,
pub merkle_tree: &'a mut IndexedMerkleTree<H, I, HEIGHT>,
}

pub type IndexedMerkleTreeZeroCopyMut22<'a, H, I> = IndexedMerkleTreeZeroCopyMut<'a, H, I, 22>;
Expand All @@ -232,15 +207,7 @@ pub type IndexedMerkleTreeZeroCopyMut40<'a, H, I> = IndexedMerkleTreeZeroCopyMut
impl<'a, H, I, const HEIGHT: usize> IndexedMerkleTreeZeroCopyMut<'a, H, I, HEIGHT>
where
H: Hasher,
I: CheckedAdd
+ CheckedSub
+ Copy
+ Clone
+ PartialOrd
+ ToBytes
+ TryFrom<usize>
+ Unsigned
+ Pod,
I: CheckedAdd + CheckedSub + Copy + Clone + PartialOrd + ToBytes + TryFrom<usize> + Unsigned,
usize: From<I>,
{
/// Casts a byte slice into wrapped `IndexedMerkleTree` structure mutable
Expand All @@ -261,7 +228,7 @@ where
pub unsafe fn struct_from_bytes_zero_copy_mut(
bytes_struct: &'a [u8],
) -> Result<Self, IndexedMerkleTreeError> {
let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<'a, H, I, HEIGHT>>();
let expected_bytes_struct_size = mem::size_of::<IndexedMerkleTree<H, I, HEIGHT>>();
if bytes_struct.len() != expected_bytes_struct_size {
return Err(IndexedMerkleTreeError::ConcurrentMerkleTree(
ConcurrentMerkleTreeError::StructBufferSize(
Expand All @@ -270,7 +237,7 @@ where
),
));
}
let tree: *mut IndexedMerkleTree<'a, H, I, HEIGHT> = bytes_struct.as_ptr() as _;
let tree: *mut IndexedMerkleTree<H, I, HEIGHT> = bytes_struct.as_ptr() as _;

Ok(Self {
merkle_tree: &mut *tree,
Expand Down Expand Up @@ -427,7 +394,7 @@ where
tree.merkle_tree.merkle_tree.changelog_length,
tree.merkle_tree.merkle_tree.current_root_index + 1,
tree.merkle_tree.merkle_tree.roots_length,
ConcurrentMerkleTree::<'a, H, HEIGHT>::canopy_size(
ConcurrentMerkleTree::<H, HEIGHT>::canopy_size(
tree.merkle_tree.merkle_tree.canopy_depth,
),
bytes_indexed_changelog,
Expand Down
2 changes: 1 addition & 1 deletion programs/account-compression/src/state/address.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ pub struct AddressMerkleTreeAccount {
pub owner: Pubkey,
/// Delegate of the Merkle tree. This will be used for program owned Merkle trees.
pub delegate: Pubkey,
pub merkle_tree_struct: [u8; 296],
pub merkle_tree_struct: [u8; 256],
pub merkle_tree_filled_subtrees: [u8; 832],
pub merkle_tree_changelog: [u8; 1220800],
pub merkle_tree_roots: [u8; 76800],
Expand Down
Loading

0 comments on commit 17cf033

Please sign in to comment.