diff --git a/librustzcash/src/rustzcash.rs b/librustzcash/src/rustzcash.rs index 0bfd5ad..6dadf67 100644 --- a/librustzcash/src/rustzcash.rs +++ b/librustzcash/src/rustzcash.rs @@ -1199,31 +1199,17 @@ fn construct_mmr_tree( ) }; - let mut peaks = Vec::new(); - for i in 0..p_len { - peaks.push(( - indices[i], - match MMREntry::from_bytes(cbranch, &nodes[i][..]) { - Ok(entry) => entry, - _ => { - return Err("Invalid encoding"); - } // error + let mut peaks: Vec<_> = indices + .iter() + .zip(nodes.iter()) + .map( + |(index, node)| match MMREntry::from_bytes(cbranch, &node[..]) { + Ok(entry) => Ok((*index, entry)), + Err(_) => Err("Invalid encoding"), }, - )); - } - - let mut extra = Vec::new(); - for i in p_len..(p_len + e_len) { - extra.push(( - indices[i], - match MMREntry::from_bytes(cbranch, &nodes[i][..]) { - Ok(entry) => entry, - _ => { - return Err("Invalid encoding"); - } // error - }, - )); - } + ) + .collect::>()?; + let extra = peaks.split_off(p_len); Ok(MMRTree::new(t_len, peaks, extra)) } @@ -1242,9 +1228,9 @@ pub extern "system" fn librustzcash_mmr_append( p_len: size_t, // New node pointer nn_ptr: *const [u8; zcash_mmr::MAX_NODE_DATA_SIZE], - // Return of root commitment (32 byte hash) - rt_ret: *mut u8, - // Return buffer for appended leaves, should be pre-allocated of log2(t_len)+1 length + // Return of root commitment + rt_ret: *mut [u8; 32], + // Return buffer for appended leaves, should be pre-allocated of ceiling(log2(t_len)) length buf_ret: *mut [c_uchar; zcash_mmr::MAX_NODE_DATA_SIZE], ) -> u32 { let new_node_bytes: &[u8; zcash_mmr::MAX_NODE_DATA_SIZE] = unsafe { @@ -1283,7 +1269,7 @@ pub extern "system" fn librustzcash_mmr_append( .root_node() .expect("Just added, should resolve always; qed"); unsafe { - slice::from_raw_parts_mut(rt_ret, 32).copy_from_slice(&root_node.data().subtree_commitment); + *rt_ret = root_node.data().subtree_commitment; for (idx, next_buf) in slice::from_raw_parts_mut(buf_ret, return_count as usize) .iter_mut() @@ -1314,8 +1300,8 @@ pub extern "system" fn librustzcash_mmr_delete( p_len: size_t, // Extra nodes loaded (for deletion) count e_len: size_t, - // Return of root commitment (32 byte hash) - rt_ret: *mut u8, + // Return of root commitment + rt_ret: *mut [u8; 32], ) -> u32 { let mut tree = match construct_mmr_tree(cbranch, t_len, ni_ptr, n_ptr, p_len, e_len) { Ok(t) => t, @@ -1332,13 +1318,11 @@ pub extern "system" fn librustzcash_mmr_delete( }; unsafe { - slice::from_raw_parts_mut(rt_ret, 32).copy_from_slice( - &tree - .root_node() - .expect("Just generated without errors, root should be resolving") - .data() - .subtree_commitment, - ); + *rt_ret = tree + .root_node() + .expect("Just generated without errors, root should be resolving") + .data() + .subtree_commitment; } truncate_len @@ -1348,7 +1332,7 @@ pub extern "system" fn librustzcash_mmr_delete( pub extern "system" fn librustzcash_mmr_hash_node( cbranch: u32, n_ptr: *const [u8; zcash_mmr::MAX_NODE_DATA_SIZE], - h_ret: *mut u8, + h_ret: *mut [u8; 32], ) -> u32 { let node_bytes: &[u8; zcash_mmr::MAX_NODE_DATA_SIZE] = unsafe { match n_ptr.as_ref() { @@ -1363,8 +1347,8 @@ pub extern "system" fn librustzcash_mmr_hash_node( }; unsafe { - slice::from_raw_parts_mut(h_ret, 32).copy_from_slice(&node.hash()[..]); + *h_ret = node.hash(); } - return 0; + 0 } diff --git a/librustzcash/src/tests/mmr.rs b/librustzcash/src/tests/mmr.rs index 6e40f7f..4107a75 100644 --- a/librustzcash/src/tests/mmr.rs +++ b/librustzcash/src/tests/mmr.rs @@ -10,8 +10,8 @@ struct TreeView { extra: Vec<(u32, Entry)>, } -fn draft(into: &mut Vec<(u32, Entry)>, vec: &Vec, peak_pos: usize, h: u32) { - let node_data = vec[peak_pos - 1].clone(); +fn draft(into: &mut Vec<(u32, Entry)>, nodes: &[NodeData], peak_pos: usize, h: u32) { + let node_data = nodes[peak_pos - 1].clone(); let peak: Entry = match h { 0 => node_data.into(), _ => Entry::new( @@ -24,34 +24,34 @@ fn draft(into: &mut Vec<(u32, Entry)>, vec: &Vec, peak_pos: usize, h: into.push(((peak_pos - 1) as u32, peak)); } -fn prepare_tree(vec: &Vec) -> TreeView { - assert!(vec.len() > 0); +fn prepare_tree(nodes: &[NodeData]) -> TreeView { + assert!(!nodes.is_empty()); - // integer log2 of (vec.len()+1), -1 - let mut h = (32 - ((vec.len() + 1) as u32).leading_zeros() - 1) - 1; + // integer log2 of (nodes.len()+1), -1 + let mut h = (32 - ((nodes.len() + 1) as u32).leading_zeros() - 1) - 1; let mut peak_pos = (1 << (h + 1)) - 1; - let mut nodes = Vec::new(); + let mut peaks = Vec::new(); // used later let mut last_peak_pos = 0; let mut last_peak_h = 0; loop { - if peak_pos > vec.len() { + if peak_pos > nodes.len() { // left child, -2^h - peak_pos = peak_pos - (1 << h); - h = h - 1; + peak_pos -= 1 << h; + h -= 1; } - if peak_pos <= vec.len() { - draft(&mut nodes, vec, peak_pos, h); + if peak_pos <= nodes.len() { + draft(&mut peaks, nodes, peak_pos, h); // save to be used in next loop last_peak_pos = peak_pos; last_peak_h = h; // right sibling - peak_pos = peak_pos + (1 << (h + 1)) - 1; + peak_pos += (1 << (h + 1)) - 1; } if h == 0 { @@ -67,28 +67,25 @@ fn prepare_tree(vec: &Vec) -> TreeView { while h > 0 { let left_pos = peak_pos - (1 << h); let right_pos = peak_pos - 1; - h = h - 1; + h -= 1; // drafting left child - draft(&mut extra, vec, left_pos, h); + draft(&mut extra, nodes, left_pos, h); // drafting right child - draft(&mut extra, vec, right_pos, h); + draft(&mut extra, nodes, right_pos, h); // continuing on right slope peak_pos = right_pos; } - TreeView { - peaks: nodes, - extra, - } + TreeView { peaks, extra } } -fn preload_tree_append(vec: &Vec) -> (Vec, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>) { - assert!(vec.len() > 0); +fn preload_tree_append(nodes: &[NodeData]) -> (Vec, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>) { + assert!(!nodes.is_empty()); - let tree_view = prepare_tree(vec); + let tree_view = prepare_tree(nodes); let mut indices = Vec::new(); let mut bytes = Vec::new(); @@ -107,11 +104,11 @@ fn preload_tree_append(vec: &Vec) -> (Vec, Vec<[u8; zcash_mmr::MA // also returns number of peaks fn preload_tree_delete( - vec: &Vec, + nodes: &[NodeData], ) -> (Vec, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>, usize) { - assert!(vec.len() > 0); + assert!(!nodes.is_empty()); - let tree_view = prepare_tree(vec); + let tree_view = prepare_tree(nodes); let mut indices = Vec::new(); let mut bytes = Vec::new(); @@ -181,7 +178,7 @@ fn append() { peaks.as_ptr(), peaks.len(), &new_node_data, - rt_ret.as_mut_ptr(), + &mut rt_ret, buf_ret.as_mut_ptr(), ); @@ -220,7 +217,7 @@ fn delete() { nodes.as_ptr(), peak_count, indices.len() - peak_count, - rt_ret.as_mut_ptr(), + &mut rt_ret, ); // Deleting from full tree of 9 height would result in cascade deleting of 10 nodes