unify code for hashing the master base table & master extension table

triton-vm/src/table/master_table.rs

@@ -16,6 +16,7 @@ use strum_macros::Display;
 use strum_macros::EnumCount as EnumCountMacro;
 use strum_macros::EnumIter;
 use twenty_first::shared_math::b_field_element::BFieldElement;
+use twenty_first::shared_math::digest::Digest;
 use twenty_first::shared_math::mpolynomial::Degree;
 use twenty_first::shared_math::other::roundup_npo2;
 use twenty_first::shared_math::polynomial::Polynomial;
@@ -266,6 +267,29 @@ where
             });
         (extended_columns, interpolation_polynomials)
     }
+
+    /// Compute a Merkle tree out of this table. Every row gives one leaf in the tree.
+    /// Therefore, constructing the tree is a two-step process:
+    /// 1. Hash each row.
+    /// 1. Construct the tree from the hashed rows.
+    ///
+    /// The function [`hash_row`](Self::hash_row) is used to hash each row.
+    fn merkle_tree(&self, maybe_profiler: &mut Option<TritonProfiler>) -> MerkleTree<StarkHasher> {
+        prof_start!(maybe_profiler, "leafs");
+        let master_matrix = self.master_matrix();
+        let all_rows = master_matrix.axis_iter(Axis(0)).into_par_iter();
+        let hashed_rows = all_rows.map(Self::hash_row).collect::<Vec<_>>();
+        prof_stop!(maybe_profiler, "leafs");
+
+        prof_start!(maybe_profiler, "Merkle tree");
+        let merkle_tree = MTMaker::from_digests(&hashed_rows);
+        prof_stop!(maybe_profiler, "Merkle tree");
+
+        merkle_tree
+    }
+
+    /// Hash one row.
+    fn hash_row(row: ArrayView1<FF>) -> Digest;
 }
 
 #[derive(Clone)]
@@ -334,6 +358,10 @@ impl MasterTable<BFieldElement> for MasterBaseTable {
     fn fri_domain(&self) -> ArithmeticDomain {
         self.fri_domain
     }
+
+    fn hash_row(row: ArrayView1<BFieldElement>) -> Digest {
+        StarkHasher::hash_varlen(row.as_slice().unwrap())
+    }
 }
 
 impl MasterTable<XFieldElement> for MasterExtTable {
@@ -366,6 +394,12 @@ impl MasterTable<XFieldElement> for MasterExtTable {
     fn fri_domain(&self) -> ArithmeticDomain {
         self.fri_domain
     }
+
+    fn hash_row(row: ArrayView1<XFieldElement>) -> Digest {
+        let interpret_xfe_as_bfes = |xfe: &XFieldElement| xfe.coefficients.to_vec();
+        let row_as_bfes = row.iter().map(interpret_xfe_as_bfes).concat();
+        StarkHasher::hash_varlen(&row_as_bfes)
+    }
 }
 
 impl MasterBaseTable {
@@ -517,29 +551,6 @@ impl MasterBaseTable {
         (master_base_table, interpolation_polynomials)
     }
 
-    /// Compute the Merkle tree of the base table. Every row is one leaf in the tree.
-    /// Therefore, constructing the tree is a two-step process:
-    /// 1. Hash each row.
-    /// 1. Construct the tree from the hashed rows.
-    pub fn merkle_tree(
-        &self,
-        maybe_profiler: &mut Option<TritonProfiler>,
-    ) -> MerkleTree<StarkHasher> {
-        prof_start!(maybe_profiler, "leafs");
-        let hashed_rows = self
-            .master_base_matrix
-            .axis_iter(Axis(0))
-            .into_par_iter()
-            .map(|row| StarkHasher::hash_varlen(&row.to_vec()))
-            .collect::<Vec<_>>();
-        prof_stop!(maybe_profiler, "leafs");
-        prof_start!(maybe_profiler, "Merkle tree");
-        let merkle_tree = MTMaker::from_digests(&hashed_rows);
-        prof_stop!(maybe_profiler, "Merkle tree");
-
-        merkle_tree
-    }
-
     /// Create a `MasterExtTable` from a `MasterBaseTable` by `.extend()`ing each individual base
     /// table. The `.extend()` for each table is specific to that table, but always involves
     /// adding some number of columns.
@@ -676,32 +687,6 @@ impl MasterExtTable {
         randomizer_polynomials
     }
 
-    pub fn merkle_tree(
-        &self,
-        maybe_profiler: &mut Option<TritonProfiler>,
-    ) -> MerkleTree<StarkHasher> {
-        prof_start!(maybe_profiler, "leafs");
-        let hashed_rows = self
-            .master_ext_matrix
-            .axis_iter(Axis(0))
-            .into_par_iter()
-            .map(|row| {
-                let contiguous_row_bfe = row
-                    .to_vec()
-                    .iter()
-                    .map(|xfe| xfe.coefficients.to_vec())
-                    .concat();
-                StarkHasher::hash_varlen(&contiguous_row_bfe)
-            })
-            .collect::<Vec<_>>();
-        prof_stop!(maybe_profiler, "leafs");
-        prof_start!(maybe_profiler, "Merkle tree");
-        let ret = MTMaker::from_digests(&hashed_rows);
-        prof_stop!(maybe_profiler, "Merkle tree");
-
-        ret
-    }
-
     fn table_slice_info(id: TableId) -> (usize, usize) {
         use TableId::*;
         match id {