chore(rln): add QoL traits to the Hasher associated type in MerkleTree trait (#238)

This commit is contained in:
Aaryamann Challani 2024-04-17 17:54:49 +05:30 committed by GitHub
parent d91a5b3568
commit 0521c7349e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 74 additions and 26 deletions

View File

@ -1,5 +1,6 @@
use criterion::{criterion_group, criterion_main, Criterion};
use hex_literal::hex;
use std::{fmt::Display, str::FromStr};
use tiny_keccak::{Hasher as _, Keccak};
use zerokit_utils::{
FullMerkleConfig, FullMerkleTree, Hasher, OptimalMerkleConfig, OptimalMerkleTree,
@ -9,34 +10,53 @@ use zerokit_utils::{
#[derive(Clone, Copy, Eq, PartialEq)]
struct Keccak256;
#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)]
struct TestFr([u8; 32]);
impl Hasher for Keccak256 {
type Fr = [u8; 32];
type Fr = TestFr;
fn default_leaf() -> Self::Fr {
[0; 32]
TestFr([0; 32])
}
fn hash(inputs: &[Self::Fr]) -> Self::Fr {
let mut output = [0; 32];
let mut hasher = Keccak::v256();
for element in inputs {
hasher.update(element);
hasher.update(element.0.as_slice());
}
hasher.finalize(&mut output);
output
TestFr(output)
}
}
impl Display for TestFr {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", String::from_utf8_lossy(self.0.as_slice()))
}
}
impl FromStr for TestFr {
type Err = std::string::FromUtf8Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(TestFr(s.as_bytes().try_into().unwrap()))
}
}
pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) {
let mut tree =
OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default()).unwrap();
OptimalMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), OptimalMerkleConfig::default())
.unwrap();
let leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000001"),
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
]
.map(|x| TestFr(x));
c.bench_function("OptimalMerkleTree::set", |b| {
b.iter(|| {
@ -71,14 +91,15 @@ pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) {
pub fn full_merkle_tree_benchmark(c: &mut Criterion) {
let mut tree =
FullMerkleTree::<Keccak256>::new(2, [0; 32], FullMerkleConfig::default()).unwrap();
FullMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), FullMerkleConfig::default()).unwrap();
let leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000001"),
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
]
.map(|x| TestFr(x));
c.bench_function("FullMerkleTree::set", |b| {
b.iter(|| {

View File

@ -21,7 +21,7 @@ use color_eyre::Result;
/// and the hash function used to initialize a Merkle Tree implementation
pub trait Hasher {
/// Type of the leaf and tree node
type Fr: Clone + Copy + Eq;
type Fr: Clone + Copy + Eq + Default + std::fmt::Debug + std::fmt::Display + FromStr;
/// Returns the default tree leaf
fn default_leaf() -> Self::Fr;

View File

@ -1,6 +1,8 @@
// Tests adapted from https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/merkle_tree.rs
#[cfg(test)]
mod test {
pub mod test {
use std::{fmt::Display, str::FromStr};
use hex_literal::hex;
use tiny_keccak::{Hasher as _, Keccak};
use zerokit_utils::{
@ -10,21 +12,38 @@ mod test {
#[derive(Clone, Copy, Eq, PartialEq)]
struct Keccak256;
#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)]
struct TestFr([u8; 32]);
impl Hasher for Keccak256 {
type Fr = [u8; 32];
type Fr = TestFr;
fn default_leaf() -> Self::Fr {
[0; 32]
TestFr([0; 32])
}
fn hash(inputs: &[Self::Fr]) -> Self::Fr {
let mut output = [0; 32];
let mut hasher = Keccak::v256();
for element in inputs {
hasher.update(element);
hasher.update(element.0.as_slice());
}
hasher.finalize(&mut output);
output
TestFr(output)
}
}
impl Display for TestFr {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", String::from_utf8_lossy(self.0.as_slice()))
}
}
impl FromStr for TestFr {
type Err = std::string::FromUtf8Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(TestFr(s.as_bytes().try_into().unwrap()))
}
}
@ -35,20 +54,24 @@ mod test {
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
]
.map(|x| TestFr(x));
let default_tree_root =
hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30");
let default_tree_root = TestFr(hex!(
"b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30"
));
let roots = [
hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95"),
hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b"),
hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c"),
hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"),
];
]
.map(|x| TestFr(x));
let mut tree =
FullMerkleTree::<Keccak256>::new(2, [0; 32], FullMerkleConfig::default()).unwrap();
FullMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), FullMerkleConfig::default())
.unwrap();
assert_eq!(tree.root(), default_tree_root);
for i in 0..leaves.len() {
tree.set(i, leaves[i]).unwrap();
@ -56,7 +79,7 @@ mod test {
}
let mut tree =
OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default())
OptimalMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), OptimalMerkleConfig::default())
.unwrap();
assert_eq!(tree.root(), default_tree_root);
for i in 0..leaves.len() {
@ -72,11 +95,13 @@ mod test {
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
]
.map(|x| TestFr(x));
// We thest the FullMerkleTree implementation
let mut tree =
FullMerkleTree::<Keccak256>::new(2, [0; 32], FullMerkleConfig::default()).unwrap();
FullMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), FullMerkleConfig::default())
.unwrap();
for i in 0..leaves.len() {
// We set the leaves
tree.set(i, leaves[i]).unwrap();
@ -101,7 +126,7 @@ mod test {
// We test the OptimalMerkleTree implementation
let mut tree =
OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default())
OptimalMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), OptimalMerkleConfig::default())
.unwrap();
for i in 0..leaves.len() {
// We set the leaves
@ -133,10 +158,11 @@ mod test {
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
]
.map(|x| TestFr(x));
let mut tree =
OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default())
OptimalMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), OptimalMerkleConfig::default())
.unwrap();
// We set the leaves
@ -145,7 +171,8 @@ mod test {
let new_leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000005"),
hex!("0000000000000000000000000000000000000000000000000000000000000006"),
];
]
.map(|x| TestFr(x));
let to_delete_indices: [usize; 2] = [0, 1];