From 5b455acbd6645a7901e2429705f2fbd233941ee9 Mon Sep 17 00:00:00 2001 From: arnaucube Date: Mon, 3 Feb 2025 18:03:45 +0100 Subject: [PATCH] add initial MerkleTree implementation (#13) Add initial MerkleTree implementation, which is a wrapper on top of Plonky2's MerkleTree, with the idea that the future iteration will replace it by the MerkleTree specified at https://0xparc.github.io/pod2/merkletree.html . --- Cargo.toml | 3 +- README.md | 3 +- book/src/merkletree.md | 8 +- rust-toolchain.toml | 3 + src/lib.rs | 19 ++++- src/merkletree.rs | 163 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 194 insertions(+), 5 deletions(-) create mode 100644 rust-toolchain.toml create mode 100644 src/merkletree.rs diff --git a/Cargo.toml b/Cargo.toml index 66e13b0..657f0c6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,8 +8,9 @@ name = "pod2" path = "src/lib.rs" [dependencies] -plonky2 = { git = "https://github.com/mir-protocol/plonky2" } +plonky2 = { git = "https://github.com/0xPolygonZero/plonky2" } hex = "0.4.3" itertools = "0.14.0" strum = "0.26" strum_macros = "0.26" +anyhow = "1.0.56" diff --git a/README.md b/README.md index 14a67a8..e7ace1e 100644 --- a/README.md +++ b/README.md @@ -2,4 +2,5 @@ Specification can be found at the [`book`](https://github.com/0xPARC/pod2/tree/main/book) directory, a live rendered version at: https://0xparc.github.io/pod2/ -Requires Rust nightly: `rustup override set nightly`. +## Usage +- Run tests: `cargo test --release` diff --git a/book/src/merkletree.md b/book/src/merkletree.md index 9c6bbaf..336a5e1 100644 --- a/book/src/merkletree.md +++ b/book/src/merkletree.md @@ -39,8 +39,12 @@ Since leaf positions are deterministic based on the key, the same approach is us ## Encoding > TODO: how key-values, nodes, merkle-proofs, ... are encoded. -## Temporary first version -The first iteration of the implementation uses a hash of the key-values concatenated, with the idea of replacing it by the MerkleTree approach described above. +## Developement plan +- short term: merkle tree as a 'precompile' in POD operations, which allows to directly verify proofs + - initial version: just a wrapper on top of the existing Plonky2's MerkleTree + - second iteration: implement the MerkleTree specified in this document +- long term exploration: + - explore feasibility of using Starky (for lookups) connected to Plonky2, which would allow doing the approach described at [https://hackmd.io/@aardvark/SkJ-wcTDJe](https://hackmd.io/@aardvark/SkJ-wcTDJe) ## Resources diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..a2d375e --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "nightly" +components = ["clippy", "rustfmt"] diff --git a/src/lib.rs b/src/lib.rs index 1ead253..c0ddee9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,13 +2,19 @@ use hex::{FromHex, FromHexError}; use plonky2::field::goldilocks_field::GoldilocksField; use plonky2::field::types::{Field, PrimeField64}; use plonky2::hash::poseidon::PoseidonHash; -use plonky2::plonk::config::Hasher; +use plonky2::plonk::config::{Hasher, PoseidonGoldilocksConfig}; +use std::cmp::Ordering; use std::fmt; pub mod backend; pub mod frontend; +pub mod merkletree; pub type F = GoldilocksField; +/// C is the Plonky2 config used in POD2 to work with Plonky2 recursion. +pub type C = PoseidonGoldilocksConfig; +/// D defines the extension degree of the field used in the Plonky2 proofs (quadratic extension). +pub const D: usize = 2; #[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)] pub struct Hash(pub [F; 4]); @@ -40,6 +46,17 @@ impl FromHex for Hash { } } +impl Ord for Hash { + fn cmp(&self, other: &Self) -> Ordering { + self.to_string().cmp(&other.to_string()) + } +} +impl PartialOrd for Hash { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)] pub struct PodId(pub Hash); pub const SELF: PodId = PodId(Hash([F::ONE, F::ZERO, F::ZERO, F::ZERO])); diff --git a/src/merkletree.rs b/src/merkletree.rs new file mode 100644 index 0000000..590dafd --- /dev/null +++ b/src/merkletree.rs @@ -0,0 +1,163 @@ +/// MerkleTree implementation for POD2. +/// +/// Current implementation is a wrapper on top of Plonky2's MerkleTree, but the future iteration +/// will replace it by the MerkleTree specified at https://0xparc.github.io/pod2/merkletree.html . +use anyhow::{anyhow, Result}; +use itertools::Itertools; +use plonky2::field::types::Field; +use plonky2::hash::{ + hash_types::HashOut, + merkle_proofs::{verify_merkle_proof, MerkleProof as PlonkyMerkleProof}, + merkle_tree::MerkleTree as PlonkyMerkleTree, + poseidon::PoseidonHash, +}; +use plonky2::plonk::config::GenericConfig; +use plonky2::plonk::config::Hasher; +use std::collections::HashMap; + +use crate::backend::Value; +use crate::{Hash, C, D, F}; + +const CAP_HEIGHT: usize = 0; + +/// MerkleTree currently implements the MerkleTree interface with a wrapper on top of Plonky2's +/// MerkleTree. A future iteration will replace it by the MerkleTree specified at +/// https://0xparc.github.io/pod2/merkletree.html . +#[derive(Clone, Debug)] +pub struct MerkleTree { + tree: PlonkyMerkleTree>::Hasher>, + // keyindex: key -> index mapping. This is just for the current plonky-tree wrapper + keyindex: HashMap, +} + +pub struct MerkleProof { + existence: bool, + index: usize, + proof: PlonkyMerkleProof>::Hasher>, +} + +impl MerkleTree { + pub fn new(kvs: HashMap) -> Self { + let mut keyindex: HashMap = HashMap::new(); + let mut leaves: Vec> = Vec::new(); + // Note: current version iterates sorting by keys of the kvs, but the merkletree defined at + // https://0xparc.github.io/pod2/merkletree.html will not need it since it will be + // deterministic based on the keys values not on the order of the keys when added into the + // tree. + for (i, (k, v)) in kvs.into_iter().sorted_by_key(|kv| kv.0).enumerate() { + let input: Vec = [k.0, v.0].concat(); + let leaf = PoseidonHash::hash_no_pad(&input).elements; + leaves.push(leaf.into()); + keyindex.insert(k, i); + } + + // pad to a power of two if needed + let leaf_empty: Vec = vec![F::ZERO, F::ZERO, F::ZERO, F::ZERO]; + for _ in leaves.len()..leaves.len().next_power_of_two() { + leaves.push(leaf_empty.clone()); + } + + let tree = PlonkyMerkleTree::>::Hasher>::new(leaves, CAP_HEIGHT); + Self { tree, keyindex } + } + + pub fn root(&self) -> Result { + if self.tree.cap.is_empty() { + return Err(anyhow!("empty tree")); + } + Ok(Hash(self.tree.cap.0[0].elements)) + } + + pub fn prove(&self, key: &Hash) -> Result { + let i = self.keyindex.get(&key).ok_or(anyhow!("key not in tree"))?; + let proof = self.tree.prove(*i); + Ok(MerkleProof { + existence: true, + index: *i, + proof, + }) + } + + pub fn prove_nonexistence(&self, _key: &Hash) -> Result { + // mock method + println!("WARNING: MerkleTree::verify_nonexistence is currently a mock"); + Ok(MerkleProof { + existence: false, + index: 0, + proof: PlonkyMerkleProof { siblings: vec![] }, + }) + } + + pub fn verify(root: Hash, proof: &MerkleProof, key: &Hash, value: &Value) -> Result<()> { + if !proof.existence { + return Err(anyhow!( + "expected proof of existence, found proof of non-existence" + )); + } + let leaf = PoseidonHash::hash_no_pad(&[key.0, value.0].concat()).elements; + let root = HashOut::from_vec(root.0.to_vec()); + verify_merkle_proof(leaf.into(), proof.index, root, &proof.proof) + } + + pub fn verify_nonexistence( + _root: Hash, + proof: &MerkleProof, + _key: &Hash, + _value: &Value, + ) -> Result<()> { + // mock method + if proof.existence { + return Err(anyhow!( + "expected proof of non-existence, found proof of existence" + )); + } + println!("WARNING: MerkleTree::verify_nonexistence is currently a mock"); + Ok(()) + } +} + +#[cfg(test)] +pub mod tests { + use super::*; + + use crate::hash_str; + + #[test] + fn test_merkletree() -> Result<()> { + let (k0, v0) = ( + hash_str("key_0".into()), + Value(hash_str("value_0".into()).0), + ); + let (k1, v1) = ( + hash_str("key_1".into()), + Value(hash_str("value_1".into()).0), + ); + let (k2, v2) = ( + hash_str("key_2".into()), + Value(hash_str("value_2".into()).0), + ); + + let mut kvs = HashMap::new(); + kvs.insert(k0, v0); + kvs.insert(k1, v1); + kvs.insert(k2, v2); + + let tree = MerkleTree::new(kvs); + + let proof = tree.prove(&k2)?; + MerkleTree::verify(tree.root()?, &proof, &k2, &v2)?; + + // expect verification to fail with different key / value + assert!(MerkleTree::verify(tree.root()?, &proof, &k2, &v0).is_err()); + assert!(MerkleTree::verify(tree.root()?, &proof, &k0, &v2).is_err()); + + // non-existence proofs + let proof_ne = tree.prove_nonexistence(&k2)?; + let _ = MerkleTree::verify_nonexistence(tree.root()?, &proof_ne, &k2, &v2)?; + + // expect verification of existence fail for nonexistence proof + let _ = MerkleTree::verify(tree.root()?, &proof_ne, &k2, &v2).is_err(); + + Ok(()) + } +}