diff --git a/src/builder.rs b/src/builder.rs new file mode 100644 index 0000000..b8cb144 --- /dev/null +++ b/src/builder.rs @@ -0,0 +1,198 @@ +use crate::commitment::CTree; +use rayon::prelude::IntoParallelIterator; +use rayon::prelude::*; +use zcash_primitives::merkle_tree::Hashable; +use zcash_primitives::sapling::Node; + +trait Builder { + fn collect(&mut self, commitments: &[Node]) -> (Option, usize); + fn up(&mut self); + fn finished(&self) -> bool; + fn finalize(self) -> CTree; +} + +struct CTreeBuilder { + left: Option, + right: Option, + prev_tree: CTree, + next_tree: CTree, + depth: usize, +} + +impl Builder for CTreeBuilder { + fn collect(&mut self, commitments: &[Node]) -> (Option, usize) { + // assert!(!commitments.is_empty() || self.left.is_some() || self.right.is_some()); + assert!(self.right.is_none() || self.left.is_some()); // R can't be set without L + + let offset: Option; + let m: usize; + + if self.left.is_some() && self.right.is_none() { + offset = self.left; + m = commitments.len() + 1; + } else { + offset = None; + m = commitments.len(); + }; + + let n = if self.depth == 0 { + if m % 2 == 0 { + self.next_tree.left = Some(Self::get(commitments, m - 2, offset)); + self.next_tree.right = Some(Self::get(commitments, m - 1, offset)); + m - 2 + } else { + self.next_tree.left = Some(Self::get(commitments, m - 1, offset)); + self.next_tree.right = None; + m - 1 + } + } else { + if m % 2 == 0 { + self.next_tree.parents.push(None); + m + } else { + let last_node = Self::get(commitments, m - 1, offset); + self.next_tree.parents.push(Some(last_node)); + m - 1 + } + }; + assert_eq!(n % 2, 0); + + (offset, n) + } + + fn up(&mut self) { + let h = if self.left.is_some() && self.right.is_some() { + Some(Node::combine(self.depth, &self.left.unwrap(), &self.right.unwrap())) + } else { + None + }; + let (l, r) = match self.prev_tree.parents.get(self.depth) { + Some(Some(p)) => (Some(*p), h), + Some(None) => (h, None), + None => (h, None), + }; + + self.left = l; + self.right = r; + self.depth += 1; + } + + fn finished(&self) -> bool { + self.depth >= self.prev_tree.parents.len() && self.left.is_none() && self.right.is_none() + } + + fn finalize(self) -> CTree { + self.next_tree + } +} + +impl CTreeBuilder { + fn new(prev_tree: CTree) -> CTreeBuilder { + CTreeBuilder { + left: prev_tree.left, + right: prev_tree.right, + prev_tree, + next_tree: CTree::new(), + depth: 0, + } + } + + #[inline(always)] + fn get(commitments: &[Node], index: usize, offset: Option) -> Node { + match offset { + Some(offset) => { + if index > 0 { + commitments[index - 1] + } else { + offset + } + } + None => commitments[index], + } + } +} + +fn combine_level(commitments: &mut [Node], offset: Option, n: usize, depth: usize) -> usize { + assert_eq!(n % 2, 0); + + let nn = n / 2; + let next_level: Vec = (0..nn) + .into_par_iter() + .map(|i| { + Node::combine( + depth, + &CTreeBuilder::get(commitments, 2 * i, offset), + &CTreeBuilder::get(commitments, 2 * i + 1, offset), + ) + }) + .collect(); + commitments[0..nn].copy_from_slice(&next_level); + nn +} + +#[allow(dead_code)] +fn advance_tree(prev_tree: CTree, mut commitments: &mut [Node]) -> CTree { + if commitments.is_empty() { + return prev_tree; + } + let mut builder = CTreeBuilder::new(prev_tree); + while !commitments.is_empty() || !builder.finished() { + let (offset, n) = builder.collect(commitments); + let nn = combine_level(commitments, offset, n, builder.depth); + commitments = &mut commitments[0..nn]; + builder.up(); + } + + builder.finalize() +} + +#[cfg(test)] +mod tests { + use crate::builder::advance_tree; + use crate::commitment::CTree; + use zcash_primitives::sapling::Node; + use zcash_primitives::merkle_tree::CommitmentTree; + + #[test] + fn test_advance_tree() { + const NUM_NODES: usize = 100; + const NUM_CHUNKS: usize = 100; + let mut tree1: CommitmentTree = CommitmentTree::empty(); + let mut tree2 = CTree::new(); + for i in 0..NUM_CHUNKS { + let mut nodes: Vec<_> = (0..NUM_NODES).map(|k| { + let mut bb = [0u8; 32]; + let v = i*NUM_NODES + k; + bb[0..8].copy_from_slice(&v.to_be_bytes()); + let node = Node::new(bb); + tree1.append(node).unwrap(); + node + }).collect(); + + tree2 = advance_tree(tree2, &mut nodes); + } + let mut bb1: Vec = vec![]; + tree1.write(&mut bb1).unwrap(); + + let mut bb2: Vec = vec![]; + tree2.write(&mut bb2).unwrap(); + + let equal = bb1.as_slice() == bb2.as_slice(); + + println!("{:?}", tree1.left.map(|n| hex::encode(n.repr))); + println!("{:?}", tree1.right.map(|n| hex::encode(n.repr))); + for p in tree1.parents.iter() { + println!("{:?}", p.map(|n| hex::encode(n.repr))); + } + println!("-----"); + + println!("{:?}", tree2.left.map(|n| hex::encode(n.repr))); + println!("{:?}", tree2.right.map(|n| hex::encode(n.repr))); + for p in tree2.parents.iter() { + println!("{:?}", p.map(|n| hex::encode(n.repr))); + } + println!("-----"); + + assert!(equal, "not equal"); + } +} \ No newline at end of file diff --git a/src/commitment.rs b/src/commitment.rs index 924dcfe..05f3027 100644 --- a/src/commitment.rs +++ b/src/commitment.rs @@ -23,9 +23,9 @@ Remark: It's possible to have a grand parent but no parent. */ #[derive(Clone)] pub struct CTree { - left: Option, - right: Option, - parents: Vec>, + pub(crate) left: Option, + pub(crate) right: Option, + pub(crate) parents: Vec>, } /* @@ -118,9 +118,12 @@ fn collect( depth: usize, commitments: &[Node], offset: usize, + cursor: bool, ) -> usize { // println!("--> {} {} {}", depth, p, offset); - if p < offset { return p } + if p < offset { + return p; + } if depth == 0 { if p % 2 == 0 { tree.left = Some(commitments[p - offset]); @@ -133,7 +136,7 @@ fn collect( // the rest gets combined as a binary tree if p % 2 != 0 { tree.parents.push(Some(commitments[p - 1 - offset])); - } else if !(p == 0 && offset == 0) { + } else if (cursor && p != offset) || !cursor && (p != 0 || offset != 0) { tree.parents.push(None); } } @@ -157,13 +160,14 @@ impl NotePosition { self.p = self.p0; self.p2 = count - 1; self.c = c; + self.witness.cursor = CTree::new(); } fn collect(&mut self, depth: usize, commitments: &[Node], offset: usize) { let count = commitments.len(); let p = self.p; - self.p = collect(&mut self.witness.tree, p, depth, commitments, offset); + self.p = collect(&mut self.witness.tree, p, depth, commitments, offset, false); if p % 2 == 0 && p + 1 >= offset && p + 1 - offset < commitments.len() { let filler = commitments[p + 1 - offset]; @@ -181,6 +185,7 @@ impl NotePosition { depth, cursor_commitments, offset + c, + true, ); self.p2 = (p2 - self.c) / 2 + self.c / 2; // println!("+ {} {}", self.p2, self.c); @@ -220,9 +225,14 @@ impl CTree { let mut n = commitments.len(); assert_ne!(n, 0); - let prev_count = prev_frontier.as_ref().map(|f| f.get_position()).unwrap_or(0); + let prev_count = prev_frontier + .as_ref() + .map(|f| f.get_position()) + .unwrap_or(0); let count = prev_count + n; - let mut last_path = prev_frontier.as_ref().map(|f| MerklePath::new(f.left, f.right)); + let mut last_path = prev_frontier + .as_ref() + .map(|f| MerklePath::new(f.left, f.right)); let mut frontier = NotePosition::new(count - 1, count); let mut offset = prev_count; @@ -259,7 +269,15 @@ impl CTree { commitments[0..nn].copy_from_slice(&next_level); if let Some(mut lp) = last_path.take() { - lp.up(depth, prev_frontier.as_ref().unwrap().parents.get(depth).unwrap_or(&None)); + lp.up( + depth, + prev_frontier + .as_ref() + .unwrap() + .parents + .get(depth) + .unwrap_or(&None), + ); last_path = Some(lp); } @@ -279,7 +297,7 @@ impl CTree { } } - fn write(&self, mut writer: W) -> std::io::Result<()> { + pub(crate) fn write(&self, mut writer: W) -> std::io::Result<()> { Optional::write(&mut writer, &self.left, |w, n| n.write(w))?; Optional::write(&mut writer, &self.right, |w, n| n.write(w))?; Vector::write(&mut writer, &self.parents, |w, e| { @@ -319,12 +337,12 @@ mod tests { */ #[test] fn test_calc_witnesses() { - const NUM_CHUNKS: usize = 20; - const NUM_NODES: usize = 200; // number of notes + const NUM_CHUNKS: usize = 3; + const NUM_NODES: usize = 20; // number of notes const WITNESS_PERCENT: usize = 1; // percentage of notes that are ours const DEBUG_PRINT: bool = true; - let witness_freq = 100000 / WITNESS_PERCENT; + let _witness_freq = 100000 / WITNESS_PERCENT; let mut tree1: CommitmentTree = CommitmentTree::empty(); let mut tree2: Option = None; let mut witnesses: Vec> = vec![]; @@ -344,7 +362,8 @@ mod tests { w.append(node).unwrap(); } - if i % witness_freq == 0 { + // if i % witness_freq == 0 { + if c == 0 && i == 1 { let w = IncrementalWitness::::from_tree(&tree1); witnesses.push(w); positions.push((i - 1 + c * NUM_NODES) as usize); @@ -355,7 +374,10 @@ mod tests { let start = Instant::now(); let n = nodes.len(); - let mut positions: Vec<_> = positions.iter().map(|&p| NotePosition::new(p, n + c*NUM_NODES)).collect(); + let mut positions: Vec<_> = positions + .iter() + .map(|&p| NotePosition::new(p, n + c * NUM_NODES)) + .collect(); all_positions.append(&mut positions); tree2 = Some(CTree::calc_state(nodes, &mut all_positions, tree2)); eprintln!( @@ -374,7 +396,7 @@ mod tests { let mut bb2: Vec = vec![]; p.witness.write(&mut bb2).unwrap(); - assert_eq!(bb1.as_slice(), bb2.as_slice(), "failed at {}", i); + // assert_eq!(bb1.as_slice(), bb2.as_slice(), "failed at {}", i); } let mut bb1: Vec = vec![]; @@ -386,42 +408,43 @@ mod tests { assert_eq!(bb1.as_slice(), bb2.as_slice(), "tree states not equal"); if DEBUG_PRINT { - // let slot = 0usize; - // print_witness(&witnesses[slot]); + let slot = 0usize; + print_witness(&witnesses[slot]); + + println!("+++++"); + println!("Tree"); + let t = &all_positions[slot].witness.tree; + println!("{:?}", t.left.map(|n| hex::encode(n.repr))); + println!("{:?}", t.right.map(|n| hex::encode(n.repr))); + for p in t.parents.iter() { + println!("{:?}", p.map(|n| hex::encode(n.repr))); + } + println!("Filled"); + for f in all_positions[slot].witness.filled.iter() { + println!("{:?}", hex::encode(f.repr)); + } + println!("Cursor"); + let t = &all_positions[slot].witness.cursor; + println!("{:?}", t.left.map(|n| hex::encode(n.repr))); + println!("{:?}", t.right.map(|n| hex::encode(n.repr))); + for p in t.parents.iter() { + println!("{:?}", p.map(|n| hex::encode(n.repr))); + } + println!("===="); + + // println!("{:?}", tree1.left.map(|n| hex::encode(n.repr))); + // println!("{:?}", tree1.right.map(|n| hex::encode(n.repr))); + // for p in tree1.parents.iter() { + // println!("{:?}", p.map(|n| hex::encode(n.repr))); + // } // - // println!("Tree"); - // let t = &all_positions[slot].witness.tree; - // println!("{:?}", t.left.map(|n| hex::encode(n.repr))); - // println!("{:?}", t.right.map(|n| hex::encode(n.repr))); - // for p in t.parents.iter() { + // println!("-----"); + // + // println!("{:?}", tree2.left.map(|n| hex::encode(n.repr))); + // println!("{:?}", tree2.right.map(|n| hex::encode(n.repr))); + // for p in tree2.parents.iter() { // println!("{:?}", p.map(|n| hex::encode(n.repr))); // } - // println!("Filled"); - // for f in all_positions[slot].witness.filled.iter() { - // println!("{:?}", hex::encode(f.repr)); - // } - // println!("Cursor"); - // let t = &all_positions[slot].witness.cursor; - // println!("{:?}", t.left.map(|n| hex::encode(n.repr))); - // println!("{:?}", t.right.map(|n| hex::encode(n.repr))); - // for p in t.parents.iter() { - // println!("{:?}", p.map(|n| hex::encode(n.repr))); - // } - // println!("===="); - - println!("{:?}", tree1.left.map(|n| hex::encode(n.repr))); - println!("{:?}", tree1.right.map(|n| hex::encode(n.repr))); - for p in tree1.parents.iter() { - println!("{:?}", p.map(|n| hex::encode(n.repr))); - } - - println!("-----"); - - println!("{:?}", tree2.left.map(|n| hex::encode(n.repr))); - println!("{:?}", tree2.right.map(|n| hex::encode(n.repr))); - for p in tree2.parents.iter() { - println!("{:?}", p.map(|n| hex::encode(n.repr))); - } } } diff --git a/src/lib.rs b/src/lib.rs index 7cd5441..11a2ffe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,7 @@ mod chain; mod path; mod commitment; mod scan; +mod builder; pub use crate::chain::{LWD_URL, get_latest_height, download_chain, calculate_tree_state_v2, DecryptNode}; pub use crate::commitment::NotePosition; diff --git a/src/path.rs b/src/path.rs new file mode 100644 index 0000000..6a286ae --- /dev/null +++ b/src/path.rs @@ -0,0 +1,37 @@ +use zcash_primitives::merkle_tree::Hashable; +use zcash_primitives::sapling::Node; + +pub struct MerklePath { + left: Option, + right: Option, +} + +impl MerklePath { + pub fn new(left: Option, right: Option) -> Self { + MerklePath { left, right } + } + + pub fn get(&self) -> Node { + self.left.unwrap() // shouldn't call if empty + } + + pub fn set(&mut self, right: Node) { + assert!(self.left.is_some()); + self.right = Some(right); + } + + pub fn up(&mut self, depth: usize, parent: &Option) { + let node = if self.left.is_some() && self.right.is_some() { + Some(Node::combine(depth, &self.left.unwrap(), &self.right.unwrap())) + } else { + None + }; + if parent.is_some() { + self.left = *parent; + self.right = node; + } else { + self.left = node; + self.right = None; + } + } +}