Browse Source
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.master
Nikolay Volf
5 years ago
committed by
str4d
8 changed files with 480 additions and 0 deletions
@ -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); |
||||||
|
} |
Binary file not shown.
Binary file not shown.
Loading…
Reference in new issue