diff --git a/contracts/programs/gummyroll/src/error.rs b/contracts/programs/gummyroll/src/error.rs index e552881d..b82622b3 100644 --- a/contracts/programs/gummyroll/src/error.rs +++ b/contracts/programs/gummyroll/src/error.rs @@ -24,6 +24,15 @@ pub enum GummyrollError { /// When using Canopy, the stored byte length should a multiple of the node's byte length (32 bytes) #[msg("Expected a different byte length for the merkle roll canopy")] CanopyLengthMismatch, + + /// Error modify pre-subtree-append data structure. + /// Your tree may have been modified during your initialization process, your proof may be invalid or out of date. Please reset the account. + #[msg("Merkle Roll Pre Append Processing Error")] + MerkleTreePreAppendProcessingError, + + /// An issue was detected with loading the provided account data for the PreAppend data structure. + #[msg("Issue zero copying concurrent merkle tree append structure")] + PreAppendZeroCopyError, } impl From<&CMTError> for GummyrollError { diff --git a/contracts/programs/gummyroll/src/lib.rs b/contracts/programs/gummyroll/src/lib.rs index 9bccd98c..b21bd304 100644 --- a/contracts/programs/gummyroll/src/lib.rs +++ b/contracts/programs/gummyroll/src/lib.rs @@ -33,7 +33,7 @@ pub mod utils; use crate::error::GummyrollError; use crate::state::{CandyWrapper, ChangeLogEvent, MerkleRollHeader}; use crate::utils::{wrap_event, ZeroCopy}; -pub use concurrent_merkle_tree::{error::CMTError, merkle_roll::MerkleRoll, state::Node}; +pub use concurrent_merkle_tree::{error::CMTError, merkle_roll::{MerkleRoll, MerkleInterface, MerkleRollPreAppend, PreAppendInterface}, state::Node}; declare_id!("GRoLLzvxpxxu2PGNJMMeZPyMxjAUH9pKqxGXV9DGiceU"); @@ -52,6 +52,47 @@ pub struct Initialize<'info> { /// See `WRAPYChf58WFCnyjXKJHtrPgzKXgHp6MD9aVDqJBbGh` pub candy_wrapper: Program<'info, CandyWrapper>, } +#[derive(Accounts)] +/// Context to modify the Pre-Append data structure including init, reset or push partition +pub struct ModifyPreAppend<'info> { + /// CHECK: validated in instruction + pub merkle_roll: UncheckedAccount<'info>, + + #[account(mut, seeds = [merkle_roll.key().as_ref()], bump)] + /// CHECK: validated in instruction, besides seeds. + pub subtree_append: UncheckedAccount<'info>, + + /// Authority over merkle_roll + /// Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs. + pub authority: Signer<'info>, + + /// Program used to emit changelogs as instruction data. + /// See `WRAPYChf58WFCnyjXKJHtrPgzKXgHp6MD9aVDqJBbGh` + pub candy_wrapper: Program<'info, CandyWrapper>, +} + +#[derive(Accounts)] +/// Context to append a subtree to this Gummyroll tree +pub struct AppendSubtree<'info> { + #[account(mut)] + /// CHECK: validated in instruction + pub merkle_roll: UncheckedAccount<'info>, + + /// CHECK: validated in instruction + pub subtree_merkle_roll: UncheckedAccount<'info>, + + /// CHECK: validated in instruction + #[account(seeds = [subtree_merkle_roll.key().as_ref()], bump)] + pub subtree_append: UncheckedAccount<'info>, + + /// Authority over merkle_roll + /// Typically a program, e.g., the Bubblegum contract validates that leaves are valid NFTs. + pub authority: Signer<'info>, + + /// Program used to emit changelogs as instruction data. + /// See `WRAPYChf58WFCnyjXKJHtrPgzKXgHp6MD9aVDqJBbGh` + pub candy_wrapper: Program<'info, CandyWrapper>, +} /// Context for inserting, appending, or replacing a leaf in the tree #[derive(Accounts)] @@ -191,30 +232,6 @@ fn fill_in_proof_from_canopy( Ok(()) } -/// This macro applies functions on a merkle roll and emits leaf information -/// needed to sync the merkle tree state with off-chain indexers. -macro_rules! merkle_roll_depth_size_apply_fn { - ($max_depth:literal, $max_size:literal, $id:ident, $bytes:ident, $func:ident, $($arg:tt)*) => { - match MerkleRoll::<$max_depth, $max_size>::load_mut_bytes($bytes) { - Ok(merkle_roll) => { - match merkle_roll.$func($($arg)*) { - Ok(_) => { - Ok(Box::::from((merkle_roll.get_change_log(), $id, merkle_roll.sequence_number))) - } - Err(err) => { - msg!("Error using concurrent merkle tree: {}", err); - err!(GummyrollError::ConcurrentMerkleTreeError) - } - } - } - Err(err) => { - msg!("Error zero copying merkle roll: {}", err); - err!(GummyrollError::ZeroCopyError) - } - } - } -} - /// This applies a given function on a merkle roll by /// allowing the compiler to infer the size of the tree based /// upon the header information stored on-chain @@ -255,42 +272,109 @@ macro_rules! merkle_roll_get_size { }; } -/// This applies a given function on a merkle roll by -/// allowing the compiler to infer the size of the tree based -/// upon the header information stored on-chain -macro_rules! merkle_roll_apply_fn { - ($header:ident, $id:ident, $bytes:ident, $func:ident, $($arg:tt)*) => { - // Note: max_buffer_size MUST be a power of 2 - match ($header.max_depth, $header.max_buffer_size) { - (3, 8) => merkle_roll_depth_size_apply_fn!(3, 8, $id, $bytes, $func, $($arg)*), - (5, 8) => merkle_roll_depth_size_apply_fn!(5, 8, $id, $bytes, $func, $($arg)*), - (14, 64) => merkle_roll_depth_size_apply_fn!(14, 64, $id, $bytes, $func, $($arg)*), - (14, 256) => merkle_roll_depth_size_apply_fn!(14, 256, $id, $bytes, $func, $($arg)*), - (14, 1024) => merkle_roll_depth_size_apply_fn!(14, 1024, $id, $bytes, $func, $($arg)*), - (14, 2048) => merkle_roll_depth_size_apply_fn!(14, 2048, $id, $bytes, $func, $($arg)*), - (20, 64) => merkle_roll_depth_size_apply_fn!(20, 64, $id, $bytes, $func, $($arg)*), - (20, 256) => merkle_roll_depth_size_apply_fn!(20, 256, $id, $bytes, $func, $($arg)*), - (20, 1024) => merkle_roll_depth_size_apply_fn!(20, 1024, $id, $bytes, $func, $($arg)*), - (20, 2048) => merkle_roll_depth_size_apply_fn!(20, 2048, $id, $bytes, $func, $($arg)*), - (24, 64) => merkle_roll_depth_size_apply_fn!(24, 64, $id, $bytes, $func, $($arg)*), - (24, 256) => merkle_roll_depth_size_apply_fn!(24, 256, $id, $bytes, $func, $($arg)*), - (24, 512) => merkle_roll_depth_size_apply_fn!(24, 512, $id, $bytes, $func, $($arg)*), - (24, 1024) => merkle_roll_depth_size_apply_fn!(24, 1024, $id, $bytes, $func, $($arg)*), - (24, 2048) => merkle_roll_depth_size_apply_fn!(24, 2048, $id, $bytes, $func, $($arg)*), - (26, 512) => merkle_roll_depth_size_apply_fn!(26, 512, $id, $bytes, $func, $($arg)*), - (26, 1024) => merkle_roll_depth_size_apply_fn!(26, 1024, $id, $bytes, $func, $($arg)*), - (26, 2048) => merkle_roll_depth_size_apply_fn!(26, 2048, $id, $bytes, $func, $($arg)*), - (30, 512) => merkle_roll_depth_size_apply_fn!(30, 512, $id, $bytes, $func, $($arg)*), - (30, 1024) => merkle_roll_depth_size_apply_fn!(30, 1024, $id, $bytes, $func, $($arg)*), - (30, 2048) => merkle_roll_depth_size_apply_fn!(30, 2048, $id, $bytes, $func, $($arg)*), +/// Returns the size of a merkle_roll's associated pre-append data structure +macro_rules! merkle_roll_append_get_size { + ($header:ident) => { + match ($header.max_depth) { + 3 => Ok(size_of::>()), + 5 => Ok(size_of::>()), + 14 => Ok(size_of::>()), + 20 => Ok(size_of::>()), + 24 => Ok(size_of::>()), + 26 => Ok(size_of::>()), + 30 => Ok(size_of::>()), _ => { - msg!("Failed to apply {} on merkle roll with max depth {} and max buffer size {}", stringify!($func), $header.max_depth, $header.max_buffer_size); + msg!( + "Failed to get size of pre append struct for max_depth {}", + $header.max_depth + ); err!(GummyrollError::MerkleRollConstantsError) } } }; } +macro_rules! merkle_roll_interface_for_size { + ($max_depth: literal, $max_buffer_size: literal, $bytes: ident) => { + match MerkleRoll::<$max_depth,$max_buffer_size>::load_mut_bytes($bytes) { + Ok(merkle_roll) => { Ok(merkle_roll as &mut dyn MerkleInterface) } + Err(err) => { + msg!("Error zero copying merkle roll: {}", err); + err!(GummyrollError::ZeroCopyError) + } + } + } +} + +fn get_merkle_roll_interface(max_depth: u32, max_buffer_size: u32, bytes: &mut [u8]) -> Result<&mut dyn MerkleInterface> { + match (max_depth, max_buffer_size) { + (3, 8) => merkle_roll_interface_for_size!(3, 8, bytes), + (5, 8) => merkle_roll_interface_for_size!(5, 8, bytes), + (14, 64) => merkle_roll_interface_for_size!(14, 64, bytes), + (14, 256) => merkle_roll_interface_for_size!(14, 256, bytes), + (14, 1024) => merkle_roll_interface_for_size!(14, 1024, bytes), + (14, 2048) => merkle_roll_interface_for_size!(14, 2048, bytes), + (20, 64) => merkle_roll_interface_for_size!(20, 64, bytes), + (20, 256) => merkle_roll_interface_for_size!(20, 256, bytes), + (20, 1024) => merkle_roll_interface_for_size!(20, 1024, bytes), + (20, 2048) => merkle_roll_interface_for_size!(20, 2048, bytes), + (24, 64) => merkle_roll_interface_for_size!(24, 64, bytes), + (24, 256) => merkle_roll_interface_for_size!(24, 256, bytes), + (24, 512) => merkle_roll_interface_for_size!(24, 512, bytes), + (24, 1024) => merkle_roll_interface_for_size!(24, 1024, bytes), + (24, 2048) => merkle_roll_interface_for_size!(24, 2048, bytes), + (26, 512) => merkle_roll_interface_for_size!(26, 512, bytes), + (26, 1024) => merkle_roll_interface_for_size!(26, 1024, bytes), + (26, 2048) => merkle_roll_interface_for_size!(26, 2048, bytes), + (30, 512) => merkle_roll_interface_for_size!(30, 512, bytes), + (30, 1024) => merkle_roll_interface_for_size!(30, 1024, bytes), + (30, 2048) => merkle_roll_interface_for_size!(30, 2048, bytes), + _ => { + msg!( + "max depth {} and max buffer size {} are an unsupported combination", + max_depth, + max_buffer_size + ); + err!(GummyrollError::MerkleRollConstantsError) + } + } +} + +macro_rules! pre_append_interface_for_size { + ($num_partitions: literal, $bytes: ident) => { + match MerkleRollPreAppend::<$num_partitions>::load_mut_bytes($bytes) { + Ok(merkle_roll_pre_append) => { Ok(merkle_roll_pre_append as &mut dyn PreAppendInterface) } + Err(err) => { + msg!("Error zero copying pre append data structure: {}", err); + err!(GummyrollError::PreAppendZeroCopyError) + } + } + } +} + +fn get_pre_append_interface(max_depth: u32, bytes: &mut [u8]) -> Result<&mut dyn PreAppendInterface> { + match max_depth { + 3 => pre_append_interface_for_size!(4, bytes), + 5 => pre_append_interface_for_size!(6, bytes), + 14 => pre_append_interface_for_size!(15, bytes), + 20 => pre_append_interface_for_size!(21, bytes), + 24 => pre_append_interface_for_size!(24, bytes), + 26 => pre_append_interface_for_size!(26, bytes), + 30 => pre_append_interface_for_size!(31, bytes), + _ => { + msg!( + "Failed to get size of append for max_depth {}", + max_depth + ); + err!(GummyrollError::MerkleRollConstantsError) + } + } +} + +fn get_changelog_after_op(merkle_roll_obj: &dyn MerkleInterface, id: Pubkey) -> Box:: { + Box::::from((merkle_roll_obj.get_change_log(), id, merkle_roll_obj.get_sequence_number())) +} + #[program] pub mod gummyroll { use super::*; @@ -326,10 +410,13 @@ pub mod gummyroll { let merkle_roll_size = merkle_roll_get_size!(header)?; let (roll_bytes, canopy_bytes) = rest.split_at_mut(merkle_roll_size); let id = ctx.accounts.merkle_roll.key(); - let change_log = merkle_roll_apply_fn!(header, id, roll_bytes, initialize,)?; + let mut merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + merkle_roll_obj.initialize(); + let change_log = get_changelog_after_op(merkle_roll_obj, id); wrap_event(change_log.try_to_vec()?, &ctx.accounts.candy_wrapper)?; emit!(*change_log); - update_canopy(canopy_bytes, header.max_depth, None) + update_canopy(canopy_bytes, header.max_depth, None); + Ok(()) } /// Note: @@ -374,20 +461,107 @@ pub mod gummyroll { let id = ctx.accounts.merkle_roll.key(); // A call is made to MerkleRoll::initialize_with_root(root, leaf, proof, index) - let change_log = merkle_roll_apply_fn!( - header, - id, - roll_bytes, - initialize_with_root, - root, - leaf, - &proof, - index - )?; + let mut merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + merkle_roll_obj.initialize_with_root(root, leaf, &proof, index); + let change_log = get_changelog_after_op(merkle_roll_obj, id); wrap_event(change_log.try_to_vec()?, &ctx.accounts.candy_wrapper)?; emit!(*change_log); update_canopy(canopy_bytes, header.max_depth, Some(change_log)) } + + // TODO(sorend): We need to handle errors thrown by the PreAppend and MerkleRoll methods + /// Initializes the subtree append data structure for a given Gummyroll tree + pub fn init_or_reset_subtree_append_account( + ctx: Context + ) -> Result<()> { + let mut merkle_roll_bytes = ctx.accounts.merkle_roll.try_borrow_mut_data()?; + + let (mut header_bytes, rest) = + merkle_roll_bytes.split_at_mut(size_of::()); + + let mut header = Box::new(MerkleRollHeader::try_from_slice(&header_bytes)?); + let merkle_roll_size = merkle_roll_get_size!(header)?; + let (roll_bytes, _) = rest.split_at_mut(merkle_roll_size); + let merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + + // Assert that subtree_append is the correct size for header.max_depth + let merkle_roll_append_size = merkle_roll_append_get_size!(header)?; + let mut subtree_append_account_bytes = ctx.accounts.subtree_append.try_borrow_mut_data()?; + let (_, mut pre_append_struct_data) = subtree_append_account_bytes.split_at_mut(8); + assert!(pre_append_struct_data.len() == merkle_roll_append_size, "Append account has incorrect size"); + let mut pre_append_obj = get_pre_append_interface(header.max_depth, pre_append_struct_data)?; + pre_append_obj.reset(merkle_roll_obj); + Ok(()) + } + + /// Push a partition to the pre-append data structure + pub fn push_pre_append_partition( + ctx: Context, + rightmost_leaf: Node, + rightmost_proof: Vec + ) -> Result<()> { + let mut merkle_roll_bytes = ctx.accounts.merkle_roll.try_borrow_mut_data()?; + + let (mut header_bytes, rest) = + merkle_roll_bytes.split_at_mut(size_of::()); + + let mut header = Box::new(MerkleRollHeader::try_from_slice(&header_bytes)?); + let merkle_roll_size = merkle_roll_get_size!(header)?; + let (roll_bytes, _) = rest.split_at_mut(merkle_roll_size); + let merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + + // The current authority for the merkle_roll must sign be a signer + assert_eq!(header.authority, ctx.accounts.authority.key()); + + // Assert that subtree_append is the correct size for header.max_depth + let merkle_roll_append_size = merkle_roll_append_get_size!(header)?; + let mut subtree_append_account_bytes = ctx.accounts.subtree_append.try_borrow_mut_data()?; + let (_, mut pre_append_struct_data) = subtree_append_account_bytes.split_at_mut(8); + assert!(pre_append_struct_data.len() == merkle_roll_append_size, "Append account has incorrect size"); + let mut pre_append_obj = get_pre_append_interface(header.max_depth, pre_append_struct_data)?; + pre_append_obj.push_partition(merkle_roll_obj, rightmost_leaf, &rightmost_proof); + Ok(()) + } + + /// Append a subtree to a larger merkle roll in a byte dense way. Requires that the subtree's pre append data structure is initialized. + pub fn append_subtree( + ctx: Context + ) -> Result<()> { + // 1. load mutable bytes for merkle_roll to append to + let mut merkle_roll_bytes = ctx.accounts.merkle_roll.try_borrow_mut_data()?; + let (mut header_bytes, rest) = + merkle_roll_bytes.split_at_mut(size_of::()); + let mut header = Box::new(MerkleRollHeader::try_from_slice(&header_bytes)?); + let merkle_roll_size = merkle_roll_get_size!(header)?; + let (roll_bytes, _) = rest.split_at_mut(merkle_roll_size); + let mut merkle_roll_receiver_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + + // 2. load bytes for subtree_merkle_roll + let mut subtree_merkle_roll_bytes = ctx.accounts.subtree_merkle_roll.try_borrow_mut_data()?; + let (mut subtree_header_bytes, subtree_rest) = + subtree_merkle_roll_bytes.split_at_mut(size_of::()); + let mut subtree_header = Box::new(MerkleRollHeader::try_from_slice(&subtree_header_bytes)?); + let subtree_merkle_roll_size = merkle_roll_get_size!(subtree_header)?; + let (subtree_roll_bytes, _) = subtree_rest.split_at_mut(subtree_merkle_roll_size); + let mut merkle_roll_to_append = get_merkle_roll_interface(subtree_header.max_depth, subtree_header.max_buffer_size, subtree_roll_bytes)?; + + // 3. load bytes for subtree_append struct + let merkle_roll_append_size = merkle_roll_append_get_size!(subtree_header)?; + let mut pre_append_struct_bytes = ctx.accounts.subtree_append.try_borrow_mut_data()?; + let (_, mut pre_append_struct_data) = pre_append_struct_bytes.split_at_mut(8); + assert!(pre_append_struct_data.len() == merkle_roll_append_size, "Append account has incorrect size"); + let pre_append_obj = get_pre_append_interface(subtree_header.max_depth, pre_append_struct_data)?; + + assert!(merkle_roll_to_append.get_sequence_number() == pre_append_obj.get_sequence_number(), "tree to append changed since partitions pushed, invalid to append"); + + merkle_roll_receiver_obj.append_subtree_packed( + &pre_append_obj.get_rightmost_proofs_as_vec(), + &pre_append_obj.get_rightmost_leaves_as_vec(), + &pre_append_obj.get_rightmost_leaves_as_vec() + ); + + Ok(()) + } /// Executes an instruction that overwrites a leaf node. /// Composing programs should check that the data hashed into previous_leaf @@ -414,17 +588,9 @@ pub mod gummyroll { fill_in_proof_from_canopy(canopy_bytes, header.max_depth, index, &mut proof)?; let id = ctx.accounts.merkle_roll.key(); // A call is made to MerkleRoll::set_leaf(root, previous_leaf, new_leaf, proof, index) - let change_log = merkle_roll_apply_fn!( - header, - id, - roll_bytes, - set_leaf, - root, - previous_leaf, - new_leaf, - &proof, - index, - )?; + let mut merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + merkle_roll_obj.set_leaf(root, previous_leaf, new_leaf, &proof, index); + let change_log = get_changelog_after_op(merkle_roll_obj, id); wrap_event(change_log.try_to_vec()?, &ctx.accounts.candy_wrapper)?; emit!(*change_log); update_canopy(canopy_bytes, header.max_depth, Some(change_log)) @@ -470,7 +636,8 @@ pub mod gummyroll { fill_in_proof_from_canopy(canopy_bytes, header.max_depth, index, &mut proof)?; let id = ctx.accounts.merkle_roll.key(); - merkle_roll_apply_fn!(header, id, roll_bytes, prove_leaf, root, leaf, &proof, index)?; + let mut merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + merkle_roll_obj.prove_leaf(root, leaf, &proof, index); Ok(()) } @@ -489,7 +656,9 @@ pub mod gummyroll { let id = ctx.accounts.merkle_roll.key(); let merkle_roll_size = merkle_roll_get_size!(header)?; let (roll_bytes, canopy_bytes) = rest.split_at_mut(merkle_roll_size); - let change_log = merkle_roll_apply_fn!(header, id, roll_bytes, append, leaf)?; + let mut merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + merkle_roll_obj.append(leaf); + let change_log = get_changelog_after_op(merkle_roll_obj, id); wrap_event(change_log.try_to_vec()?, &ctx.accounts.candy_wrapper)?; emit!(*change_log); update_canopy(canopy_bytes, header.max_depth, Some(change_log)) @@ -519,16 +688,9 @@ pub mod gummyroll { fill_in_proof_from_canopy(canopy_bytes, header.max_depth, index, &mut proof)?; // A call is made to MerkleRoll::fill_empty_or_append let id = ctx.accounts.merkle_roll.key(); - let change_log = merkle_roll_apply_fn!( - header, - id, - roll_bytes, - fill_empty_or_append, - root, - leaf, - &proof, - index, - )?; + let mut merkle_roll_obj = get_merkle_roll_interface(header.max_depth, header.max_buffer_size, roll_bytes)?; + merkle_roll_obj.fill_empty_or_append(root, leaf, &proof, index); + let change_log = get_changelog_after_op(merkle_roll_obj, id); wrap_event(change_log.try_to_vec()?, &ctx.accounts.candy_wrapper)?; emit!(*change_log); update_canopy(canopy_bytes, header.max_depth, Some(change_log)) diff --git a/contracts/programs/gummyroll/src/state/mod.rs b/contracts/programs/gummyroll/src/state/mod.rs index e250b53c..33ea1b53 100644 --- a/contracts/programs/gummyroll/src/state/mod.rs +++ b/contracts/programs/gummyroll/src/state/mod.rs @@ -2,7 +2,7 @@ //! use anchor_lang::prelude::*; use borsh::{BorshDeserialize, BorshSerialize}; -use concurrent_merkle_tree::state::{ChangeLog, Node}; +use concurrent_merkle_tree::state::{ChangeLog, ChangeLogInterface, Node, Path}; #[derive(AnchorDeserialize, AnchorSerialize, Clone, Copy, Debug)] pub struct PathNode { @@ -41,30 +41,30 @@ pub struct ChangeLogEvent { /// Bitmap of node parity (used when hashing) pub index: u32, } -// ChangeLog -impl From<(Box>, Pubkey, u64)> + +impl From<(Box, Pubkey, u64)> for Box { - fn from(log_info: (Box>, Pubkey, u64)) -> Self { + fn from(log_info: (Box, Pubkey, u64)) -> Self { let (changelog, tree_id, seq) = log_info; - let path_len = changelog.path.len() as u32; + let path_len = changelog.get_path_as_vec().len() as u32; let mut path: Vec = changelog - .path + .get_path_as_vec() .iter() .enumerate() .map(|(lvl, n)| { PathNode::new( *n, - (1 << (path_len - lvl as u32)) + (changelog.index >> lvl), + (1 << (path_len - lvl as u32)) + (changelog.get_index() >> lvl), ) }) .collect(); - path.push(PathNode::new(changelog.root, 1)); + path.push(PathNode::new(changelog.get_root(), 1)); Box::new(ChangeLogEvent { id: tree_id, path, seq, - index: changelog.index, + index: changelog.get_index(), }) } } diff --git a/contracts/programs/gummyroll/src/utils.rs b/contracts/programs/gummyroll/src/utils.rs index f6f114cb..26a08a00 100644 --- a/contracts/programs/gummyroll/src/utils.rs +++ b/contracts/programs/gummyroll/src/utils.rs @@ -6,7 +6,7 @@ use anchor_lang::{ solana_program::{msg, program::invoke, program_error::ProgramError}, }; use bytemuck::{Pod, PodCastError}; -use concurrent_merkle_tree::merkle_roll::MerkleRoll; +use concurrent_merkle_tree::merkle_roll::{MerkleRoll, MerkleRollPreAppend}; use std::any::type_name; use std::mem::size_of; @@ -37,6 +37,11 @@ impl ZeroCopy { } +impl ZeroCopy + for MerkleRollPreAppend +{ +} + pub fn error_msg(data_len: usize) -> impl Fn(PodCastError) -> ProgramError { move |_: PodCastError| -> ProgramError { msg!( diff --git a/lib/concurrent-merkle-tree/src/error.rs b/lib/concurrent-merkle-tree/src/error.rs index 80dd52f4..2144de87 100644 --- a/lib/concurrent-merkle-tree/src/error.rs +++ b/lib/concurrent-merkle-tree/src/error.rs @@ -35,4 +35,12 @@ pub enum CMTError { "Valid proof was passed to a leaf, but it's value has changed since the proof was issued" )] LeafContentsModified, + + /// Attempted to append subtree of invalid size for the current state of the tree + #[error("Cannot append subtree with invalid size")] + SubtreeInvalidSize, + + /// Attempted to index into proof beyond rightmost bound + #[error("Invalid index into rightmost_proof")] + InvalidProofAccessOrWrite, } diff --git a/lib/concurrent-merkle-tree/src/merkle_roll.rs b/lib/concurrent-merkle-tree/src/merkle_roll.rs index f774b4ca..2e2f2602 100644 --- a/lib/concurrent-merkle-tree/src/merkle_roll.rs +++ b/lib/concurrent-merkle-tree/src/merkle_roll.rs @@ -1,7 +1,7 @@ use crate::{ error::CMTError, - state::{ChangeLog, Node, Path, EMPTY}, - utils::{empty_node, empty_node_cached, fill_in_proof, recompute, hash_to_parent}, + state::{ChangeLog, ChangeLogInterface, Node, Path, EMPTY}, + utils::{empty_node, empty_node_cached, fill_in_proof, hash_to_parent, recompute}, }; use bytemuck::{Pod, Zeroable}; pub(crate) use log_compute; @@ -16,6 +16,68 @@ fn check_bounds(max_depth: usize, max_buffer_size: usize) { // This will return true if MAX_BUFFER_SIZE is a power of 2 or if it is 0 assert!(max_buffer_size & (max_buffer_size - 1) == 0); } + +pub trait MerkleInterface { + fn get_sequence_number(&self) -> u64; + fn get_rightmost_proof_node_at_index(&self, i: usize) -> Result<[u8; 32], CMTError>; + fn set_rightmost_proof_node_at_index(&mut self, i: usize, node_to_write: [u8; 32]) -> Result<(), CMTError>; + fn get_rightmost_proof_index(&self) -> u32; + fn get_rightmost_proof_leaf(&self) -> [u8; 32]; + fn get_max_depth(&self) -> usize; + fn initialize(&mut self) -> Result; + + fn initialize_with_root( + &mut self, + root: Node, + rightmost_leaf: Node, + proof_vec: &Vec, + index: u32, + ) -> Result; + fn get_change_log(&self) -> Box; + + fn prove_leaf( + &mut self, + current_root: Node, + leaf: Node, + proof_vec: &Vec, + leaf_index: u32, + ) -> Result; + + fn append(&mut self, node: Node) -> Result; + + fn append_subtree_direct( + &mut self, + subtree_root: Node, + subtree_rightmost_leaf: Node, + subtree_rightmost_index: u32, + subtree_rightmost_proof: &Vec, + ) -> Result; + + fn append_subtree_packed( + &mut self, + subtree_proofs: &Vec>, + subtree_rightmost_leaves: &Vec, + subtree_roots: &Vec + ) -> Result; + + fn fill_empty_or_append( + &mut self, + current_root: Node, + leaf: Node, + proof_vec: &Vec, + index: u32, + ) -> Result; + + fn set_leaf( + &mut self, + current_root: Node, + previous_leaf: Node, + new_leaf: Node, + proof_vec: &Vec, + index: u32, + ) -> Result; +} + /// Tracks updates to off-chain Merkle tree /// /// Allows for concurrent writes to same merkle tree so long as proof @@ -41,18 +103,41 @@ unsafe impl Pod { } -impl MerkleRoll { - pub fn new() -> Self { - Self { - sequence_number: 0, - active_index: 0, - buffer_size: 0, - change_logs: [ChangeLog::::default(); MAX_BUFFER_SIZE], - rightmost_proof: Path::::default(), +impl MerkleInterface for MerkleRoll { + fn get_sequence_number(&self) -> u64 { + return self.sequence_number; + } + + fn get_rightmost_proof_node_at_index(&self, i: usize) -> Result<[u8; 32], CMTError> { + if i < self.rightmost_proof.proof.len() { + return Ok(self.rightmost_proof.proof[i]); + } else { + return Err(CMTError::InvalidProofAccessOrWrite); } } - pub fn initialize(&mut self) -> Result { + fn set_rightmost_proof_node_at_index(&mut self, i: usize, node_to_write: [u8; 32]) -> Result<(), CMTError> { + if i < self.rightmost_proof.proof.len() { + self.rightmost_proof.proof[i] = node_to_write; + return Ok(()); + } else { + return Err(CMTError::InvalidProofAccessOrWrite); + } + } + + fn get_rightmost_proof_leaf(&self) -> [u8; 32] { + return self.rightmost_proof.leaf; + } + + fn get_rightmost_proof_index(&self) -> u32 { + return self.rightmost_proof.index; + } + + fn get_max_depth(&self) -> usize { + return MAX_DEPTH; + } + + fn initialize(&mut self) -> Result { check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE); let mut rightmost_proof = Path::default(); let mut empty_node_cache = Box::new([Node::default(); MAX_DEPTH]); @@ -72,7 +157,7 @@ impl MerkleRoll MerkleRoll Box> { + fn get_change_log(&self) -> Box { Box::new(self.change_logs[self.active_index as usize]) } - pub fn prove_leaf( + fn prove_leaf( &mut self, current_root: Node, leaf: Node, @@ -129,41 +214,26 @@ impl MerkleRoll Result { - let old_root = recompute(EMPTY, &proof, 0); - if old_root == empty_node(MAX_DEPTH as u32) { - self.try_apply_proof(old_root, EMPTY, leaf, &mut proof, 0, false) - } else { - return Err(CMTError::TreeAlreadyInitialized); - } - } - - /// Basic operation that always succeeds - pub fn append(&mut self, mut node: Node) -> Result { + fn append(&mut self, node: Node) -> Result { + let mut mut_node = node; check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE); - if node == EMPTY { + if mut_node == EMPTY { return Err(CMTError::CannotAppendEmptyNode); } if self.rightmost_proof.index >= 1 << MAX_DEPTH { return Err(CMTError::TreeFull); } if self.rightmost_proof.index == 0 { - return self.initialize_tree_from_append(node, self.rightmost_proof.proof); + return self.initialize_tree_from_append(mut_node, self.rightmost_proof.proof); } - let leaf = node.clone(); + let leaf = mut_node.clone(); let intersection = self.rightmost_proof.index.trailing_zeros() as usize; let mut change_list = [EMPTY; MAX_DEPTH]; let mut intersection_node = self.rightmost_proof.leaf; let mut empty_node_cache = Box::new([Node::default(); MAX_DEPTH]); for i in 0..MAX_DEPTH { - change_list[i] = node; + change_list[i] = mut_node; if i < intersection { // Compute proof to the appended node from empty nodes let sibling = empty_node_cached::(i as u32, &mut empty_node_cache); @@ -172,10 +242,98 @@ impl MerkleRoll> i) & 1 == 0, ); - hash_to_parent(&mut node, &sibling, true); + hash_to_parent(&mut mut_node, &sibling, true); self.rightmost_proof.proof[i] = sibling; } else if i == intersection { // Compute the where the new node intersects the main tree + hash_to_parent(&mut mut_node, &intersection_node, false); + self.rightmost_proof.proof[intersection] = intersection_node; + } else { + // Update the change list path up to the root + hash_to_parent( + &mut mut_node, + &self.rightmost_proof.proof[i], + ((self.rightmost_proof.index - 1) >> i) & 1 == 0, + ); + } + } + self.update_state_from_append(mut_node, change_list, self.rightmost_proof.index, leaf)?; + Ok(mut_node) + } + + /// Append subtree to current tree + fn append_subtree_direct( + &mut self, + subtree_root: Node, + subtree_rightmost_leaf: Node, + subtree_rightmost_index: u32, + subtree_rightmost_proof: &Vec, + ) -> Result { + check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE); + if self.rightmost_proof.index >= 1 << MAX_DEPTH { + return Err(CMTError::TreeFull); + } + + // If the rightmost proof is empty, then we just append the rightmost leaf + if subtree_rightmost_proof.len() == 0 { + return self.append(subtree_rightmost_leaf); + } + + // Confirm that subtree_rightmost_proof is valid + if recompute( + subtree_rightmost_leaf, + &subtree_rightmost_proof[..], + subtree_rightmost_index - 1, + ) != subtree_root + { + return Err(CMTError::InvalidProof); + } + + let intersection = self.rightmost_proof.index.trailing_zeros() as usize; + + if self.rightmost_proof.index == 0 { + // @dev: If the tree has only empty leaves, then there are many sizes of tree which can be appended, but they cannot be larger than the tree itself + if subtree_rightmost_proof.len() >= MAX_DEPTH { + return Err(CMTError::SubtreeInvalidSize); + } + return self.initialize_tree_from_subtree_append( + subtree_root, + subtree_rightmost_leaf, + subtree_rightmost_index, + subtree_rightmost_proof, + ); + } else { + // @dev: At any given time (other than initialization), there is only one valid size of subtree that can be appended + if subtree_rightmost_proof.len() != intersection { + return Err(CMTError::SubtreeInvalidSize); + } + } + + let leaf = subtree_rightmost_leaf.clone(); + let mut change_list = [EMPTY; MAX_DEPTH]; + let mut intersection_node = self.rightmost_proof.leaf; + + // This will be mutated into the new root after the append by gradually hashing this node with the RMP to the subtree, then the critical node, and then the rest of the RMP to this tree. + let mut node = subtree_rightmost_leaf; + + for i in 0..MAX_DEPTH { + change_list[i] = node; + if i < intersection { + // Compute proof to the appended node from empty nodes + hash_to_parent( + &mut intersection_node, + &self.rightmost_proof.proof[i], + ((self.rightmost_proof.index - 1) >> i) & 1 == 0, + ); + hash_to_parent( + &mut node, + &subtree_rightmost_proof[i], + ((subtree_rightmost_index - 1) >> i) & 1 == 0, + ); + self.rightmost_proof.proof[i] = subtree_rightmost_proof[i]; + } else if i == intersection { + // Compute where the new node intersects the main tree + assert!(node == subtree_root); hash_to_parent(&mut node, &intersection_node, false); self.rightmost_proof.proof[intersection] = intersection_node; } else { @@ -187,19 +345,76 @@ impl MerkleRoll::new(node, change_list, self.rightmost_proof.index); - self.rightmost_proof.index = self.rightmost_proof.index + 1; - self.rightmost_proof.leaf = leaf; + self.update_state_from_append( + node, + change_list, + self.rightmost_proof.index + subtree_rightmost_index - 1, + leaf, + )?; Ok(node) } + /// @dev: this function appends a subtree of any size at any time (as long as there is sufficient capacity in the large tree, but irrespective of the next free index) + /// this requires the subtree (with 2^n leaves) to be unpacked into n+1 sub-subtrees of size (# leaves), 1,1,..,2^(n-2),2^(n-1) + /// as the consumer of this function you must pass the root, right most leaf and right most proof to the sub-subtree of size 2^(n-1) [as the last element in each of the vector parameters], etc. where the first element is a tree of size 1 + /// @notice: in all cases, this function is much slower and more expensive than its counterpart "append_subtree_direct". The tradeoff is that this + /// method will always succeed as long as the larger tree has enough total capacity to fit the subtree being appended, where as a direct append will fail + /// if the subtree cannot be appended by appending its complete topological form starting with the next available index in the larger tree + fn append_subtree_packed( + &mut self, + subtree_proofs: &Vec>, + subtree_rightmost_leaves: &Vec, + subtree_roots: &Vec + ) -> Result { + // 1. determine index range for each sub-subtree to append of size 2^k + // a. find the index at which to insert 2^(n-1): m + // b. compute m-l (leftmost index). Insert trees based on the binary representation of m-l (starting with the lsbs) -> each 1 in the binary representation of m-l maps to a sub-subtree to be appended + let n = subtree_proofs.len()-1; + let largest_power: u32 = 1 << (n-1); + + // TODO(sorend): to be more efficient we can use a u32 as a bitset rather than a vector of bool == vector of char (8x inefficient) + let mut index_appended = vec![false; n+1]; + // find the first index in constant time => some i st i % (2^(n-1)) == 0 + let mut largest_tree_insertion_index = self.rightmost_proof.index; + let remainder = largest_tree_insertion_index % largest_power; + if remainder != 0 { + largest_tree_insertion_index = largest_tree_insertion_index + (largest_power - remainder); + } + + let first_difference = largest_tree_insertion_index - self.rightmost_proof.index; + let num_bits = first_difference.count_ones() + first_difference.count_zeros(); + + // Append trees ascending up to largest_tree_insertion_index + for i in 0..num_bits { + // If the i'th lsb is a 1, then we want to insert the appropriate tree + if first_difference >> i & 1 == 1 { + let ind = (i+1) as usize; + self.append_subtree_direct(subtree_roots[ind], subtree_rightmost_leaves[ind], 1 << subtree_proofs[ind].len(), &subtree_proofs[ind])?; + index_appended[ind] = true; + } + } + + // 2. append the largest sub-subtree at index m + // Append the largest sub-subtree at the largest tree insertion index + self.append_subtree_direct(subtree_roots[n], subtree_rightmost_leaves[n], 1 << subtree_proofs[n].len(), &subtree_proofs[n])?; + index_appended[n] = true; + + // 3. append the rest of the remaining sub-subtrees in order from largest to smallest which have not yet been inserted + // Iterate over the remaining sub-subtrees that need to be appended in order from largest to smallest (indices [n-1, 0]) + let mut root = EMPTY; + for i in (0..n).rev() { + if !index_appended[i] { + root = self.append_subtree_direct(subtree_roots[i], subtree_rightmost_leaves[i], 1 << subtree_proofs[i].len(), &subtree_proofs[i])?; + } + } + + Ok(root) + } + /// Convenience function for `set_leaf` /// On write conflict: /// Will append - pub fn fill_empty_or_append( + fn fill_empty_or_append( &mut self, current_root: Node, leaf: Node, @@ -223,7 +438,7 @@ impl MerkleRoll MerkleRoll MerkleRoll { + + pub fn new() -> Self { + Self { + sequence_number: 0, + active_index: 0, + buffer_size: 0, + change_logs: [ChangeLog::::default(); MAX_BUFFER_SIZE], + rightmost_proof: Path::::default(), + } + } + + /// Only used to initialize right most path for a completely empty tree + #[inline(always)] + fn initialize_tree_from_append( + &mut self, + leaf: Node, + mut proof: [Node; MAX_DEPTH], + ) -> Result { + let old_root = recompute(EMPTY, &proof, 0); + if old_root == empty_node(MAX_DEPTH as u32) { + self.try_apply_proof(old_root, EMPTY, leaf, &mut proof, 0, false) + } else { + return Err(CMTError::TreeAlreadyInitialized); + } + } + + fn update_state_from_append( + &mut self, + root: Node, + change_list: [Node; MAX_DEPTH], + rmp_index: u32, + rmp_leaf: Node, + ) -> Result<(), CMTError> { + self.update_internal_counters(); + self.change_logs[self.active_index as usize] = + ChangeLog::::new(root, change_list, rmp_index); + self.rightmost_proof.index = rmp_index + 1; + self.rightmost_proof.leaf = rmp_leaf; + Ok(()) + } + + fn initialize_tree_from_subtree_append( + &mut self, + subtree_root: Node, + subtree_rightmost_leaf: Node, + subtree_rightmost_index: u32, + subtree_rightmost_proof: &Vec, + ) -> Result { + let leaf = subtree_rightmost_leaf.clone(); + let mut change_list = [EMPTY; MAX_DEPTH]; + + // This will be mutated into the new root after the append by gradually hashing this node with the RMP to the subtree, then the critical node, and then the rest of the RMP to this tree. + let mut node = subtree_rightmost_leaf; + for i in 0..MAX_DEPTH { + change_list[i] = node; + if i < subtree_rightmost_proof.len() { + // Hash up to subtree_root using subtree_rmp, to create accurate change_list + hash_to_parent( + &mut node, + &subtree_rightmost_proof[i], + ((subtree_rightmost_index - 1) >> i) & 1 == 0, + ); + self.rightmost_proof.proof[i] = subtree_rightmost_proof[i]; + } else { + // Compute where the new node intersects the main tree + if i == subtree_rightmost_proof.len() { + assert!(node == subtree_root); + } + + hash_to_parent(&mut node, &self.rightmost_proof.proof[i], true); + // No need to update the RMP anymore + } + } + self.update_state_from_append( + node, + change_list, + self.rightmost_proof.index + subtree_rightmost_index - 1, + leaf, + )?; + Ok(node) + } /// Modifies the `proof` for leaf at `leaf_index` /// in place by fast-forwarding the given `proof` through the @@ -342,7 +641,7 @@ impl MerkleRoll MerkleRoll= 1 << merkle_roll.get_max_depth(), "Cannot begin constructing pre-append data structure before tree is full"); +} + +pub trait PreAppendInterface { + fn get_sequence_number(&self) -> u64; + fn get_rightmost_proofs_as_vec(&self) -> Vec>; + fn get_rightmost_leaves_as_vec(&self) -> Vec; + fn get_partition_roots_as_vec(&self) -> Vec; + fn get_initialized(&self) -> bool; + fn reset(&mut self, merkle_roll: &dyn MerkleInterface) -> Result<(), CMTError>; + fn push_partition(&mut self, merkle_roll: &dyn MerkleInterface, rightmost_leaf: Node, rightmost_proof: &Vec) -> Result<(), CMTError>; +} +/// Stores merkle tree partitions in storage due to transaction size constraints. +/// +/// If a user wants to append their tree as a subtree to some other tree. They should populate this structure. +/// This structure should be filled by providing partitions with the following number of leaves in order: 1,1,2,4,...2^(k+1) where 2^k is the number of leaves in the associated merkle_roll struct +/// The subtrees should be passed in order from smallest to largest, greedily taking the smaller trees from the rightmost side of the tree, such that the root of each partitioned tree aligns with the ith index into the rightmost proof of the larger tree +#[derive(Copy, Clone)] +pub struct MerkleRollPreAppend { + /// The next depth of tree which needs to be initialized + pub next_index_to_initialize: u32, + /// Whether or not the structure has been fully initialized with NUM_PARTITIONS trees + pub initialized: u32, + /// The sequence number of the MerkleRoll when the initialization of this struct began + pub sequence_number: u64, + /// Right most proofs to each partitioned subtree + pub rightmost_proofs: [Path; NUM_PARTITIONS], + /// Rightmost leaves of each partitioned subtree + pub partition_rightmost_leaves: [Node; NUM_PARTITIONS], + /// Roots of partitioned trees + pub partition_roots: [Node; NUM_PARTITIONS] +} + +unsafe impl Zeroable + for MerkleRollPreAppend +{ +} +unsafe impl Pod + for MerkleRollPreAppend +{ +} + +impl PreAppendInterface for MerkleRollPreAppend { + fn get_sequence_number(&self) -> u64 { + return self.sequence_number; + } + + fn get_rightmost_proofs_as_vec(&self) -> Vec> { + let proofs_vec: Vec> = vec![vec![], vec![]]; + for i in 2..NUM_PARTITIONS { + let mut proof: Vec = vec![]; + for j in 0..i-1 { + proof.push(self.rightmost_proofs[i].proof[j]); + } + } + return proofs_vec; + } + + fn get_rightmost_leaves_as_vec(&self) -> Vec { + let mut leaf_vec: Vec = vec![]; + for leaf in self.partition_rightmost_leaves { + leaf_vec.push(leaf); + } + return leaf_vec; + } + + fn get_partition_roots_as_vec(&self) -> Vec { + let mut root_vec: Vec = vec![]; + for root in self.partition_roots { + root_vec.push(root); + } + return root_vec; + } + + fn get_initialized(&self) -> bool { + return self.initialized == 1; + } + + /// Reset the pre-append data structure to a zero-state configuration. + /// @dev: doubles as an initialization function + fn reset(&mut self, merkle_roll: &dyn MerkleInterface) -> Result<(), CMTError> { + assert_valid_pre_append_state(merkle_roll, NUM_PARTITIONS); + self.set_empty_values(merkle_roll)?; + Ok(()) + } + + /// @dev: Push a new tree partition. We expect partitions to be pushed with num leaves 1,1,2,4.. IN THAT ORDER. The first partitioned tree should be the rightmost_leaf of merkle_roll associated with this struct. + /// The second partition pushed should be the second rightmost leaf in the tree. Followed by the rightmost subtree of size 2 which has not yet been included etc. + /// While you are pushing partitions to this struct, the associated merkle_roll should not change. If it does it would invalidate some of the pushed state, and you should call reset and begin pushing partitions from the beginning. + fn push_partition(&mut self, merkle_roll: &dyn MerkleInterface, rightmost_leaf: Node, rightmost_proof: &Vec) -> Result<(), CMTError> { + assert_valid_pre_append_state(merkle_roll, NUM_PARTITIONS); + assert!(merkle_roll.get_sequence_number() == self.sequence_number, "Tree has been modified while pushing proofs. State invalid, please reset."); + let index = self.next_index_to_initialize as usize; + if index == 0 { + assert!(rightmost_proof.len() == 0); + assert!(merkle_roll.get_rightmost_proof_leaf() == rightmost_leaf); + } + else if index == 1 { + assert!(rightmost_proof.len() == 0); + assert!(merkle_roll.get_rightmost_proof_node_at_index(0)? == rightmost_leaf); + } + else { + assert!(rightmost_proof.len() == index - 1); + assert!(recompute(rightmost_leaf, &rightmost_proof[..], (1 << index) - 1) == merkle_roll.get_rightmost_proof_node_at_index(index-1)?); + // For the single leaves, we don't need to update their roots since they are just leaves by themselves + self.partition_roots[index] = merkle_roll.get_rightmost_proof_node_at_index(index-1)?; + } + // Copy the supplied proof into the data structure. Note that indices from rightmost_proof.len() -> NUM_PARTITIONS will be garbage and should be ignored in whatever is passed to append_direct + for i in 0..rightmost_proof.len() { + self.rightmost_proofs[index].proof[i] = rightmost_proof[i]; + } + self.partition_rightmost_leaves[index] = rightmost_leaf; + self.next_index_to_initialize += 1; + if self.next_index_to_initialize as usize >= NUM_PARTITIONS { + self.initialized = 1; + } + Ok(()) + } +} + +impl MerkleRollPreAppend { + pub fn new(merkle_roll: &dyn MerkleInterface) -> Self { + Self { + rightmost_proofs: [Path::::default(); NUM_PARTITIONS], + partition_rightmost_leaves: [Node::default(); NUM_PARTITIONS], + partition_roots: [Node::default(); NUM_PARTITIONS], + next_index_to_initialize: 0, + initialized: 0, + sequence_number: merkle_roll.get_sequence_number() + } + } + + fn set_empty_values(&mut self, merkle_roll: &dyn MerkleInterface) -> Result<(), CMTError> { + assert_valid_pre_append_state(merkle_roll, NUM_PARTITIONS); + self.rightmost_proofs = [Path::::default(); NUM_PARTITIONS]; + self.partition_rightmost_leaves = [Node::default(); NUM_PARTITIONS]; + self.partition_roots = [Node::default(); NUM_PARTITIONS]; + self.next_index_to_initialize = 0; + self.initialized = 0; + self.sequence_number = merkle_roll.get_sequence_number(); + Ok(()) + } +} diff --git a/lib/concurrent-merkle-tree/src/state.rs b/lib/concurrent-merkle-tree/src/state.rs index 81a2b047..00e197c5 100644 --- a/lib/concurrent-merkle-tree/src/state.rs +++ b/lib/concurrent-merkle-tree/src/state.rs @@ -1,5 +1,10 @@ use crate::utils::hash_to_parent; +pub trait ChangeLogInterface { + fn get_root(&self) -> Node; + fn get_path_as_vec(&self) -> Vec; + fn get_index(&self) -> u32; +} #[derive(Copy, Clone, Debug, PartialEq)] /// Stores proof for a given Merkle root update #[repr(C)] @@ -13,6 +18,20 @@ pub struct ChangeLog { pub _padding: u32, } +impl ChangeLogInterface for ChangeLog { + fn get_root(&self) -> Node { + return self.root; + } + + fn get_path_as_vec(&self) -> Vec { + return self.path.to_vec(); + } + + fn get_index(&self) -> u32 { + return self.index; + } +} + impl ChangeLog { pub fn default() -> Self { Self { @@ -37,7 +56,12 @@ impl ChangeLog { } /// Sets all change log values from a leaf and valid proof - pub fn replace_and_recompute_path(&mut self, index: u32, mut node: Node, proof: &[Node]) -> Node { + pub fn replace_and_recompute_path( + &mut self, + index: u32, + mut node: Node, + proof: &[Node], + ) -> Node { self.index = index; for (i, sibling) in proof.iter().enumerate() { self.path[i] = node; diff --git a/lib/concurrent-merkle-tree/tests/tests.rs b/lib/concurrent-merkle-tree/tests/tests.rs index 78486e11..e7caffb9 100644 --- a/lib/concurrent-merkle-tree/tests/tests.rs +++ b/lib/concurrent-merkle-tree/tests/tests.rs @@ -1,6 +1,7 @@ use concurrent_merkle_tree::error::CMTError; -use concurrent_merkle_tree::merkle_roll::MerkleRoll; +use concurrent_merkle_tree::merkle_roll::{MerkleRoll, MerkleInterface}; use concurrent_merkle_tree::state::{Node, EMPTY}; +use concurrent_merkle_tree::utils::{recompute, hash_to_parent}; use merkle_tree_reference::MerkleTree; use rand::thread_rng; use rand::{self, Rng}; @@ -9,19 +10,22 @@ use tokio; const DEPTH: usize = 14; const BUFFER_SIZE: usize = 64; -fn setup() -> (MerkleRoll, MerkleTree) { - // On-chain merkle change-record - let merkle = MerkleRoll::::new(); - - // Init off-chain Merkle tree with corresponding # of leaves +fn make_empty_offchain_tree_of_depth(depth: usize) -> MerkleTree { let mut leaves = vec![]; - for _ in 0..(1 << DEPTH) { + for _ in 0..(1 << depth) { let leaf = EMPTY; leaves.push(leaf); } - // Off-chain merkle tree - let reference_tree = MerkleTree::new(leaves); + MerkleTree::new(leaves) +} + +fn setup() -> (MerkleRoll, MerkleTree) { + // On-chain merkle change-record + let merkle = MerkleRoll::::new(); + + // Init off-chain Merkle tree with corresponding # of leaves + let reference_tree = make_empty_offchain_tree_of_depth(DEPTH); (merkle, reference_tree) } @@ -32,7 +36,7 @@ async fn test_initialize() { merkle_roll.initialize().unwrap(); assert_eq!( - merkle_roll.get_change_log().root, + merkle_roll.get_change_log().get_root(), off_chain_tree.get_root(), "Init failed to set root properly" ); @@ -49,7 +53,7 @@ async fn test_append() { merkle_roll.append(leaf).unwrap(); off_chain_tree.add_leaf(leaf, i); assert_eq!( - merkle_roll.get_change_log().root, + merkle_roll.get_change_log().get_root(), off_chain_tree.get_root(), "On chain tree failed to update properly on an append", ); @@ -61,6 +65,315 @@ async fn test_append() { ); } +#[tokio::test(threaded_scheduler)] +async fn test_append_complete_subtree() { + let (mut merkle_roll, mut off_chain_tree) = setup(); + let mut rng = thread_rng(); + merkle_roll.initialize().unwrap(); + + // insert eight leaves into the large tree + for i in 0..8 { + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + } + + // create a simple subtree to append of depth three + let mut onchain_subtree = MerkleRoll::<3, 8>::new(); + onchain_subtree.initialize().unwrap(); + + // completely fill the subtree with unique leaves, and also append them to the off-chain tree + for i in 8..16 { + let leaf = rng.gen::<[u8; 32]>(); + onchain_subtree.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + } + + // append the on_chain subtree to the merkle_roll + merkle_roll + .append_subtree_direct( + onchain_subtree.get_change_log().get_root(), + onchain_subtree.rightmost_proof.leaf, + onchain_subtree.rightmost_proof.index, + &onchain_subtree.rightmost_proof.proof.to_vec(), + ) + .unwrap(); + + // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..15 appended one by one + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + + // Show that we can still append to the large tree after performing a subtree append + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, 16); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Failed to append accurately to merkle roll after subtree append", + ); +} + +#[tokio::test(threaded_scheduler)] +async fn test_append_incomplete_subtree() { + let (mut merkle_roll, mut off_chain_tree) = setup(); + let mut rng = thread_rng(); + merkle_roll.initialize().unwrap(); + + // insert four leaves into the large tree + for i in 0..4 { + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + } + + // create a simple subtree to append of depth three + let mut onchain_subtree = MerkleRoll::<2, 8>::new(); + onchain_subtree.initialize().unwrap(); + + // append leaves to the subtree, and also append them to the off-chain tree + // note: this gives us a partially filled tree, the other two leaves are empty nodes + for i in 4..6 { + let leaf = rng.gen::<[u8; 32]>(); + onchain_subtree.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + } + + // append the on_chain subtree to the merkle_roll + merkle_roll + .append_subtree_direct( + onchain_subtree.get_change_log().get_root(), + onchain_subtree.rightmost_proof.leaf, + onchain_subtree.rightmost_proof.index, + &onchain_subtree.rightmost_proof.proof.to_vec(), + ) + .unwrap(); + + // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..5 appended one by one + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + + // Show that we can still append to the large tree after performing a subtree append + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, 6); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Failed to append accurately to merkle roll after subtree append", + ); +} + +#[tokio::test(threaded_scheduler)] +async fn test_append_subtree_to_empty_tree() { + let (mut merkle_roll, mut off_chain_tree) = setup(); + let mut rng = thread_rng(); + merkle_roll.initialize().unwrap(); + + // create a simple subtree to append of depth two + let mut onchain_subtree = MerkleRoll::<2, 8>::new(); + onchain_subtree.initialize().unwrap(); + + // append leaves to the subtree, and also append them to the off-chain tree + for i in 0..4 { + let leaf = rng.gen::<[u8; 32]>(); + onchain_subtree.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + } + + // append the on_chain subtree to the merkle_roll + merkle_roll + .append_subtree_direct( + onchain_subtree.get_change_log().get_root(), + onchain_subtree.rightmost_proof.leaf, + onchain_subtree.rightmost_proof.index, + &onchain_subtree.rightmost_proof.proof.to_vec(), + ) + .unwrap(); + + // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..4 appended one by one + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + + // Show that we can still append to the large tree after performing a subtree append + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, 4); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Failed to append accurately to merkle roll after subtree append", + ); +} + +// Working on this, lets make it simpler +#[tokio::test(threaded_scheduler)] +async fn test_append_complete_subtree_tightly_packed_depth_three() { + let (mut merkle_roll, mut off_chain_tree) = setup(); + let mut rng = thread_rng(); + merkle_roll.initialize().unwrap(); + + // insert one leaf into the larger tree + for i in 0..1 { + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Large tree failed to update properly on an append", + ); + } + + let (mut small_off_chain_tree) = make_empty_offchain_tree_of_depth(3); + + // completely fill the subtree with unique leaves, and also append them to the off-chain tree + let mut leaves_in_small_tree = vec![]; + for i in 0..8 { + let leaf = rng.gen::<[u8; 32]>(); + small_off_chain_tree.add_leaf(leaf, i); + leaves_in_small_tree.push(leaf); + } + + // The order of the leaves will be shuffled around based on the tight append algo. The order below was manually calculated based on the test case situation + off_chain_tree.add_leaf(leaves_in_small_tree[6], 1); + off_chain_tree.add_leaf(leaves_in_small_tree[4], 2); + off_chain_tree.add_leaf(leaves_in_small_tree[5], 3); + off_chain_tree.add_leaf(leaves_in_small_tree[0], 4); + off_chain_tree.add_leaf(leaves_in_small_tree[1], 5); + off_chain_tree.add_leaf(leaves_in_small_tree[2], 6); + off_chain_tree.add_leaf(leaves_in_small_tree[3], 7); + off_chain_tree.add_leaf(leaves_in_small_tree[7], 8); + + // Mock the creation of the pre-append data structure + let subtree_proofs: Vec> = vec![vec![], vec![], vec![small_off_chain_tree.get_node(4)], vec![small_off_chain_tree.get_node(2), small_off_chain_tree.get_proof_of_leaf(3)[1]]]; + let subtree_rmls: Vec = vec![small_off_chain_tree.get_node(7), small_off_chain_tree.get_node(6), small_off_chain_tree.get_node(5), small_off_chain_tree.get_node(3)]; + let subtree_roots: Vec = vec![small_off_chain_tree.get_node(7), small_off_chain_tree.get_node(6), small_off_chain_tree.get_proof_of_leaf(7)[1], small_off_chain_tree.get_proof_of_leaf(7)[2]]; + + // append the small_subtree to merkle_roll + merkle_roll + .append_subtree_packed( + &subtree_proofs, + &subtree_rmls, + &subtree_roots + ) + .unwrap(); + + // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..9 appended one by one + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + + // The index of the rmp to the on_chain_merkle roll should be equivalent to us having performed nine appends, since the append should be dense + assert_eq!( + merkle_roll.rightmost_proof.index, + 9, + "On chain append was not tightly packed" + ); + + // Show that we can still append to the large tree after performing a subtree append + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, 9); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Failed to append accurately to merkle roll after subtree append", + ); +} + +#[tokio::test(threaded_scheduler)] +async fn test_append_complete_subtree_tightly_packed_depth_one() { + let (mut merkle_roll, mut off_chain_tree) = setup(); + let mut rng = thread_rng(); + merkle_roll.initialize().unwrap(); + + // insert one leaf into the larger tree + for i in 0..1 { + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, i); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Large tree failed to update properly on an append", + ); + } + + let (mut small_off_chain_tree) = make_empty_offchain_tree_of_depth(1); + + // completely fill the subtree with unique leaves + let mut leaves_in_small_tree = vec![]; + for i in 0..2 { + let leaf = rng.gen::<[u8; 32]>(); + small_off_chain_tree.add_leaf(leaf, i); + leaves_in_small_tree.push(leaf); + } + + // Append the leaves of the subtree to the bigger offchain tree in the order that the tightly packed algo will apply them + off_chain_tree.add_leaf(leaves_in_small_tree[1], 1); + off_chain_tree.add_leaf(leaves_in_small_tree[0], 2); + + // Mock the creation of the pre-append data structure + let subtree_proofs: Vec> = vec![vec![], vec![]]; + let subtree_rmls: Vec = vec![small_off_chain_tree.get_node(0), small_off_chain_tree.get_node(1)]; + let subtree_roots: Vec = vec![small_off_chain_tree.get_node(0), small_off_chain_tree.get_node(1)]; + + // append the small_merkle_roll to merkle_roll + merkle_roll + .append_subtree_packed( + &subtree_proofs, + &subtree_rmls, + &subtree_roots + ) + .unwrap(); + + // The index of the rmp to the on_chain_merkle roll should be equivalent to us having performed two appends, since the append should be dense + assert_eq!( + merkle_roll.rightmost_proof.index, + 3, + "On chain append was not tightly packed" + ); + + // The result should be that the merkle_roll's new root is the same as the root of the off-chain tree which had leaves 0..15 appended one by one + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "On chain tree failed to update properly on an append", + ); + + // Show that we can still append to the large tree after performing a subtree append + let leaf = rng.gen::<[u8; 32]>(); + merkle_roll.append(leaf).unwrap(); + off_chain_tree.add_leaf(leaf, 3); + assert_eq!( + merkle_roll.get_change_log().get_root(), + off_chain_tree.get_root(), + "Failed to append accurately to merkle roll after subtree append", + ); +} + #[tokio::test(threaded_scheduler)] async fn test_prove_leaf() { let (mut merkle_roll, mut off_chain_tree) = setup(); @@ -142,7 +455,7 @@ async fn test_initialize_with_root() { .unwrap(); assert_eq!( - merkle_roll.get_change_log().root, + merkle_roll.get_change_log().get_root(), tree.get_root(), "Init failed to set root properly" ); @@ -202,7 +515,7 @@ async fn test_replaces() { tree.add_leaf(leaf, i); merkle_roll.append(leaf).unwrap(); } - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); // Replace leaves in order for i in 0..(1 << DEPTH) { @@ -217,7 +530,7 @@ async fn test_replaces() { ) .unwrap(); tree.add_leaf(leaf, i); - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); } // Replaces leaves in a random order by 4x capacity @@ -235,7 +548,7 @@ async fn test_replaces() { ) .unwrap(); tree.add_leaf(leaf, index); - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); } } @@ -261,7 +574,7 @@ async fn test_mixed() { tree.add_leaf(leaf, i); merkle_roll.append(leaf).unwrap(); } - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); // Replaces leaves in a random order by 4x capacity let mut last_rmp = merkle_roll.rightmost_proof; @@ -289,15 +602,12 @@ async fn test_mixed() { .unwrap(); tree.add_leaf(leaf, index); } - if merkle_roll.get_change_log().root != tree.get_root() { + if merkle_roll.get_change_log().get_root() != tree.get_root() { let last_active_index: usize = (merkle_roll.active_index as usize + BUFFER_SIZE - 1) % BUFFER_SIZE; - println!("{:?}", &last_rmp); - println!("{:?}", &merkle_roll.change_logs[last_active_index]); - println!("{:?}", &merkle_roll.get_change_log()) } last_rmp = merkle_roll.rightmost_proof; - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); } } @@ -315,7 +625,7 @@ async fn test_append_bug_repro_1() { tree.add_leaf(leaf, i); merkle_roll.append(leaf).unwrap(); } - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); // Replace the rightmost leaf let leaf_0 = rng.gen::<[u8; 32]>(); @@ -340,13 +650,13 @@ async fn test_append_bug_repro_1() { tree_size += 1; // Now compare something - if merkle_roll.get_change_log().root != tree.get_root() { + if merkle_roll.get_change_log().get_root() != tree.get_root() { let last_active_index: usize = (merkle_roll.active_index as usize + BUFFER_SIZE - 1) % BUFFER_SIZE; println!("{:?}", &last_rmp); } last_rmp = merkle_roll.rightmost_proof; - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); } #[tokio::test(threaded_scheduler)] @@ -363,7 +673,7 @@ async fn test_append_bug_repro_2() { tree.add_leaf(leaf, i); merkle_roll.append(leaf).unwrap(); } - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); // Replace the rightmost leaf let mut leaf = rng.gen::<[u8; 32]>(); @@ -389,11 +699,11 @@ async fn test_append_bug_repro_2() { tree_size += 1; // Now compare something - if merkle_roll.get_change_log().root != tree.get_root() { + if merkle_roll.get_change_log().get_root() != tree.get_root() { let last_active_index: usize = (merkle_roll.active_index as usize + BUFFER_SIZE - 1) % BUFFER_SIZE; println!("{:?}", &last_rmp); } last_rmp = merkle_roll.rightmost_proof; - assert_eq!(merkle_roll.get_change_log().root, tree.get_root()); + assert_eq!(merkle_roll.get_change_log().get_root(), tree.get_root()); }