Optimize hash_node

stwo_cairo_verifier/src/vcs/hasher.cairo

@@ -1,27 +1,35 @@
 use core::array::ArrayTrait;
-use core::option::OptionTrait;
-use core::poseidon::poseidon_hash_span;
+use core::hash::HashStateTrait;
+use core::poseidon::{poseidon_hash_span, hades_permutation, HashState};
 use stwo_cairo_verifier::BaseField;
 
-// A Merkle node hash is a hash of:
-//   [left_child_hash, right_child_hash], column0_value, column1_value, ...
-// "[]" denotes optional values.
-// The largest Merkle layer has no left and right child hashes. The rest of the layers have
-// children hashes.
-// At each layer, the tree may have multiple columns of the same length as the layer.
-// Each node in that layer contains one value from each column.
+/// 8 M31 elements fit in a hash, since 31*8 = 242 < 252.
+const M31_ELEMENETS_IN_HASH: usize = 8;
+
+const M31_ELEMENETS_IN_HASH_MINUS1: usize = M31_ELEMENETS_IN_HASH - 1;
+
+/// Equals `2^31`.
+const M31_IN_HASH_SHIFT: felt252 = 0x80000000;
+
+/// Equals `(2^31)^4`.
+const M31_IN_HASH_SHIFT_POW_4: felt252 = 0x10000000000000000000000000000000;
+
+/// A Merkle node hash is a hash of:
+///   [left_child_hash, right_child_hash], column0_value, column1_value, ...
+/// "[]" denotes optional values.
+/// The largest Merkle layer has no left and right child hashes. The rest of the layers have
+/// children hashes.
+/// At each layer, the tree may have multiple columns of the same length as the layer.
+/// Each node in that layer contains one value from each column.
 pub trait MerkleHasher {
     type Hash;
-    // Hashes a single Merkle node.
+
+    /// Hashes a single Merkle node.
     fn hash_node(
         children_hashes: Option<(Self::Hash, Self::Hash)>, column_values: Array<BaseField>,
     ) -> Self::Hash;
 }
 
-// 8 M31 elements fit in a hash, since 31*8 = 242 < 252.
-const M31_ELEMENETS_IN_HASH: usize = 8;
-const M31_ELEMENETS_IN_HASH_MINUS1: usize = M31_ELEMENETS_IN_HASH - 1;
-const M31_IN_HASH_SHIFT: felt252 = 0x80000000; // 2**31.
 pub impl PoseidonMerkleHasher of MerkleHasher {
     type Hash = felt252;
 
@@ -30,30 +38,54 @@ pub impl PoseidonMerkleHasher of MerkleHasher {
     ) -> Self::Hash {
         let mut hash_array: Array<felt252> = Default::default();
         if let Option::Some((x, y)) = children_hashes {
+            // Most often a node has no column values.
+            if column_values.len() == 0 {
+                // Inlined and simplified `poseidon_hash_span(...)` for better performance.
+                // TODO: Posiedon2 here { == hades_permutation(x, y, 2); }
+                let (s0, s1, s2) = hades_permutation(x, y, 0);
+                let hash_state = HashState { s0, s1, s2, odd: false };
+                return hash_state.finalize();
+            }
             hash_array.append(x);
             hash_array.append(y);
+        } else {
+            // Most often a single QM31 column commitment due to FRI.
+            // TODO(andrew): Implement non-mixed degree merkle for FRI decommitments.
+            if let Option::Some(values) = column_values.span().try_into() {
+                // Inlined and simplified `poseidon_hash_span(...)` for better performance.
+                let [v0, v1, v2, v3]: [BaseField; 4] = (*values).unbox();
+                let mut word = v0.inner.into();
+                word = word * M31_IN_HASH_SHIFT + v1.inner.into();
+                word = word * M31_IN_HASH_SHIFT + v2.inner.into();
+                word = word * M31_IN_HASH_SHIFT + v3.inner.into();
+                word = word * M31_IN_HASH_SHIFT_POW_4;
+                let (hash, _, _) = hades_permutation(word, 1, 0);
+                return hash;
+            }
         }
 
         // Pad column_values to a multiple of 8.
         let mut pad_len = M31_ELEMENETS_IN_HASH_MINUS1
             - ((column_values.len() + M31_ELEMENETS_IN_HASH_MINUS1) % M31_ELEMENETS_IN_HASH);
-        while pad_len > 0 {
-            column_values.append(core::num::traits::Zero::zero());
-            pad_len = M31_ELEMENETS_IN_HASH_MINUS1
-                - ((column_values.len() + M31_ELEMENETS_IN_HASH_MINUS1) % M31_ELEMENETS_IN_HASH);
+        let padding = core::num::traits::Zero::zero();
+        for _ in 0..pad_len {
+            column_values.append(padding);
         };
 
-        while !column_values.is_empty() {
-            let mut word = 0;
-            // Hash M31_ELEMENETS_IN_HASH = 8 M31 elements into a single field element.
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
-            word = word * M31_IN_HASH_SHIFT + column_values.pop_front().unwrap().inner.into();
+        let mut column_values = column_values.span();
+
+        // TODO(andrew): Measure performance diff and consider inlining `poseidon_hash_span(..)`
+        // functionality here to do all packing and hashing in a single pass.
+        while let Option::Some(values) = column_values.multi_pop_front::<8>() {
+            let [v0, v1, v2, v3, v4, v5, v6, v7] = (*values).unbox();
+            let mut word = v0.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v1.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v2.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v3.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v4.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v5.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v6.inner.into();
+            word = word * M31_IN_HASH_SHIFT + v7.inner.into();
             hash_array.append(word);
         };