From c08060e46aece32c385318f4611b975f5a0e12c7 Mon Sep 17 00:00:00 2001
From: Luca Giussani <lucagiussani@horizenlabs.io>
Date: Tue, 9 Nov 2021 11:44:55 +0100
Subject: [PATCH] Fix matrix balancing

---
 src/ahp/constraint_systems.rs | 30 +++++++++++++++---------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/src/ahp/constraint_systems.rs b/src/ahp/constraint_systems.rs
index 40f6b7a..c134ecd 100644
--- a/src/ahp/constraint_systems.rs
+++ b/src/ahp/constraint_systems.rs
@@ -34,20 +34,20 @@ fn to_matrix_helper<F: Field>(
     new_matrix
 }
 
-fn balance_matrices<F: Field>(a_matrix: &mut Matrix<F>, b_matrix: &mut Matrix<F>) {
-    let mut a_density: usize = a_matrix.iter().map(|row| row.len()).sum();
-    let mut b_density: usize = b_matrix.iter().map(|row| row.len()).sum();
-    let mut max_density = std::cmp::max(a_density, b_density);
-    let mut a_is_denser = a_density == max_density;
+// TODO: write a test to check that `balance_matrices` improves the balancing of the matrices
+// A and B by distributing the non-zero elements (more or less) evenly between the two.
+fn balance_matrices<F: Field>(a_matrix: &mut Vec<Vec<(F, VarIndex)>>, b_matrix: &mut Vec<Vec<(F, VarIndex)>>) {
+    let mut a_weight: usize = a_matrix.iter().map(|row| row.len()).sum();
+    let mut b_weight: usize = b_matrix.iter().map(|row| row.len()).sum();
     for (a_row, b_row) in a_matrix.iter_mut().zip(b_matrix) {
-        if a_is_denser {
-            let a_row_size = a_row.len();
-            let b_row_size = b_row.len();
+        let a_row_weight = a_row.len();
+        let b_row_weight = b_row.len();
+        if (a_weight < b_weight && a_row_weight < b_row_weight) ||
+            (a_weight > b_weight && a_row_weight > b_row_weight)
+        {
             std::mem::swap(a_row, b_row);
-            a_density = a_density - a_row_size + b_row_size;
-            b_density = b_density - b_row_size + a_row_size;
-            max_density = std::cmp::max(a_density, b_density);
-            a_is_denser = a_density == max_density;
+            a_weight = a_weight - a_row_weight + b_row_weight;
+            b_weight = b_weight - b_row_weight + a_row_weight;
         }
     }
 }
@@ -108,10 +108,10 @@ pub(crate) struct IndexerConstraintSystem<F: Field> {
 
 impl<F: Field> IndexerConstraintSystem<F> {
     pub(crate) fn process_matrices(&mut self) {
-        let mut a = to_matrix_helper(&self.a, self.num_input_variables);
-        let mut b = to_matrix_helper(&self.b, self.num_input_variables);
+        balance_matrices(&mut self.a, &mut self.b);
+        let a = to_matrix_helper(&self.a, self.num_input_variables);
+        let b = to_matrix_helper(&self.b, self.num_input_variables);
         let c = to_matrix_helper(&self.c, self.num_input_variables);
-        balance_matrices(&mut a, &mut b);
         self.a_matrix = Some(a);
         self.b_matrix = Some(b);
         self.c_matrix = Some(c);