Skip to content

Commit

Permalink
✨ Merkle tree generation (#178)
Browse files Browse the repository at this point in the history
- [x] Add poseidon version of the merkle tree
- [x] Add proof computation
- [x] Clean and add corresponding tests

## Pull Request type

<!-- Please try to limit your pull request to one type; submit multiple
pull requests if needed. -->

Please check the type of change your PR introduces:

- [ ] Bugfix
- [x] Feature
- [ ] Code style update (formatting, renaming)
- [ ] Refactoring (no functional changes, no API changes)
- [ ] Build-related changes
- [ ] Documentation content changes
- [ ] Other (please describe):

## What is the current behavior?

<!-- Please describe the current behavior that you are modifying, or
link to a relevant issue. -->

Issue Number: #169  
Issue Number: #170 

## What is the new behavior?

<!-- Please describe the behavior or changes that are being added by
this PR. -->

- A new implementation of the MerkleTree using poseidon hash method
- Add a `compute_proof` method to the MerkleTreeTrait to generate a
proof based on leaves and an index
- Rename the previous merkle tree type from MerkleTree to
MerkleTreeLegacy
- Merkle tree tests refactoring to clarify tests and add tests according
to the new feature

## Does this introduce a breaking change?

- [x] Yes
- [ ] No

<!-- If this does introduce a breaking change, please describe the
impact and migration path for existing applications below. -->

Before the change you could just have:

```cairo
use alexandria_data_structures::merkle_tree::MerkleTreeTrait;
let mut merkle_tree = MerkleTreeTrait::new();
```

Now you must specify which kind of merkle tree you want (the legacy
version or the new one):

```cairo
// The new version, using the poseidon hash method
use alexandria_data_structures::merkle_tree::{MerkleTree, HashMethod, MerkleTreeTrait};
let hash_method = HashMethod::Poseidon(());
let mut merkle_tree = MerkleTreeTrait::new(hash_method);
```

```cairo
// The legacy version, using the pedersen hash method
use alexandria_data_structures::merkle_tree::{MerkleTree, HashMethod, MerkleTreeTrait};
let hash_method = HashMethod::Poseidon(());
let mut merkle_tree = MerkleTreeTrait::new(hash_method);
```

## Other information

Work load: ~ 8 hours

---------

Co-authored-by: LucasLvy <[email protected]>
Co-authored-by: Lucas @ StarkWare <[email protected]>
  • Loading branch information
3 people authored Sep 18, 2023
1 parent fa3f7e7 commit eb93e9b
Show file tree
Hide file tree
Showing 3 changed files with 330 additions and 94 deletions.
1 change: 1 addition & 0 deletions src/data_structures/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ The Merkle tree algorithm is a cryptographic hashing algorithm used to create a
The purpose of the Merkle tree algorithm is to provide a way to verify the integrity and authenticity of large amounts of data without needing to store the entire data set. The algorithm has applications in various areas of computer science, including cryptocurrency, file sharing, and database management.
The Merkle tree algorithm is also used in creating digital signatures and verifying the authenticity of transactions.
By providing a secure and efficient way to verify data integrity, the Merkle tree algorithm is an important tool in cryptography and information security.
A generic implementation is available to manage both pedersen (legacy) and poseidon hash methods.

## [Queue](./src/queue.cairo)

Expand Down
226 changes: 197 additions & 29 deletions src/data_structures/src/merkle_tree.cairo
Original file line number Diff line number Diff line change
@@ -1,64 +1,134 @@
//! MerkleTree implementation.
//!
//! # Example
//! # Examples
//!
//! ```
//! // This version uses the pedersen hash method because the PedersenHasherImpl is in the scope.
//! use alexandria_data_structures::merkle_tree::{Hasher, MerkleTree, pedersen::PedersenHasherImpl, MerkleTreeTrait};
//!
//! // Create a new merkle tree instance.
//! let mut merkle_tree: MerkleTree<Hasher> = MerkleTreeTrait::new();
//! let mut proof = array![element_1, element_2];
//! // Compute the merkle root.
//! let root = merkle_tree.compute_root(leaf, proof);
//! ```
//!
//! ```
//! use alexandria::data_structures::merkle_tree::MerkleTreeTrait;
//! // This version uses the poseidon hash method because the PoseidonHasherImpl is in the scope.
//! use alexandria_data_structures::merkle_tree::{ Hasher, MerkleTree, poseidon::PoseidonHasherImpl, MerkleTreeTrait };
//!
//! // Create a new merkle tree instance.
//! let mut merkle_tree = MerkleTreeTrait::new();
//! let mut proof = array![];
//! proof.append(element_1);
//! proof.append(element_2);
//! let mut merkle_tree: MerkleTree<PoseidonHasher> = MerkleTreeTrait::new();
//! let mut proof = array![element_1, element_2];
//! // Compute the merkle root.
//! let root = merkle_tree.compute_root(leaf, proof);
//! ```

// Core lib imports
use array::SpanTrait;
use hash::LegacyHash;
use traits::Into;
use array::{ArrayTrait, SpanTrait};
use traits::{Into, Copy, Drop};

/// Hasher trait.

trait HasherTrait<T> {
fn new() -> T;
fn hash(ref self: T, data1: felt252, data2: felt252) -> felt252;
}


// Hasher representations.

#[derive(Drop, Copy)]
struct Hasher {}

/// Hasher impls.

mod pedersen {
use pedersen::PedersenTrait;
use hash::HashStateTrait;
use super::{Hasher, HasherTrait};

impl PedersenHasherImpl of HasherTrait<Hasher> {
fn new() -> Hasher {
Hasher {}
}
fn hash(ref self: Hasher, data1: felt252, data2: felt252) -> felt252 {
let mut state = PedersenTrait::new(data1);
state = state.update(data2);
state.finalize()
}
}
}

mod poseidon {
use poseidon::PoseidonTrait;
use hash::HashStateTrait;
use super::{Hasher, HasherTrait};

impl PoseidonHasherImpl of HasherTrait<Hasher> {
fn new() -> Hasher {
Hasher {}
}
fn hash(ref self: Hasher, data1: felt252, data2: felt252) -> felt252 {
let mut state = PoseidonTrait::new();
state = state.update(data1);
state = state.update(data2);
state.finalize()
}
}
}

/// MerkleTree representation.

#[derive(Drop)]
struct MerkleTree {}
struct MerkleTree<T> {
hasher: T
}

/// MerkleTree trait.
trait MerkleTreeTrait {
trait MerkleTreeTrait<T> {
/// Create a new merkle tree instance.
fn new() -> MerkleTree;
fn new() -> MerkleTree<T>;
/// Compute the merkle root of a given proof.
fn compute_root(ref self: MerkleTree, current_node: felt252, proof: Span<felt252>) -> felt252;
fn compute_root(
ref self: MerkleTree<T>, current_node: felt252, proof: Span<felt252>
) -> felt252;
/// Verify a merkle proof.
fn verify(ref self: MerkleTree, root: felt252, leaf: felt252, proof: Span<felt252>) -> bool;
fn verify(ref self: MerkleTree<T>, root: felt252, leaf: felt252, proof: Span<felt252>) -> bool;
/// Compute a merkle proof of given leaves and at a given index.
fn compute_proof(ref self: MerkleTree<T>, leaves: Array<felt252>, index: u32) -> Span<felt252>;
}

/// MerkleTree implementation.
impl MerkleTreeImpl of MerkleTreeTrait {
/// MerkleTree Legacy implementation.
impl MerkleTreeImpl<
T, impl THasher: HasherTrait<T>, impl TCopy: Copy<T>, impl TDrop: Drop<T>
> of MerkleTreeTrait<T> {
/// Create a new merkle tree instance.
#[inline(always)]
fn new() -> MerkleTree {
MerkleTree {}
fn new() -> MerkleTree<T> {
MerkleTree { hasher: HasherTrait::new() }
}

/// Compute the merkle root of a given proof.
/// Compute the merkle root of a given proof using the generic T hasher.
/// # Arguments
/// * `current_node` - The current node of the proof.
/// * `proof` - The proof.
/// # Returns
/// The merkle root.
fn compute_root(
ref self: MerkleTree, mut current_node: felt252, mut proof: Span<felt252>
ref self: MerkleTree<T>, mut current_node: felt252, mut proof: Span<felt252>
) -> felt252 {
loop {
match proof.pop_front() {
Option::Some(proof_element) => {
// Compute the hash of the current node and the current element of the proof.
// We need to check if the current node is smaller than the current element of the proof.
// If it is, we need to swap the order of the hash.
if Into::<felt252, u256>::into(current_node) < (*proof_element).into() {
current_node = LegacyHash::hash(current_node, *proof_element);
} else {
current_node = LegacyHash::hash(*proof_element, current_node);
}
current_node =
if Into::<felt252, u256>::into(current_node) < (*proof_element).into() {
self.hasher.hash(current_node, *proof_element)
} else {
self.hasher.hash(*proof_element, current_node)
};
},
Option::None => {
break current_node;
Expand All @@ -67,17 +137,115 @@ impl MerkleTreeImpl of MerkleTreeTrait {
}
}

/// Verify a merkle proof.
/// Verify a merkle proof using the generic T hasher.
/// # Arguments
/// * `root` - The merkle root.
/// * `leaf` - The leaf to verify.
/// * `proof` - The proof.
/// # Returns
/// True if the proof is valid, false otherwise.
fn verify(
ref self: MerkleTree, root: felt252, leaf: felt252, mut proof: Span<felt252>
ref self: MerkleTree<T>, root: felt252, mut leaf: felt252, mut proof: Span<felt252>
) -> bool {
let computed_root = self.compute_root(leaf, proof);
let computed_root = loop {
match proof.pop_front() {
Option::Some(proof_element) => {
// Compute the hash of the current node and the current element of the proof.
// We need to check if the current node is smaller than the current element of the proof.
// If it is, we need to swap the order of the hash.
leaf =
if Into::<felt252, u256>::into(leaf) < (*proof_element).into() {
self.hasher.hash(leaf, *proof_element)
} else {
self.hasher.hash(*proof_element, leaf)
};
},
Option::None => {
break leaf;
},
};
};
computed_root == root
}

/// Compute a merkle proof of given leaves and at a given index using the generic T hasher.
/// # Arguments
/// * `leaves` - The sorted leaves.
/// * `index` - The index of the given.
/// # Returns
/// The merkle proof.
fn compute_proof(
ref self: MerkleTree<T>, mut leaves: Array<felt252>, index: u32
) -> Span<felt252> {
let mut proof: Array<felt252> = array![];
compute_proof(leaves, self.hasher, index, ref proof);
proof.span()
}
}

/// Helper function to compute a merkle proof of given leaves and at a given index.
/// # Arguments
/// * `nodes` - The sorted nodes.
/// * `index` - The index of the given.
/// * `hasher` - The hasher to use.
/// * `proof` - The proof array to fill.
fn compute_proof<T, impl THasher: HasherTrait<T>, impl TDrop: Drop<T>>(
mut nodes: Array<felt252>, mut hasher: T, index: u32, ref proof: Array<felt252>
) {
// Break if we have reached the top of the tree
if nodes.len() == 1 {
return;
}

// If odd number of nodes, add a null virtual leaf
if nodes.len() % 2 != 0 {
nodes.append(0);
}

// Compute next level
let mut next_level: Array<felt252> = get_next_level(nodes.span(), ref hasher);

// Find neighbor node
let mut index_parent = 0;
let mut i = 0;
loop {
if i == index {
index_parent = i / 2;
if i % 2 == 0 {
proof.append(*nodes.at(i + 1));
} else {
proof.append(*nodes.at(i - 1));
}
break;
}
i += 1;
};

compute_proof(next_level, hasher, index_parent, ref proof)
}

/// Helper function to compute the next layer of a merkle tree providing a layer of nodes.
/// # Arguments
/// * `nodes` - The sorted nodes.
/// * `hasher` - The hasher to use.
/// # Returns
/// The next layer of nodes.
fn get_next_level<T, impl THasher: HasherTrait<T>, impl TDrop: Drop<T>>(
mut nodes: Span<felt252>, ref hasher: T
) -> Array<felt252> {
let mut next_level: Array<felt252> = array![];
loop {
if nodes.is_empty() {
break;
}
let left = *nodes.pop_front().expect('Index out of bounds');
let right = *nodes.pop_front().expect('Index out of bounds');
let node = if Into::<felt252, u256>::into(left) < right.into() {
hasher.hash(left, right)
} else {
hasher.hash(right, left)
};
next_level.append(node);
};
next_level
}
Loading

0 comments on commit eb93e9b

Please sign in to comment.