feat: patricia hash tree

Signed-off-by: Dori Medini <[email protected]>

Cargo.toml

@@ -11,6 +11,7 @@ description = "Starknet Rust types related to computation and execution."
 testing = []
 
 [dependencies]
+bitvec = "1.0.1"
 cairo-lang-starknet-classes = "2.6.0"
 derive_more = "0.99.17"
 hex = "0.4.3"

src/lib.rs

@@ -11,6 +11,7 @@ pub mod deprecated_contract_class;
 pub mod external_transaction;
 pub mod hash;
 pub mod internal_transaction;
+pub mod patricia_hash;
 pub mod serde_utils;
 pub mod state;
 pub mod transaction;

src/patricia_hash.rs

@@ -0,0 +1,133 @@
+//! Patricia hash tree implementation.
+//!
+//! Supports root hash calculation for Stark felt values, keyed by consecutive 64 bits numbers,
+//! starting from 0.
+//!
+//! Each edge is marked with one or more bits.
+//! The key of a node is the concatenation of the edges' marks in the path from the root to this
+//! node.
+//! The input keys are in the leaves, and each leaf is an input key.
+//!
+//! The edges coming out of an internal node with a key `K` are:
+//! - If there are input keys that start with 'K0...' and 'K1...', then two edges come out, marked
+//! with '0' and '1' bits.
+//! - Otherwise, a single edge mark with 'Z' is coming out. 'Z' is the longest string, such that all
+//! the input keys that start with 'K...' start with 'KZ...' as well. Note, the order of the input
+//! keys in this implementation forces 'Z' to be a zeros string.
+//!
+//! Hash of a node depends on the number of edges coming out of it:
+//! - A leaf: The hash is the input value of its key.
+//! - A single edge: pedersen_hash(child_hash, edge_mark) + edge_length.
+//! - '0' and '1' edges: pedersen_hash(zero_child_hash, one_child_hash).
+
+#[cfg(test)]
+#[path = "patricia_hash_test.rs"]
+mod patricia_hash_test;
+
+use bitvec::prelude::{BitArray, Msb0};
+use starknet_crypto::FieldElement;
+
+use crate::hash::{pedersen_hash, StarkFelt};
+
+const TREE_HEIGHT: u8 = 64;
+type BitPath = BitArray<[u8; 8], Msb0>;
+
+// An entry in a Patricia tree.
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
+struct Entry {
+    key: BitPath,
+    value: StarkFelt,
+}
+
+// A sub-tree is defined by a sub-sequence of leaves with a common ancestor at the specified height,
+// with no other leaves under it besides these.
+#[derive(Debug)]
+struct SubTree<'a> {
+    leaves: &'a [Entry],
+    // Levels from the root.
+    height: u8,
+}
+
+enum SubTreeSplitting {
+    // Number of '0' bits that all the keys start with.
+    CommonZerosPrefix(u8),
+    // The index of the first key that starts with a '1' bit.
+    PartitionPoint(usize),
+}
+
+/// Calculates Patricia hash root on the given values.
+/// The values are keyed by consecutive numbers, starting from 0.
+pub fn calculate_root(values: Vec<StarkFelt>) -> StarkFelt {
+    if values.is_empty() {
+        return StarkFelt::ZERO;
+    }
+    let leaves: Vec<Entry> = values
+        .into_iter()
+        .zip(0u64..)
+        .map(|(felt, idx)| Entry { key: idx.to_be_bytes().into(), value: felt })
+        .collect();
+    get_hash(SubTree { leaves: &leaves[..], height: 0_u8 })
+}
+
+// Recursive hash calculation. There are 3 cases:
+// - Leaf: The sub tree height is maximal. It should contain exactly one entry.
+// - Edge: All the keys start with a longest common ('0's) prefix. NOTE: We assume that the keys are
+// a continuous range, and hence the case of '1's in the longest common prefix is impossible.
+// - Binary: Some keys start with '0' bit and some start with '1' bit.
+fn get_hash(sub_tree: SubTree<'_>) -> StarkFelt {
+    if sub_tree.height == TREE_HEIGHT {
+        return sub_tree.leaves.first().expect("a leaf should not be empty").value;
+    }
+    match get_splitting(&sub_tree) {
+        SubTreeSplitting::CommonZerosPrefix(n_zeros) => get_edge_hash(sub_tree, n_zeros),
+        SubTreeSplitting::PartitionPoint(partition_point) => {
+            get_binary_hash(sub_tree, partition_point)
+        }
+    }
+}
+
+// Hash on a '0's sequence with the bottom sub tree.
+fn get_edge_hash(sub_tree: SubTree<'_>, n_zeros: u8) -> StarkFelt {
+    let child_hash =
+        get_hash(SubTree { leaves: sub_tree.leaves, height: sub_tree.height + n_zeros });
+    let child_and_path_hash = pedersen_hash(&child_hash, &StarkFelt::ZERO);
+    StarkFelt::from(FieldElement::from(child_and_path_hash) + FieldElement::from(n_zeros))
+}
+
+// Hash on both sides: starts with '0' bit and starts with '1' bit.
+// Assumes: 0 < partition point < sub_tree.len().
+fn get_binary_hash(sub_tree: SubTree<'_>, partition_point: usize) -> StarkFelt {
+    let zero_hash = get_hash(SubTree {
+        leaves: &sub_tree.leaves[..partition_point],
+        height: sub_tree.height + 1,
+    });
+    let one_hash = get_hash(SubTree {
+        leaves: &sub_tree.leaves[partition_point..],
+        height: sub_tree.height + 1,
+    });
+    pedersen_hash(&zero_hash, &one_hash)
+}
+
+// Returns the manner the keys of a subtree are splitting: some keys start with '1' or all keys
+// start with '0'.
+fn get_splitting(sub_tree: &SubTree<'_>) -> SubTreeSplitting {
+    let mut height = sub_tree.height;
+
+    let first_one_bit_index =
+        sub_tree.leaves.partition_point(|entry| !entry.key[usize::from(height)]);
+    if first_one_bit_index < sub_tree.leaves.len() {
+        return SubTreeSplitting::PartitionPoint(first_one_bit_index);
+    }
+
+    height += 1;
+    let mut n_zeros = 1;
+
+    while height < TREE_HEIGHT {
+        if sub_tree.leaves.last().expect("sub tree should not be empty").key[usize::from(height)] {
+            break;
+        }
+        n_zeros += 1;
+        height += 1;
+    }
+    SubTreeSplitting::CommonZerosPrefix(n_zeros)
+}

src/patricia_hash_test.rs

@@ -0,0 +1,32 @@
+use super::calculate_root;
+use crate::hash::StarkFelt;
+
+// The expected roots were calculated by the starkware-libs/cairo-lang repository. These are the
+// roots of PatriciaTree objects with the same leaves.
+#[test]
+fn test_patricia() {
+    let root =
+        calculate_root(vec![StarkFelt::from(1_u8), StarkFelt::from(2_u8), StarkFelt::from(3_u8)]);
+    let expected_root =
+        StarkFelt::try_from("0x231e110514ca3a27707cd6c365e00685142d43b03d26f6274db51cbfa91aa1c")
+            .unwrap();
+    assert_eq!(root, expected_root);
+}
+
+#[test]
+fn test_edge_patricia() {
+    let root = calculate_root(vec![StarkFelt::from(1_u8)]);
+    let expected_root =
+        StarkFelt::try_from("0x268a9d47dde48af4b6e2c33932ed1c13adec25555abaa837c376af4ea2f8ad4")
+            .unwrap();
+    assert_eq!(root, expected_root);
+}
+
+#[test]
+fn test_binary_patricia() {
+    let root = calculate_root(vec![StarkFelt::from(1_u8), StarkFelt::from(2_u8)]);
+    let expected_root =
+        StarkFelt::try_from("0x599927f1181d5633c6f680dbf039534de49c44e0b9903c5305b2582dfd6a56a")
+            .unwrap();
+    assert_eq!(root, expected_root);
+}