diff --git a/src/builder.rs b/src/builder.rs index b8cb144..3e16ee3 100644 --- a/src/builder.rs +++ b/src/builder.rs @@ -1,14 +1,14 @@ -use crate::commitment::CTree; +use crate::commitment::{CTree, Witness}; 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); +trait Builder { + fn collect(&mut self, commitments: &[Node], context: &C) -> usize; fn up(&mut self); fn finished(&self) -> bool; - fn finalize(self) -> CTree; + fn finalize(self) -> T; } struct CTreeBuilder { @@ -16,11 +16,13 @@ struct CTreeBuilder { right: Option, prev_tree: CTree, next_tree: CTree, + start: usize, depth: usize, + offset: Option, } -impl Builder for CTreeBuilder { - fn collect(&mut self, commitments: &[Node]) -> (Option, usize) { +impl Builder for CTreeBuilder { + fn collect(&mut self, commitments: &[Node], _context: &()) -> 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 @@ -57,7 +59,8 @@ impl Builder for CTreeBuilder { }; assert_eq!(n % 2, 0); - (offset, n) + self.offset = offset; + n } fn up(&mut self) { @@ -74,6 +77,10 @@ impl Builder for CTreeBuilder { self.left = l; self.right = r; + + assert!(self.start % 2 == 0 || self.offset.is_some()); + self.start /= 2; + self.depth += 1; } @@ -88,12 +95,15 @@ impl Builder for CTreeBuilder { impl CTreeBuilder { fn new(prev_tree: CTree) -> CTreeBuilder { + let start = prev_tree.get_position(); CTreeBuilder { left: prev_tree.left, right: prev_tree.right, prev_tree, next_tree: CTree::new(), + start, depth: 0, + offset: None, } } @@ -130,35 +140,112 @@ fn combine_level(commitments: &mut [Node], offset: Option, n: usize, depth nn } -#[allow(dead_code)] -fn advance_tree(prev_tree: CTree, mut commitments: &mut [Node]) -> CTree { - if commitments.is_empty() { - return prev_tree; +struct WitnessBuilder { + witness: Witness, + p: usize, + inside: bool, +} + +impl WitnessBuilder { + fn new(tree_builder: &CTreeBuilder, prev_witness: Witness, count: usize) -> WitnessBuilder { + let position = prev_witness.position; + let inside = position >= tree_builder.start && position < tree_builder.start + count; + WitnessBuilder { + witness: prev_witness, + p: position, + inside, + } } - 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(); +} + +impl Builder for WitnessBuilder { + fn collect(&mut self, commitments: &[Node], context: &CTreeBuilder) -> usize { + let offset = context.offset; + let depth = context.depth; + + if self.inside { + let rp = self.p - context.start; + if depth == 0 { + if self.p % 2 == 1 { + self.witness.tree.left = Some(CTreeBuilder::get(commitments, rp - 1, offset)); + self.witness.tree.right = Some(CTreeBuilder::get(commitments, rp, offset)); + } + else { + self.witness.tree.left = Some(CTreeBuilder::get(commitments, rp, offset)); + self.witness.tree.right = None; + } + } + else { + if self.p % 2 == 1 { + self.witness.tree.parents.push(Some(CTreeBuilder::get(commitments, rp, offset))); + } + else { + self.witness.tree.parents.push(None); + } + } + } + + // TODO: update filled + + 0 } - builder.finalize() + fn up(&mut self) { + self.p /= 2; + } + + fn finished(&self) -> bool { + false + } + + fn finalize(self) -> Witness { + self.witness + } +} + +#[allow(dead_code)] +fn advance_tree(prev_tree: CTree, prev_witnesses: &[Witness], mut commitments: &mut [Node]) -> (CTree, Vec) { + if commitments.is_empty() { + return (prev_tree, prev_witnesses.to_vec()); + } + let mut builder = CTreeBuilder::new(prev_tree); + let mut witness_builders: Vec<_> = prev_witnesses.iter().map(|witness| { + WitnessBuilder::new(&builder, witness.clone(), commitments.len()) + }).collect(); + while !commitments.is_empty() || !builder.finished() { + let n = builder.collect(commitments, &()); + for b in witness_builders.iter_mut() { + b.collect(commitments, &builder); + } + let nn = combine_level(commitments, builder.offset, n, builder.depth); + commitments = &mut commitments[0..nn]; + builder.up(); + for b in witness_builders.iter_mut() { + b.up(); + } + } + + let tree = builder.finalize(); + let witnesses = witness_builders.into_iter().map(|b| b.finalize()).collect(); + (tree, witnesses) } #[cfg(test)] mod tests { - use crate::builder::advance_tree; - use crate::commitment::CTree; + use crate::builder::{advance_tree, WitnessBuilder}; + use crate::commitment::{CTree, Witness}; use zcash_primitives::sapling::Node; - use zcash_primitives::merkle_tree::CommitmentTree; + use zcash_primitives::merkle_tree::{CommitmentTree, IncrementalWitness}; + use crate::print::{print_witness, print_tree}; #[test] fn test_advance_tree() { - const NUM_NODES: usize = 100; - const NUM_CHUNKS: usize = 100; + const NUM_NODES: usize = 10; + const NUM_CHUNKS: usize = 10; let mut tree1: CommitmentTree = CommitmentTree::empty(); let mut tree2 = CTree::new(); + let mut ws: Vec> = vec![]; + let mut ws2: Vec = vec![]; for i in 0..NUM_CHUNKS { let mut nodes: Vec<_> = (0..NUM_NODES).map(|k| { let mut bb = [0u8; 32]; @@ -166,10 +253,20 @@ mod tests { bb[0..8].copy_from_slice(&v.to_be_bytes()); let node = Node::new(bb); tree1.append(node).unwrap(); + for w in ws.iter_mut() { + w.append(node).unwrap(); + } + if v == 55 { + let w = IncrementalWitness::from_tree(&tree1); + ws.push(w); + ws2.push(Witness::new(v)); + } node }).collect(); - tree2 = advance_tree(tree2, &mut nodes); + let (new_tree, new_witnesses) = advance_tree(tree2, &ws2, &mut nodes); + tree2 = new_tree; + ws2 = new_witnesses; } let mut bb1: Vec = vec![]; tree1.write(&mut bb1).unwrap(); @@ -179,20 +276,25 @@ mod tests { 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))); - } + print_tree(&ws[0].tree); println!("-----"); + let tree2 = &ws2[0].tree; + + // 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"); + // + // assert!(equal, "not equal"); } } \ No newline at end of file diff --git a/src/commitment.rs b/src/commitment.rs index 05f3027..9f6816a 100644 --- a/src/commitment.rs +++ b/src/commitment.rs @@ -77,14 +77,16 @@ gets pushed into `filled`. */ #[derive(Clone)] pub struct Witness { - tree: CTree, // commitment tree at the moment the witness is created: immutable + pub(crate) position: usize, + pub(crate) tree: CTree, // commitment tree at the moment the witness is created: immutable filled: Vec, // as more nodes are added, levels get filled up: won't change anymore cursor: CTree, // partial tree which still updates when nodes are added } impl Witness { - pub fn new() -> Witness { + pub fn new(position: usize) -> Witness { Witness { + position, tree: CTree::new(), filled: vec![], cursor: CTree::new(), @@ -151,7 +153,7 @@ impl NotePosition { p: position, p2: count - 1, c, - witness: Witness::new(), + witness: Witness::new(position), } } @@ -305,7 +307,7 @@ impl CTree { }) } - fn get_position(&self) -> usize { + pub(crate) fn get_position(&self) -> usize { let mut p = 0usize; for parent in self.parents.iter().rev() { if parent.is_some() {