Fix partition hashing (#338)

air/src/options.rs

@@ -377,7 +377,11 @@ impl PartitionOptions {
 
     /// Returns the size of each partition used when committing to the main and auxiliary traces as
     /// well as the constraint evaluation trace.
+    /// The returned size is given in terms of number of columns in the field `E`.
     pub fn partition_size<E: FieldElement>(&self, num_columns: usize) -> usize {
+        if self.num_partitions == 1 && self.min_partition_size == 1 {
+            return num_columns;
+        }
         let base_elements_per_partition = cmp::max(
             (num_columns * E::EXTENSION_DEGREE).div_ceil(self.num_partitions as usize),
             self.min_partition_size as usize,

prover/src/lib.rs

@@ -219,6 +219,7 @@ pub trait Prover {
     /// Builds and returns the auxiliary trace.
     #[allow(unused_variables)]
     #[maybe_async]
+    #[instrument(skip_all)]
     fn build_aux_trace<E>(
         &self,
         main_trace: &Self::Trace,

prover/src/matrix/row_matrix.rs

@@ -188,7 +188,7 @@ impl<E: FieldElement> RowMatrix<E> {
         // allocate vector to store row hashes
         let mut row_hashes = unsafe { uninit_vector::<H::Digest>(self.num_rows()) };
 
-        if partition_size == self.num_cols() * E::EXTENSION_DEGREE {
+        if partition_size == self.num_cols() {
             // iterate though matrix rows, hashing each row
             batch_iter_mut!(
                 &mut row_hashes,

verifier/src/channel.rs

@@ -442,7 +442,7 @@ where
     E: FieldElement,
     H: ElementHasher<BaseField = E::BaseField>,
 {
-    if partition_size == row.len() * E::EXTENSION_DEGREE {
+    if partition_size == row.len() {
         H::hash_elements(row)
     } else {
         let mut buffer = vec![H::Digest::default(); partition_size];