tree::testing: Fix and test arb_tree().

This commit is contained in:
therealyingtong 2021-06-10 13:54:42 +08:00
parent 2d0afe9357
commit b3daeb0861
3 changed files with 95 additions and 59 deletions

View File

@ -522,7 +522,7 @@ pub mod testing {
/// Produce a random valid Orchard bundle.
fn arb_bundle_inputs(sk: SpendingKey)
(
n_notes in 1..30,
n_notes in 1usize..30,
n_recipients in 1..30,
)
(

View File

@ -47,11 +47,6 @@ impl NoteCommitment {
pub struct ExtractedNoteCommitment(pub(super) pallas::Base);
impl ExtractedNoteCommitment {
/// ExtractedNoteCommitment for an uncommitted note.
pub fn uncommitted() -> Self {
Self(pallas::Base::zero())
}
/// Deserialize the extracted note commitment from a byte array.
pub fn from_bytes(bytes: &[u8; 32]) -> CtOption<Self> {
pallas::Base::from_bytes(bytes).map(ExtractedNoteCommitment)

View File

@ -40,6 +40,14 @@ impl MerklePath {
}
}
/// <https://zips.z.cash/protocol/protocol.pdf#orchardmerklecrh>
/// The layer with 2^n nodes is called "layer n":
/// - leaves are at layer MERKLE_DEPTH_ORCHARD = 32;
/// - the root is at layer 0.
/// `l_star` is MERKLE_DEPTH_ORCHARD - layer - 1.
/// - when hashing two leaves, we produce a node on the layer above the leaves, i.e.
/// layer = 31, l_star = 0
/// - when hashing to the final root, we produce the anchor with layer = 0, l_star = 31.
pub fn root(&self, cmx: ExtractedNoteCommitment) -> Anchor {
let node = self
.auth_path
@ -73,6 +81,13 @@ fn cond_swap(swap: bool, node: pallas::Base, sibling: pallas::Base) -> Pair {
}
// <https://zips.z.cash/protocol/protocol.pdf#orchardmerklecrh>
// The layer with 2^n nodes is called "layer n":
// - leaves are at layer MERKLE_DEPTH_ORCHARD = 32;
// - the root is at layer 0.
// `l_star` is MERKLE_DEPTH_ORCHARD - layer - 1.
// - when hashing two leaves, we produce a node on the layer above the leaves, i.e.
// layer = 31, l_star = 0
// - when hashing to the final root, we produce the anchor with layer = 0, l_star = 31.
fn hash_layer(l_star: usize, pair: Pair) -> pallas::Base {
// MerkleCRH Sinsemilla hash domain.
let domain = HashDomain::new(MERKLE_CRH_PERSONALIZATION);
@ -104,14 +119,13 @@ fn hash_layer(l_star: usize, pair: Pair) -> pallas::Base {
pub mod testing {
use lazy_static::lazy_static;
use std::convert::TryInto;
use std::iter;
use crate::{
constants::{MERKLE_CRH_PERSONALIZATION, MERKLE_DEPTH_ORCHARD},
constants::MERKLE_DEPTH_ORCHARD,
note::{commitment::ExtractedNoteCommitment, testing::arb_note, Note},
primitives::sinsemilla::{i2lebsp_k, HashDomain, K},
value::{testing::arb_positive_note_value, MAX_NOTE_VALUE},
};
use ff::{PrimeField, PrimeFieldBits};
use pasta_curves::pallas;
use proptest::collection::vec;
@ -121,95 +135,113 @@ pub mod testing {
lazy_static! {
static ref EMPTY_ROOTS: Vec<pallas::Base> = {
// MerkleCRH Sinsemilla hash domain.
let domain = HashDomain::new(MERKLE_CRH_PERSONALIZATION);
let mut v = vec![pallas::Base::zero()];
for l_star in 0..MERKLE_DEPTH_ORCHARD {
let next = domain
.hash(
std::iter::empty()
.chain(i2lebsp_k(l_star).iter().copied().take(K))
.chain(
v[l_star]
.to_le_bits()
.iter()
.by_val()
.take(pallas::Base::NUM_BITS as usize)
)
.chain(
v[l_star]
.to_le_bits()
.iter()
.by_val()
.take(pallas::Base::NUM_BITS as usize)
),
iter::empty()
.chain(Some(pallas::Base::zero()))
.chain(
(0..MERKLE_DEPTH_ORCHARD).scan(pallas::Base::zero(), |state, l_star| {
*state = hash_layer(
l_star,
Pair {
left: *state,
right: *state,
},
);
Some(*state)
}),
)
.unwrap();
v.push(next);
}
v
.collect()
};
}
prop_compose! {
/// Generates an arbitrary Merkle tree of with `n_notes` nonempty leaves.
pub fn arb_tree(n_notes: i32)
pub fn arb_tree(n_notes: usize)
(
// generate note values that we're certain won't exceed MAX_NOTE_VALUE in total
notes in vec(
arb_positive_note_value(MAX_NOTE_VALUE / n_notes as u64).prop_flat_map(arb_note),
n_notes as usize
n_notes
),
)
-> (Vec<(Note, MerklePath)>, Anchor) {
// Inefficient algorithm to build a perfect subtree containing all notes.
let perfect_subtree_depth = (n_notes as f64).log2().ceil() as usize;
let commitments: Vec<ExtractedNoteCommitment> = notes.iter().map(|note| {
let n_leaves = 1 << perfect_subtree_depth;
let commitments: Vec<Option<ExtractedNoteCommitment>> = notes.iter().map(|note| {
let cmx: ExtractedNoteCommitment = note.commitment().into();
cmx
Some(cmx)
}).collect();
let padded_leaves = {
let mut padded_leaves = commitments.clone();
let pad = (0..((1 << perfect_subtree_depth) - n_notes as usize)).map(
|_| ExtractedNoteCommitment::uncommitted()
let pad = (0..(n_leaves - n_notes)).map(
|_| None
).collect::<Vec<_>>();
padded_leaves.extend_from_slice(&pad);
padded_leaves
};
let tree = {
let mut tree: Vec<Vec<pallas::Base>> = vec![padded_leaves.into_iter().map(|leaf| *leaf).collect()];
let perfect_subtree = {
let mut perfect_subtree: Vec<Vec<Option<pallas::Base>>> = vec![
padded_leaves.iter().map(|cmx| cmx.map(|cmx| *cmx)).collect()
];
// <https://zips.z.cash/protocol/protocol.pdf#orchardmerklecrh>
// The layer with 2^n nodes is called "layer n":
// - leaves are at layer MERKLE_DEPTH_ORCHARD = 32;
// - the root is at layer 0.
// `l_star` is MERKLE_DEPTH_ORCHARD - layer - 1.
// - when hashing two leaves, we produce a node on the layer above the leaves, i.e.
// layer = 31, l_star = 0
// - when hashing to the final root, we produce the anchor with layer = 0, l_star = 31.
for height in 1..perfect_subtree_depth {
let inner_nodes = (0..(perfect_subtree_depth >> height)).map(|pos| {
hash_layer(height, Pair {
left: tree[height - 1][pos * 2],
right: tree[height - 1][pos * 2 + 1],
})
let l_star = height - 1;
let inner_nodes = (0..(n_leaves >> height)).map(|pos| {
let left = perfect_subtree[height - 1][pos * 2];
let right = perfect_subtree[height - 1][pos * 2 + 1];
match (left, right) {
(None, None) => None,
(Some(left), None) => {
let right = EMPTY_ROOTS[height - 1];
Some(hash_layer(l_star, Pair {left, right}))
},
(Some(left), Some(right)) => {
Some(hash_layer(l_star, Pair {left, right}))
},
(None, Some(_)) => {
unreachable!("The perfect subtree is left-packed.")
}
}
}).collect();
tree.push(inner_nodes);
perfect_subtree.push(inner_nodes);
};
tree
perfect_subtree
};
// Get Merkle path for each note commitment
let auth_paths = {
let mut auth_paths: Vec<MerklePath> = Vec::new();
for (pos, _) in commitments.iter().enumerate() {
// Initialize the authentication path to the path for an empty tree.
let mut auth_path: [pallas::Base; MERKLE_DEPTH_ORCHARD] = (0..MERKLE_DEPTH_ORCHARD).map(|idx| EMPTY_ROOTS[idx]).collect::<Vec<_>>().try_into().unwrap();
let mut layer_pos = pos;
for height in 0..perfect_subtree_depth {
let is_right_sibling = layer_pos & 1 == 1;
let sibling = if is_right_sibling {
tree[height][layer_pos - 1]
// This node is the right sibling, so we need its left sibling at the current height.
perfect_subtree[height][layer_pos - 1]
} else {
tree[height][layer_pos + 1]
// This node is the left sibling, so we need its right sibling at the current height.
perfect_subtree[height][layer_pos + 1]
};
auth_path[height] = sibling;
if let Some(sibling) = sibling {
auth_path[height] = sibling;
}
layer_pos = (layer_pos - is_right_sibling as usize) / 2;
};
@ -220,11 +252,7 @@ pub mod testing {
};
// Compute anchor for this tree
let anchor = auth_paths[0].root(commitments[0]);
for (cmx, auth_path) in commitments.iter().zip(auth_paths.iter()) {
let computed_anchor = auth_path.root(*cmx);
assert_eq!(anchor, computed_anchor);
}
let anchor = auth_paths[0].root(notes[0].commitment().into());
(
notes.into_iter().zip(auth_paths.into_iter()).map(|(note, auth_path)| (note, auth_path)).collect(),
@ -232,4 +260,17 @@ pub mod testing {
)
}
}
proptest! {
#[allow(clippy::redundant_closure)]
#[test]
fn tree(
(notes_and_auth_paths, anchor) in (1usize..4).prop_flat_map(|n_notes| arb_tree(n_notes))
) {
for (note, auth_path) in notes_and_auth_paths.iter() {
let computed_anchor = auth_path.root(note.commitment().into());
assert_eq!(anchor, computed_anchor);
}
}
}
}