witness filled

src/builder.rs

@@ -8,7 +8,7 @@ trait Builder<T, C> {
     fn collect(&mut self, commitments: &[Node], context: &C) -> usize;
     fn up(&mut self);
     fn finished(&self) -> bool;
-    fn finalize(self) -> T;
+    fn finalize(self, context: &C) -> T;
 }
 
 struct CTreeBuilder {
@@ -39,11 +39,11 @@ impl Builder<CTree, ()> for CTreeBuilder {
 
         let n = if self.depth == 0 {
             if m % 2 == 0 {
-                self.next_tree.left = Some(Self::get(commitments, m - 2, offset));
+                self.next_tree.left = Some(*Self::get(commitments, m - 2, &offset));
-                self.next_tree.right = Some(Self::get(commitments, m - 1, 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.left = Some(*Self::get(commitments, m - 1, &offset));
                 self.next_tree.right = None;
                 m - 1
             }
@@ -52,8 +52,8 @@ impl Builder<CTree, ()> for CTreeBuilder {
                 self.next_tree.parents.push(None);
                 m
             } else {
-                let last_node = Self::get(commitments, m - 1, offset);
+                let last_node = Self::get(commitments, m - 1, &offset);
-                self.next_tree.parents.push(Some(last_node));
+                self.next_tree.parents.push(Some(*last_node));
                 m - 1
             }
         };
@@ -65,7 +65,11 @@ impl Builder<CTree, ()> for CTreeBuilder {
 
     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()))
+            Some(Node::combine(
+                self.depth,
+                &self.left.unwrap(),
+                &self.right.unwrap(),
+            ))
         } else {
             None
         };
@@ -88,7 +92,7 @@ impl Builder<CTree, ()> for CTreeBuilder {
         self.depth >= self.prev_tree.parents.len() && self.left.is_none() && self.right.is_none()
     }
 
-    fn finalize(self) -> CTree {
+    fn finalize(self, _context: &()) -> CTree {
         self.next_tree
     }
 }
@@ -108,18 +112,50 @@ impl CTreeBuilder {
     }
 
     #[inline(always)]
-    fn get(commitments: &[Node], index: usize, offset: Option<Node>) -> Node {
+    fn get_opt<'a>(
-        match offset {
+        commitments: &'a [Node],
-            Some(offset) => {
+        index: usize,
-                if index > 0 {
+        offset: &'a Option<Node>,
-                    commitments[index - 1]
+    ) -> Option<&'a Node> {
-                } else {
+        if offset.is_some() {
-                    offset
+            if index > 0 {
-                }
+                commitments.get(index - 1)
+            } else {
+                offset.as_ref()
             }
-            None => commitments[index],
+        } else {
+            commitments.get(index)
         }
     }
+
+    #[inline(always)]
+    fn get<'a>(commitments: &'a [Node], index: usize, offset: &'a Option<Node>) -> &'a Node {
+        Self::get_opt(commitments, index, offset).unwrap()
+    }
+
+    fn adjusted_start(&self) -> usize {
+        if self.offset.is_some() {
+            self.start - 1
+        } else {
+            self.start
+        }
+    }
+
+    fn adjusted_end(&self, commitments: &[Node]) -> usize {
+        self.start + commitments.len()
+    }
+
+    fn clone_trimmed(&self, mut depth: usize) -> CTree {
+        if depth == 0 {
+            return CTree::new()
+        }
+        let mut tree = self.next_tree.clone();
+        while depth > 0 && depth <= tree.parents.len() && tree.parents[depth - 1].is_none() {
+            depth -= 1;
+        }
+        tree.parents.truncate(depth);
+        tree
+    }
 }
 
 fn combine_level(commitments: &mut [Node], offset: Option<Node>, n: usize, depth: usize) -> usize {
@@ -131,8 +167,8 @@ fn combine_level(commitments: &mut [Node], offset: Option<Node>, n: usize, depth
         .map(|i| {
             Node::combine(
                 depth,
-                &CTreeBuilder::get(commitments, 2 * i, offset),
+                CTreeBuilder::get(commitments, 2 * i, &offset),
-                &CTreeBuilder::get(commitments, 2 * i + 1, offset),
+                CTreeBuilder::get(commitments, 2 * i + 1, &offset),
             )
         })
         .collect();
@@ -163,30 +199,52 @@ impl Builder<Witness, CTreeBuilder> for WitnessBuilder {
         let offset = context.offset;
         let depth = context.depth;
 
+        if depth == 1 {
+            context.offset.map(|c| println!("{}", hex::encode(c.repr)));
+            for c in commitments.iter() {
+                println!("{}", hex::encode(c.repr));
+            }
+            println!("===");
+        }
+
+        let tree = &mut self.witness.tree;
+
         if self.inside {
-            let rp = self.p - context.start;
+            let rp = self.p - context.adjusted_start();
             if depth == 0 {
                 if self.p % 2 == 1 {
-                    self.witness.tree.left = Some(CTreeBuilder::get(commitments, rp - 1, offset));
+                    tree.left = Some(*CTreeBuilder::get(commitments, rp - 1, &offset));
-                    self.witness.tree.right = Some(CTreeBuilder::get(commitments, rp, offset));
+                    tree.right = Some(*CTreeBuilder::get(commitments, rp, &offset));
+                } else {
+                    tree.left = Some(*CTreeBuilder::get(commitments, rp, &offset));
+                    tree.right = None;
                 }
-                else {
+            } 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)));
+                    tree.parents
-                }
+                        .push(Some(*CTreeBuilder::get(commitments, rp - 1, &offset)));
-                else {
+                } else if self.p != 0 {
-                    self.witness.tree.parents.push(None);
+                    tree.parents.push(None);
                 }
             }
         }
 
-        // TODO: update filled
+        let p1 = self.p + 1;
-
+        let has_p1 = p1 >= context.adjusted_start() && p1 < context.adjusted_end(commitments);
+        // println!("D{}: {} {} {} {}", depth, p1, context.start, context.start + commitments.len(), context.offset.is_some());
+        if has_p1 {
+            let p1 = CTreeBuilder::get(commitments, p1 - context.adjusted_start(), &offset);
+            println!("P1 {} {}", depth, hex::encode(p1.repr));
+            if depth == 0 {
+                if tree.right.is_none() {
+                    self.witness.filled.push(*p1);
+                }
+            } else {
+                if depth - 1 >= tree.parents.len() || tree.parents[depth - 1].is_none() {
+                    self.witness.filled.push(*p1);
+                }
+            }
+        }
         0
     }
 
@@ -198,20 +256,48 @@ impl Builder<Witness, CTreeBuilder> for WitnessBuilder {
         false
     }
 
-    fn finalize(self) -> Witness {
+    fn finalize(mut self, context: &CTreeBuilder) -> Witness {
+        let tree = &self.witness.tree;
+        let mut num_filled = self.witness.filled.len();
+
+        let mut depth = 0;
+        loop {
+            let is_none = if depth == 0 { // check if this level is occupied
+                tree.right.is_none()
+            } else {
+                depth > tree.parents.len() || tree.parents[depth - 1].is_none()
+            };
+            if is_none {
+                if num_filled > 0 {
+                    num_filled -= 1; // we filled it
+                }
+                else {
+                    break
+                }
+            }
+            depth += 1;
+            // loop terminates because we are eventually going to run out of ancestors and filled
+        }
+
+        self.witness.cursor = context.clone_trimmed(depth - 1);
         self.witness
     }
 }
 
 #[allow(dead_code)]
-fn advance_tree(prev_tree: CTree, prev_witnesses: &[Witness], mut commitments: &mut [Node]) -> (CTree, Vec<Witness>) {
+fn advance_tree(
+    prev_tree: CTree,
+    prev_witnesses: &[Witness],
+    mut commitments: &mut [Node],
+) -> (CTree, Vec<Witness>) {
     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| {
+    let mut witness_builders: Vec<_> = prev_witnesses
-        WitnessBuilder::new(&builder, witness.clone(), commitments.len())
+        .iter()
-    }).collect();
+        .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() {
@@ -225,49 +311,61 @@ fn advance_tree(prev_tree: CTree, prev_witnesses: &[Witness], mut commitments: &
         }
     }
 
-    let tree = builder.finalize();
+    let witnesses = witness_builders
-    let witnesses = witness_builders.into_iter().map(|b| b.finalize()).collect();
+        .into_iter()
+        .map(|b| b.finalize(&builder))
+        .collect();
+    let tree = builder.finalize(&());
     (tree, witnesses)
 }
 
 #[cfg(test)]
+#[allow(unused_imports)]
 mod tests {
-    use crate::builder::{advance_tree, WitnessBuilder};
+    use crate::builder::advance_tree;
     use crate::commitment::{CTree, Witness};
-    use zcash_primitives::sapling::Node;
+    use crate::print::{print_tree, print_witness};
     use zcash_primitives::merkle_tree::{CommitmentTree, IncrementalWitness};
-    use crate::print::{print_witness, print_tree};
+    use zcash_primitives::sapling::Node;
 
     #[test]
     fn test_advance_tree() {
         const NUM_NODES: usize = 10;
-        const NUM_CHUNKS: usize = 10;
+        const NUM_CHUNKS: usize = 5;
+        const WITNESS_PERCENT: f64 = 100.0; // percentage of notes that are ours
+        const DEBUG_PRINT: bool = true;
+        let witness_freq = (100.0 / WITNESS_PERCENT) as usize;
+
         let mut tree1: CommitmentTree<Node> = CommitmentTree::empty();
         let mut tree2 = CTree::new();
         let mut ws: Vec<IncrementalWitness<Node>> = vec![];
         let mut ws2: Vec<Witness> = vec![];
         for i in 0..NUM_CHUNKS {
-            let mut nodes: Vec<_> = (0..NUM_NODES).map(|k| {
+            let mut nodes: Vec<_> = vec![];
+            for j in 0..NUM_NODES {
                 let mut bb = [0u8; 32];
-                let v = i*NUM_NODES + k;
+                let v = i * NUM_NODES + j;
                 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 {
+                // if v % witness_freq == 0 {
+                if v == 37 {
                     let w = IncrementalWitness::from_tree(&tree1);
                     ws.push(w);
                     ws2.push(Witness::new(v));
                 }
-                node
+                nodes.push(node);
-            }).collect();
+            }
 
             let (new_tree, new_witnesses) = advance_tree(tree2, &ws2, &mut nodes);
             tree2 = new_tree;
             ws2 = new_witnesses;
         }
+
+        // check final state
         let mut bb1: Vec<u8> = vec![];
         tree1.write(&mut bb1).unwrap();
 
@@ -276,25 +374,64 @@ mod tests {
 
         let equal = bb1.as_slice() == bb2.as_slice();
 
-        print_tree(&ws[0].tree);
+        println!("# witnesses = {}", ws.len());
-        println!("-----");
 
-        let tree2 = &ws2[0].tree;
+        // check witnesses
+        let mut failed_index: Option<usize> = None;
+        for (i, (w1, w2)) in ws.iter().zip(&ws2).enumerate() {
+            let mut bb1: Vec<u8> = vec![];
+            w1.write(&mut bb1).unwrap();
 
-        // println!("{:?}", tree1.left.map(|n| hex::encode(n.repr)));
+            let mut bb2: Vec<u8> = vec![];
-        // println!("{:?}", tree1.right.map(|n| hex::encode(n.repr)));
+            w2.write(&mut bb2).unwrap();
-        // for p in tree1.parents.iter() {
+
-        //     println!("{:?}", p.map(|n| hex::encode(n.repr)));
+            if bb1.as_slice() != bb2.as_slice() {
-        // }
+                failed_index = Some(i);
-        // println!("-----");
+            }
-        //
+        }
-        println!("{:?}", tree2.left.map(|n| hex::encode(n.repr)));
+
-        println!("{:?}", tree2.right.map(|n| hex::encode(n.repr)));
+        if DEBUG_PRINT && (failed_index.is_some() || !equal) {
-        for p in tree2.parents.iter() {
+            let i = failed_index.unwrap();
-            println!("{:?}", p.map(|n| hex::encode(n.repr)));
+            println!("FAILED AT {}", i);
+            print_witness(&ws[i]);
+
+            // println!("-----");
+            // println!("Final-----");
+            //
+            // 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!("-----");
+
+            // 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!("----- {}", ws2[i].position);
+            let tree2 = &ws2[i].tree;
+            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!("-----");
+            let filled2 = &ws2[i].filled;
+            println!("Filled");
+            for f in filled2.iter() {
+                println!("{:?}", hex::encode(f.repr));
+            }
+            println!("Cursor");
+            let cursor2 = &ws2[i].cursor;
+            println!("{:?}", cursor2.left.map(|n| hex::encode(n.repr)));
+            println!("{:?}", cursor2.right.map(|n| hex::encode(n.repr)));
+            for p in cursor2.parents.iter() {
+                println!("{:?}", p.map(|n| hex::encode(n.repr)));
+            }
+
+            assert!(false);
         }
-        println!("-----");
-        //
-        // assert!(equal, "not equal");
     }
-}
+}

src/commitment.rs

@@ -77,10 +77,10 @@ gets pushed into `filled`.
 */
 #[derive(Clone)]
 pub struct Witness {
-    pub(crate) position: usize,
+    pub position: usize,
-    pub(crate) tree: CTree,       // commitment tree at the moment the witness is created: immutable
+    pub tree: CTree,       // commitment tree at the moment the witness is created: immutable
-    filled: Vec<Node>, // as more nodes are added, levels get filled up: won't change anymore
+    pub filled: Vec<Node>, // as more nodes are added, levels get filled up: won't change anymore
-    cursor: CTree,     // partial tree which still updates when nodes are added
+    pub cursor: CTree,     // partial tree which still updates when nodes are added
 }
 
 impl Witness {
@@ -299,6 +299,10 @@ impl CTree {
         }
     }
 
+    pub fn is_empty(&self) -> bool {
+        self.left.is_none() && self.right.is_none()
+    }
+
     pub(crate) fn write<W: 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))?;
@@ -391,7 +395,7 @@ mod tests {
 
         println!("# witnesses = {}", all_positions.len());
 
-        for (i, (w, p)) in witnesses.iter().zip(&all_positions).enumerate() {
+        for (_i, (w, p)) in witnesses.iter().zip(&all_positions).enumerate() {
             let mut bb1: Vec<u8> = vec![];
             w.write(&mut bb1).unwrap();