MMR trees API (#118)

The main design goals of this mmr implementation were

1. Avoid database callbacks. As it is implemented, calling side must just smartly pre-load MMR nodes from the database (about log2(tree length) for append, twice as much for deletion).

2. Reuse as much code/logic between rust and c++ clients.

3. Close to zero memory consumption.

librustzcash/Cargo.toml

@@ -28,6 +28,7 @@ libc = "0.2"
 pairing = { version = "0.15.0", path = "../pairing" }
 lazy_static = "1"
 rand_core = "0.5.1"
+zcash_mmr = { git = "https://github.com/nikvolf/zcash-mmr" }
 zcash_primitives = { version = "0.1.0", path = "../zcash_primitives" }
 zcash_proofs = { version = "0.1.0", path = "../zcash_proofs" }
 

librustzcash/include/librustzcash.h

@@ -307,6 +307,33 @@ extern "C" {
         unsigned char *j_ret,
         unsigned char *addr_ret
     );
+
+    uint32_t librustzcash_mmr_append(
+        uint32_t cbranch,
+        uint32_t t_len,
+        const uint32_t *ni_ptr,
+        const unsigned char *n_ptr,
+        size_t p_len,
+        const unsigned char *nn_ptr,
+        unsigned char *rt_ret,
+        unsigned char *buf_ret
+    );
+
+    uint32_t librustzcash_mmr_delete(
+        uint32_t cbranch,
+        uint32_t t_len,
+        const uint32_t *ni_ptr,
+        const unsigned char *n_ptr,
+        size_t p_len,
+        size_t e_len,
+        unsigned char *rt_ret
+    );
+
+    uint32_t librustzcash_mmr_hash_node(
+        uint32_t cbranch,
+        const unsigned char *n_ptr,
+        unsigned char *h_ret
+    );
 }
 
 #endif // LIBRUSTZCASH_INCLUDE_H_

librustzcash/src/rustzcash.rs

@@ -65,6 +65,8 @@ use zcash_proofs::{
     sprout,
 };
 
+use zcash_mmr::{Entry as MMREntry, NodeData as MMRNodeData, Tree as MMRTree};
+
 #[cfg(test)]
 mod tests;
 
@@ -1173,3 +1175,196 @@ pub extern "C" fn librustzcash_zip32_xfvk_address(
 
     true
 }
+
+fn construct_mmr_tree(
+    // Consensus branch id
+    cbranch: u32,
+    // Length of tree in array representation
+    t_len: u32,
+
+    // Indices of provided tree nodes, length of p_len+e_len
+    ni_ptr: *const u32,
+    // Provided tree nodes data, length of p_len+e_len
+    n_ptr: *const [c_uchar; zcash_mmr::MAX_ENTRY_SIZE],
+
+    // Peaks count
+    p_len: size_t,
+    // Extra nodes loaded (for deletion) count
+    e_len: size_t,
+) -> Result<MMRTree, &'static str> {
+    let (indices, nodes) = unsafe {
+        (
+            slice::from_raw_parts(ni_ptr, p_len + e_len),
+            slice::from_raw_parts(n_ptr, p_len + e_len),
+        )
+    };
+
+    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 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))
+}
+
+#[no_mangle]
+pub extern "system" fn librustzcash_mmr_append(
+    // Consensus branch id
+    cbranch: u32,
+    // Length of tree in array representation
+    t_len: u32,
+    // Indices of provided tree nodes, length of p_len
+    ni_ptr: *const u32,
+    // Provided tree nodes data, length of p_len
+    n_ptr: *const [c_uchar; zcash_mmr::MAX_ENTRY_SIZE],
+    // Peaks count
+    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
+    buf_ret: *mut [c_uchar; zcash_mmr::MAX_NODE_DATA_SIZE],
+) -> u32 {
+    let new_node_bytes: &[u8; zcash_mmr::MAX_NODE_DATA_SIZE] = unsafe {
+        match nn_ptr.as_ref() {
+            Some(r) => r,
+            None => {
+                return 0;
+            } // Null pointer passed, error
+        }
+    };
+
+    let mut tree = match construct_mmr_tree(cbranch, t_len, ni_ptr, n_ptr, p_len, 0) {
+        Ok(t) => t,
+        _ => {
+            return 0;
+        } // error
+    };
+
+    let node = match MMRNodeData::from_bytes(cbranch, &new_node_bytes[..]) {
+        Ok(node) => node,
+        _ => {
+            return 0;
+        } // error
+    };
+
+    let appended = match tree.append_leaf(node) {
+        Ok(appended) => appended,
+        _ => {
+            return 0;
+        }
+    };
+
+    let return_count = appended.len();
+
+    let root_node = tree
+        .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);
+
+        for (idx, next_buf) in slice::from_raw_parts_mut(buf_ret, return_count as usize)
+            .iter_mut()
+            .enumerate()
+        {
+            tree.resolve_link(appended[idx])
+                .expect("This was generated by the tree and thus resolvable; qed")
+                .data()
+                .write(&mut &mut next_buf[..])
+                .expect("Write using cursor with enough buffer size cannot fail; qed");
+        }
+    }
+
+    return_count as u32
+}
+
+#[no_mangle]
+pub extern "system" fn librustzcash_mmr_delete(
+    // Consensus branch id
+    cbranch: u32,
+    // Length of tree in array representation
+    t_len: u32,
+    // Indices of provided tree nodes, length of p_len+e_len
+    ni_ptr: *const u32,
+    // Provided tree nodes data, length of p_len+e_len
+    n_ptr: *const [c_uchar; zcash_mmr::MAX_ENTRY_SIZE],
+    // Peaks count
+    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,
+) -> u32 {
+    let mut tree = match construct_mmr_tree(cbranch, t_len, ni_ptr, n_ptr, p_len, e_len) {
+        Ok(t) => t,
+        _ => {
+            return 0;
+        } // error
+    };
+
+    let truncate_len = match tree.truncate_leaf() {
+        Ok(v) => v,
+        _ => {
+            return 0;
+        } // Error
+    };
+
+    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,
+        );
+    }
+
+    truncate_len
+}
+
+#[no_mangle]
+pub extern "system" fn librustzcash_mmr_hash_node(
+    cbranch: u32,
+    n_ptr: *const [u8; zcash_mmr::MAX_NODE_DATA_SIZE],
+    h_ret: *mut u8,
+) -> u32 {
+    let node_bytes: &[u8; zcash_mmr::MAX_NODE_DATA_SIZE] = unsafe {
+        match n_ptr.as_ref() {
+            Some(r) => r,
+            None => return 1,
+        }
+    };
+
+    let node = match MMRNodeData::from_bytes(cbranch, &node_bytes[..]) {
+        Ok(n) => n,
+        _ => return 1, // error
+    };
+
+    unsafe {
+        slice::from_raw_parts_mut(h_ret, 32).copy_from_slice(&node.hash()[..]);
+    }
+
+    return 0;
+}

librustzcash/src/tests/mmr.rs

@@ -0,0 +1,228 @@
+use zcash_mmr::{Entry, EntryLink, NodeData};
+
+use crate::{librustzcash_mmr_append, librustzcash_mmr_delete};
+
+const NODE_DATA_16L: &[u8] = include_bytes!("./res/tree16.dat");
+const NODE_DATA_1023L: &[u8] = include_bytes!("./res/tree1023.dat");
+
+struct TreeView {
+    peaks: Vec<(u32, Entry)>,
+    extra: Vec<(u32, Entry)>,
+}
+
+fn draft(into: &mut Vec<(u32, Entry)>, vec: &Vec<NodeData>, peak_pos: usize, h: u32) {
+    let node_data = vec[peak_pos - 1].clone();
+    let peak: Entry = match h {
+        0 => node_data.into(),
+        _ => Entry::new(
+            node_data,
+            EntryLink::Stored((peak_pos - (1 << h) - 1) as u32),
+            EntryLink::Stored((peak_pos - 2) as u32),
+        ),
+    };
+
+    into.push(((peak_pos - 1) as u32, peak));
+}
+
+fn prepare_tree(vec: &Vec<NodeData>) -> TreeView {
+    assert!(vec.len() > 0);
+
+    // integer log2 of (vec.len()+1), -1
+    let mut h = (32 - ((vec.len() + 1) as u32).leading_zeros() - 1) - 1;
+    let mut peak_pos = (1 << (h + 1)) - 1;
+    let mut nodes = Vec::new();
+
+    // used later
+    let mut last_peak_pos = 0;
+    let mut last_peak_h = 0;
+
+    loop {
+        if peak_pos > vec.len() {
+            // left child, -2^h
+            peak_pos = peak_pos - (1 << h);
+            h = h - 1;
+        }
+
+        if peak_pos <= vec.len() {
+            draft(&mut nodes, vec, 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;
+        }
+
+        if h == 0 {
+            break;
+        }
+    }
+
+    // for deletion, everything on the right slope of the last peak should be pre-loaded
+    let mut extra = Vec::new();
+    let mut h = last_peak_h;
+    let mut peak_pos = last_peak_pos;
+
+    while h > 0 {
+        let left_pos = peak_pos - (1 << h);
+        let right_pos = peak_pos - 1;
+        h = h - 1;
+
+        // drafting left child
+        draft(&mut extra, vec, left_pos, h);
+
+        // drafting right child
+        draft(&mut extra, vec, right_pos, h);
+
+        // continuing on right slope
+        peak_pos = right_pos;
+    }
+
+    TreeView {
+        peaks: nodes,
+        extra,
+    }
+}
+
+fn preload_tree_append(vec: &Vec<NodeData>) -> (Vec<u32>, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>) {
+    assert!(vec.len() > 0);
+
+    let tree_view = prepare_tree(vec);
+
+    let mut indices = Vec::new();
+    let mut bytes = Vec::new();
+
+    for (idx, entry) in tree_view.peaks.into_iter() {
+        let mut buf = [0u8; zcash_mmr::MAX_ENTRY_SIZE];
+        entry
+            .write(&mut &mut buf[..])
+            .expect("Cannot fail if enough buffer length");
+        indices.push(idx);
+        bytes.push(buf);
+    }
+
+    (indices, bytes)
+}
+
+// also returns number of peaks
+fn preload_tree_delete(
+    vec: &Vec<NodeData>,
+) -> (Vec<u32>, Vec<[u8; zcash_mmr::MAX_ENTRY_SIZE]>, usize) {
+    assert!(vec.len() > 0);
+
+    let tree_view = prepare_tree(vec);
+
+    let mut indices = Vec::new();
+    let mut bytes = Vec::new();
+
+    let peak_count = tree_view.peaks.len();
+
+    for (idx, entry) in tree_view
+        .peaks
+        .into_iter()
+        .chain(tree_view.extra.into_iter())
+    {
+        let mut buf = [0u8; zcash_mmr::MAX_ENTRY_SIZE];
+        entry
+            .write(&mut &mut buf[..])
+            .expect("Cannot fail if enough buffer length");
+        indices.push(idx);
+        bytes.push(buf);
+    }
+
+    (indices, bytes, peak_count)
+}
+
+fn load_nodes(bytes: &'static [u8]) -> Vec<NodeData> {
+    let mut res = Vec::new();
+    let mut cursor = std::io::Cursor::new(bytes);
+    while (cursor.position() as usize) < bytes.len() {
+        let node_data =
+            zcash_mmr::NodeData::read(0, &mut cursor).expect("Statically checked to be correct");
+        res.push(node_data);
+    }
+
+    res
+}
+
+#[test]
+fn append() {
+    let nodes = load_nodes(NODE_DATA_16L);
+    let (indices, peaks) = preload_tree_append(&nodes);
+
+    let mut rt_ret = [0u8; 32];
+
+    let mut buf_ret = Vec::<[u8; zcash_mmr::MAX_NODE_DATA_SIZE]>::with_capacity(32);
+
+    let mut new_node_data = [0u8; zcash_mmr::MAX_NODE_DATA_SIZE];
+    let new_node = NodeData {
+        consensus_branch_id: 0,
+        subtree_commitment: [0u8; 32],
+        start_time: 101,
+        end_time: 110,
+        start_target: 190,
+        end_target: 200,
+        start_sapling_root: [0u8; 32],
+        end_sapling_root: [0u8; 32],
+        subtree_total_work: Default::default(),
+        start_height: 10,
+        end_height: 10,
+        shielded_tx: 13,
+    };
+    new_node
+        .write(&mut &mut new_node_data[..])
+        .expect("Failed to write node data");
+
+    let result = librustzcash_mmr_append(
+        0,
+        nodes.len() as u32,
+        indices.as_ptr(),
+        peaks.as_ptr(),
+        peaks.len(),
+        &new_node_data,
+        rt_ret.as_mut_ptr(),
+        buf_ret.as_mut_ptr(),
+    );
+
+    unsafe {
+        buf_ret.set_len(result as usize);
+    }
+
+    assert_eq!(result, 2);
+
+    let new_node_1 =
+        NodeData::from_bytes(0, &buf_ret[0][..]).expect("Failed to reconstruct return node #1");
+
+    let new_node_2 =
+        NodeData::from_bytes(0, &buf_ret[1][..]).expect("Failed to reconstruct return node #2");
+
+    assert_eq!(new_node_1.start_height, 10);
+    assert_eq!(new_node_1.end_height, 10);
+
+    // this is combined new node (which is `new_node_1`) + the one which was there before (for block #9)
+    assert_eq!(new_node_2.start_height, 9);
+    assert_eq!(new_node_2.end_height, 10);
+    assert_eq!(new_node_2.shielded_tx, 27);
+}
+
+#[test]
+fn delete() {
+    let nodes = load_nodes(NODE_DATA_1023L);
+    let (indices, nodes, peak_count) = preload_tree_delete(&nodes);
+
+    let mut rt_ret = [0u8; 32];
+
+    let result = librustzcash_mmr_delete(
+        0,
+        nodes.len() as u32,
+        indices.as_ptr(),
+        nodes.as_ptr(),
+        peak_count,
+        indices.len() - peak_count,
+        rt_ret.as_mut_ptr(),
+    );
+
+    // Deleting from full tree of 9 height would result in cascade deleting of 10 nodes
+    assert_eq!(result, 10);
+}

librustzcash/src/tests/mod.rs

@@ -4,6 +4,7 @@ use super::JUBJUB;
 
 mod key_agreement;
 mod key_components;
+mod mmr;
 mod notes;
 mod signatures;