witness tree

This commit is contained in:
Hanh 2021-06-21 11:36:53 +08:00
parent dcfc8af0e1
commit 9283e9940b
2 changed files with 139 additions and 35 deletions

View File

@ -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<Node>, usize);
trait Builder<T, C> {
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<Node>,
prev_tree: CTree,
next_tree: CTree,
start: usize,
depth: usize,
offset: Option<Node>,
}
impl Builder for CTreeBuilder {
fn collect(&mut self, commitments: &[Node]) -> (Option<Node>, usize) {
impl Builder<CTree, ()> 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<Node>, 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<Witness, CTreeBuilder> 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<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| {
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<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 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<u8> = 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");
}
}

View File

@ -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<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
}
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() {