Merge pull request #180 from str4d/mmr-cleanups

MMR API cleanups

librustzcash/src/rustzcash.rs

@@ -1199,31 +1199,17 @@ fn construct_mmr_tree(
         )
     };
 
-    let mut peaks = Vec::new();
+    let mut peaks: Vec<_> = indices
-    for i in 0..p_len {
+        .iter()
-        peaks.push((
+        .zip(nodes.iter())
-            indices[i],
+        .map(
-            match MMREntry::from_bytes(cbranch, &nodes[i][..]) {
+            |(index, node)| match MMREntry::from_bytes(cbranch, &node[..]) {
-                Ok(entry) => entry,
+                Ok(entry) => Ok((*index, entry)),
-                _ => {
+                Err(_) => Err("Invalid encoding"),
-                    return Err("Invalid encoding");
-                } // error
             },
-        ));
+        )
-    }
+        .collect::<Result<_, _>>()?;
-
+    let extra = peaks.split_off(p_len);
-    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
-            },
-        ));
-    }
 
     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)
+    // Return of root commitment
-    rt_ret: *mut u8,
+    rt_ret: *mut [u8; 32],
-    // Return buffer for appended leaves, should be pre-allocated of log2(t_len)+1 length
+    // 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)
+    // Return of root commitment
-    rt_ret: *mut u8,
+    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(
+        *rt_ret = tree
-            &tree
+            .root_node()
-                .root_node()
+            .expect("Just generated without errors, root should be resolving")
-                .expect("Just generated without errors, root should be resolving")
+            .data()
-                .data()
+            .subtree_commitment;
-                .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
 }

librustzcash/src/tests/mmr.rs

@@ -10,8 +10,8 @@ struct TreeView {
     extra: Vec<(u32, Entry)>,
 }
 
-fn draft(into: &mut Vec<(u32, Entry)>, vec: &Vec<NodeData>, peak_pos: usize, h: u32) {
+fn draft(into: &mut Vec<(u32, Entry)>, nodes: &[NodeData], peak_pos: usize, h: u32) {
-    let node_data = vec[peak_pos - 1].clone();
+    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<NodeData>, peak_pos: usize, h:
     into.push(((peak_pos - 1) as u32, peak));
 }
 
-fn prepare_tree(vec: &Vec<NodeData>) -> TreeView {
+fn prepare_tree(nodes: &[NodeData]) -> TreeView {
-    assert!(vec.len() > 0);
+    assert!(!nodes.is_empty());
 
-    // integer log2 of (vec.len()+1), -1
+    // integer log2 of (nodes.len()+1), -1
-    let mut h = (32 - ((vec.len() + 1) as u32).leading_zeros() - 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);
+            peak_pos -= 1 << h;
-            h = h - 1;
+            h -= 1;
         }
 
-        if peak_pos <= vec.len() {
+        if peak_pos <= nodes.len() {
-            draft(&mut nodes, vec, peak_pos, h);
+            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<NodeData>) -> 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 {
+    TreeView { peaks, extra }
-        peaks: nodes,
-        extra,
-    }
 }
 
-fn preload_tree_append(vec: &Vec<NodeData>) -> (Vec<u32>, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>) {
+fn preload_tree_append(nodes: &[NodeData]) -> (Vec<u32>, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>) {
-    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();
@@ -107,11 +104,11 @@ fn preload_tree_append(vec: &Vec<NodeData>) -> (Vec<u32>, Vec<[u8; zcash_mmr::MA
 
 // also returns number of peaks
 fn preload_tree_delete(
-    vec: &Vec<NodeData>,
+    nodes: &[NodeData],
 ) -> (Vec<u32>, 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