diff --git a/Cargo.toml b/Cargo.toml index 3f68b5e..3369570 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,8 +1,8 @@ [package] name = "poseidon252" -version = "0.4.0" +version = "0.5.0" authors = [ - "zer0 ", "vlopes11 ", "CPerezz " + "zer0 ", "vlopes11 ", "CPerezz ", "Kristoffer Ström " ] edition = "2018" diff --git a/src/lib.rs b/src/lib.rs index 6ec683a..765dda2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -70,7 +70,7 @@ //! ### Zero Knowledge Merkle Opening Proof example: //! ```no_run //! use std::borrow::Borrow; -//! use poseidon252::{StorageScalar, PoseidonAnnotation}; +//! use poseidon252::{StorageScalar, PoseidonTree}; //! use poseidon252::merkle_proof::merkle_opening_gadget; //! use dusk_plonk::commitment_scheme::kzg10::PublicParameters; //! use dusk_plonk::constraint_system::{Variable, StandardComposer}; @@ -83,9 +83,9 @@ //! let pub_params = PublicParameters::setup(1 << 17, &mut rand::thread_rng()).unwrap(); //! let (ck, vk) = pub_params.trim(1 << 16).unwrap(); //! // Generate a tree with random scalars inside. -//! let mut nstack: NStack<_, PoseidonAnnotation, Blake2b> = NStack::new(); +//! let mut ptree: PoseidonTree<_, Blake2b> = PoseidonTree::new(17); //! for i in 0..1024u64 { -//! nstack.push(StorageScalar(Scalar::from(i as u64))) +//! ptree.push(StorageScalar(Scalar::from(i as u64))) //! .unwrap(); //! } //! @@ -98,24 +98,17 @@ //! // In this case, the key X corresponds to the Scalar(X). //! // We're supposing that we're provided with a Kelvin::Branch to perform //! // the proof. -//! let branch = nstack.get(*i).unwrap().unwrap(); +//! let branch = ptree.poseidon_branch(*i).unwrap().unwrap(); //! //! // Get tree root. -//! let root = StorageScalar::from(branch -//! .levels() -//! .first() -//! .unwrap() -//! .annotation() -//! .unwrap() -//! .to_owned() -//! .borrow()); +//! let root = ptree.root().unwrap(); //! //! // Add the proven leaf value to the Constraint System //! let proven_leaf = composer.add_input(Scalar::from(*i)); //! //! // Print inside of the Composer Constraint System the Merkle Proof -//! // with all of the needed checks. Using branch length of 17 -//! merkle_opening_gadget(&mut composer, branch, proven_leaf, root.0.into(), 17); +//! // with all of the needed checks. +//! merkle_opening_gadget(&mut composer, branch, proven_leaf, root); //! //! // Since we don't use all of the wires, we set some dummy constraints to avoid Committing //! // to zero polynomials. @@ -137,16 +130,15 @@ //! ### Standard Merkle Opening Proof example: //! ```no_run //! use std::borrow::Borrow; -//! use poseidon252::{StorageScalar, PoseidonAnnotation}; +//! use poseidon252::{StorageScalar, PoseidonTree}; //! use poseidon252::merkle_proof::merkle_opening_scalar_verification; //! use dusk_bls12_381::Scalar; //! use kelvin::{Blake2b, Compound}; -//! use nstack::NStack; //! //! // Generate a tree with random scalars inside. -//! let mut nstack: NStack<_, PoseidonAnnotation, Blake2b> = NStack::new(); +//! let mut ptree: PoseidonTree<_, Blake2b> = PoseidonTree::new(17); //! for i in 0..1024u64 { -//! nstack.push(StorageScalar(Scalar::from(i as u64))) +//! ptree.push(StorageScalar(Scalar::from(i as u64))) //! .unwrap(); //! } //! @@ -157,24 +149,16 @@ //! // In this case, the key X corresponds to the Scalar(X). //! // We're supposing that we're provided with a Kelvin::Branch to perform //! // the proof. -//! let branch = nstack.get(i).unwrap().unwrap(); +//! let branch = ptree.poseidon_branch(i).unwrap().unwrap(); //! //! // Get tree root. -//! let root = StorageScalar::from(branch -//! .levels() -//! .first() -//! .unwrap() -//! .annotation() -//! .unwrap() -//! .to_owned() -//! .borrow()); +//! let root = ptree.root().unwrap(); //! //! // Verify the `Branch`. Use a branch length of 17. //! assert!(merkle_opening_scalar_verification( //! branch, -//! root.0.into(), +//! root, //! Scalar::from(i), -//! 17, //! )); //! } //! ``` @@ -207,6 +191,10 @@ pub mod merkle_proof; /// Reference implementation for the Poseidon Sponge hash function pub mod sponge; +/// The module handling posedion-trees +pub mod tree; +pub use tree::PoseidonTree; + /// Maximum arity supported for trees. /// /// This is due to the fact that actually we rely in Hades252 crate diff --git a/src/merkle_proof/poseidon_branch.rs b/src/merkle_proof/poseidon_branch.rs index 17e0caa..3d1146b 100644 --- a/src/merkle_proof/poseidon_branch.rs +++ b/src/merkle_proof/poseidon_branch.rs @@ -143,6 +143,21 @@ impl PoseidonBranch { } } +/// Applies the extension padding n times to a scalar +pub(crate) fn extend_scalar(mut scalar: Scalar, n: usize) -> Scalar { + for _ in 0..n { + let flag = Scalar::from(0b1000); + let mut leaves = [Scalar::zero(); ARITY + 1]; + + leaves[0] = flag; + leaves[1] = scalar; + + let level = PoseidonLevel { leaves, offset: 1 }; + scalar = hash::merkle_level_hash_without_bitflags(&level); + } + scalar +} + #[derive(Debug, Clone, PartialEq)] /// Represents a Merkle-Tree Level inside of a `PoseidonBranch`. /// It stores the leaves as `Scalar` and the offset which represents diff --git a/src/merkle_proof/proof.rs b/src/merkle_proof/proof.rs index 4993c61..d2161fb 100644 --- a/src/merkle_proof/proof.rs +++ b/src/merkle_proof/proof.rs @@ -1,13 +1,11 @@ //! Merkle-tree hashing functions using Poseidon252 //! -use super::poseidon_branch::{PoseidonBranch, PoseidonLevel}; +use super::poseidon_branch::PoseidonBranch; use crate::merkle_lvl_hash::hash::*; -use crate::ARITY; -use crate::{PoseidonAnnotation, StorageScalar}; + use dusk_bls12_381::Scalar; use dusk_plonk::constraint_system::{StandardComposer, Variable}; use hades252::WIDTH; -use nstack::NStack; /// Provided a `kelvin::Branch`, a `&mut StandardComposer`, a leaf value and a root, print inside of the /// constraint system a Merkle Tree Proof that hashes up from the searched leaf in kelvin until @@ -17,21 +15,12 @@ use nstack::NStack; /// /// NOTE: The root of the `Branch` (root of the Merkle tree) will be set as Public Input so we /// can re-use the circuits that rely on this gadget. -pub fn merkle_opening_gadget( +pub fn merkle_opening_gadget( composer: &mut StandardComposer, - branch: kelvin::Branch, H>, + branch: PoseidonBranch, proven_leaf: Variable, proven_root: Scalar, - branch_length: usize, -) where - H: kelvin::ByteHash, -{ - // Generate a `PoseidonBranch` from the kelvin Branch. - let mut branch = PoseidonBranch::from(&branch); - - let n_extensions = branch.extend(branch_length); - let proven_root = extend_scalar(proven_root, n_extensions); - +) { // Allocate space for each level Variables that will be generated. let mut lvl_vars = [composer.zero_var; WIDTH]; // Allocate space for the last level computed hash as a variable to compare @@ -102,21 +91,6 @@ pub fn merkle_opening_gadget( assert_eq!(branch.root, proven_root); } -/// Applies the extension padding n times to the scalar -fn extend_scalar(mut scalar: Scalar, n: usize) -> Scalar { - for _ in 0..n { - let flag = Scalar::from(0b1000); - let mut leaves = [Scalar::zero(); ARITY + 1]; - - leaves[0] = flag; - leaves[1] = scalar; - - let level = PoseidonLevel { leaves, offset: 1 }; - scalar = merkle_level_hash_without_bitflags(&level); - } - scalar -} - /// Provided a `PoseidonBranch` and a Merkle Tree root, verify that /// the path to the root is correct. /// @@ -124,22 +98,13 @@ fn extend_scalar(mut scalar: Scalar, n: usize) -> Scalar { /// /// This hashing-chain is performed using Poseidon hashing algorithm /// and relies on the `Hades252` permutation. -pub fn merkle_opening_scalar_verification( - branch: kelvin::Branch, H>, +pub fn merkle_opening_scalar_verification( + branch: PoseidonBranch, root: Scalar, leaf: Scalar, - branch_length: usize, -) -> bool -where - H: kelvin::ByteHash, -{ - let mut branch = PoseidonBranch::from(&branch); - let n_extensions = branch.extend(branch_length); - - let extended_root = extend_scalar(root, n_extensions); - +) -> bool { // Check that the root is indeed the one that we think - if branch.root != extended_root { + if branch.root != root { return false; }; // Allocate space for the last level computed hash as a variable to compare @@ -187,20 +152,18 @@ where mod tests { use super::*; use crate::hashing_utils::scalar_storage::StorageScalar; - use crate::PoseidonAnnotation; + use crate::PoseidonTree; use dusk_plonk::commitment_scheme::kzg10::PublicParameters; use dusk_plonk::fft::EvaluationDomain; - use kelvin::{Blake2b, Compound}; + use kelvin::Blake2b; use merlin::Transcript; - use nstack::NStack; - use std::borrow::Borrow; #[test] fn scalar_merkle_proof() { // Generate a tree with random scalars inside. - let mut nstack: NStack<_, PoseidonAnnotation, Blake2b> = NStack::new(); + let mut ptree: PoseidonTree<_, Blake2b> = PoseidonTree::new(17); for i in 0..1024u64 { - nstack.push(StorageScalar(Scalar::from(i as u64))).unwrap(); + ptree.push(StorageScalar(Scalar::from(i as u64))).unwrap(); } for i in 0..1024u64 { @@ -210,24 +173,15 @@ mod tests { // In this case, the key X corresponds to the Scalar(X). // We're supposing that we're provided with a Kelvin::Branch to perform // the proof. - let branch = nstack.get(i).unwrap().unwrap(); + let branch = ptree.poseidon_branch(i).unwrap().unwrap(); // Get tree root. - let root = StorageScalar::from( - branch - .levels() - .first() - .unwrap() - .annotation() - .unwrap() - .borrow(), - ); + let root = ptree.root().unwrap(); assert!(merkle_opening_scalar_verification( branch, - root.0.into(), + root, Scalar::from(i), - 17, )); } } @@ -239,9 +193,9 @@ mod tests { PublicParameters::setup(1 << 17, &mut rand::thread_rng()).unwrap(); let (ck, vk) = pub_params.trim(1 << 16).unwrap(); // Generate a tree with random scalars inside. - let mut nstack: NStack<_, PoseidonAnnotation, Blake2b> = NStack::new(); + let mut ptree: PoseidonTree<_, Blake2b> = PoseidonTree::new(17); for i in 0..1024u64 { - nstack.push(StorageScalar(Scalar::from(i as u64))).unwrap(); + ptree.push(StorageScalar(Scalar::from(i as u64))).unwrap(); } let mut composer_sizes = vec![]; @@ -255,29 +209,15 @@ mod tests { // In this case, the key X corresponds to the Scalar(X). // We're supposing that we're provided with a Kelvin::Branch to perform // the proof. - let branch = nstack.get(*i).unwrap().unwrap(); + let branch = ptree.poseidon_branch(*i).unwrap().unwrap(); // Get tree root. - let root = StorageScalar::from( - branch - .levels() - .first() - .unwrap() - .annotation() - .unwrap() - .borrow(), - ); + let root = ptree.root().unwrap(); // Add the proven leaf value to the Constraint System let proven_leaf = composer.add_input(Scalar::from(*i)); - merkle_opening_gadget( - &mut composer, - branch, - proven_leaf, - root.0.into(), - 17, - ); + merkle_opening_gadget(&mut composer, branch, proven_leaf, root); // Since we don't use all of the wires, we set some dummy constraints to avoid Committing // to zero polynomials. diff --git a/src/tree/mod.rs b/src/tree/mod.rs new file mode 100644 index 0000000..322f857 --- /dev/null +++ b/src/tree/mod.rs @@ -0,0 +1,128 @@ +use std::io; + +use dusk_bls12_381::Scalar; +use kelvin::{ + annotation, + annotations::{Cardinality, Count}, + Branch, BranchMut, ByteHash, Compound, Content, +}; +use nstack::NStack; + +use crate::merkle_lvl_hash::hash::*; +use crate::merkle_proof::poseidon_branch::extend_scalar; +use crate::ARITY; +use crate::{PoseidonBranch, PoseidonLevel, StorageScalar}; + +annotation! { + /// The annotation for the PoseidonTree + pub struct PoseidonAnnotation { + scalar: StorageScalar, + count: Cardinality, + } +} + +/// A zk-friendly datastructure to store elements +pub struct PoseidonTree +where + T: Content, + for<'any> &'any T: Into, + H: ByteHash, +{ + branch_depth: usize, + inner: NStack, +} + +impl PoseidonTree +where + T: Content, + for<'any> &'any T: Into, + H: ByteHash, +{ + /// Constructs a new empty PoseidonTree + pub fn new(depth: usize) -> Self { + PoseidonTree { + branch_depth: depth, + inner: Default::default(), + } + } + + /// Returns the scalar root-hash of the poseidon tree + /// + /// This includes padding the value to the correct branch length equivalent + pub fn root(&self) -> io::Result { + if let Some(ann) = self.inner.annotation() { + let borrow: &StorageScalar = ann.borrow(); + let scalar: Scalar = borrow.clone().into(); + + // FIXME, depth could be inferred from the cardinality + if let Some(branch) = self.get(0)? { + let depth = branch.levels().len(); + Ok(extend_scalar(scalar, self.branch_depth - depth)) + } else { + unreachable!("Annotation in empty tree") + } + } else { + // empty case, use an empty level for hashing + let leaves = [Scalar::zero(); ARITY + 1]; + let level = PoseidonLevel { leaves, offset: 0 }; + let root = merkle_level_hash_without_bitflags(&level); + Ok(extend_scalar(root, self.branch_depth)) + } + } + + /// Returns a poseidon branch pointing at the specific index + /// + /// This includes padding the value to the correct branch length equivalent + pub fn poseidon_branch( + &self, + idx: u64, + ) -> io::Result> { + Ok(self.inner.get(idx)?.map(|ref branch| { + let mut pbranch: PoseidonBranch = branch.into(); + pbranch.extend(self.branch_depth); + pbranch + })) + } + + /// Push a new item onto the tree + pub fn push(&mut self, t: T) -> io::Result { + let idx = self.inner.count(); + self.inner.push(t)?; + Ok(idx) + } + + /// Get a branch reference to the element at index `idx`, if any + pub fn get( + &self, + idx: u64, + ) -> io::Result, H>>> { + self.inner.get(idx) + } + + /// Get a mutable branch reference to the element at index `idx`, if any + pub fn get_mut( + &mut self, + idx: u64, + ) -> io::Result, H>>> + { + self.inner.get_mut(idx) + } +} + +#[cfg(test)] +mod test { + use super::*; + use kelvin::Blake2b; + + #[test] + fn insert() { + let mut tree = PoseidonTree::<_, Blake2b>::new(17); + + for i in 0..128u64 { + let idx = tree.push(StorageScalar::from(i)).unwrap(); + assert_eq!(idx, i); + } + + assert!(true) + } +}