From e25b37d64c1aa175661cb8402ef2707e41f60773 Mon Sep 17 00:00:00 2001
From: Ohad Agadi <ohad.a@starkware.co>
Date: Tue, 3 Dec 2024 10:14:23 +0200
Subject: [PATCH] relation tracker bug fix

---
 .../constraint_framework/relation_tracker.rs  | 31 +++++++++++++++++--
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/crates/prover/src/constraint_framework/relation_tracker.rs b/crates/prover/src/constraint_framework/relation_tracker.rs
index 3866df39a..f372fdc59 100644
--- a/crates/prover/src/constraint_framework/relation_tracker.rs
+++ b/crates/prover/src/constraint_framework/relation_tracker.rs
@@ -176,7 +176,10 @@ impl<'a> EvalAtRow for RelationTrackerEvaluator<'a> {
                 if cannonical_index >= self.n_rows {
                     continue;
                 }
-                let values = values.iter().map(|v| v[j]).collect_vec();
+                let mut values = values.iter().map(|v| v[j]).collect_vec();
+                while values.last().is_some_and(|v| v.is_zero()) {
+                    values.pop();
+                }
                 let mult = mult[j].to_m31_array()[0];
                 self.entries.push(RelationTrackerEntry {
                     relation: relation.clone(),
@@ -194,9 +197,15 @@ impl RelationSummary {
     /// Returns the sum of every entry's yields and uses.
     /// The result is a map from relation name to a list of values(M31 vectors) and their sum.
     pub fn summarize_relations(entries: &[RelationTrackerEntry]) -> Self {
+        let mut entry_by_relation = HashMap::new();
+        for entry in entries {
+            entry_by_relation
+                .entry(entry.relation.clone())
+                .or_insert_with(Vec::new)
+                .push(entry);
+        }
         let mut summary = vec![];
-        let relations = entries.iter().group_by(|entry| entry.relation.clone());
-        for (relation, entries) in &relations {
+        for (relation, entries) in entry_by_relation {
             let mut relation_sums: HashMap<Vec<_>, M31> = HashMap::new();
             for entry in entries {
                 let mult = relation_sums
@@ -216,6 +225,22 @@ impl RelationSummary {
             .find(|(name, _)| name == relation)
             .map(|(_, entries)| entries.as_slice())
     }
+
+    pub fn cleaned(self) -> Self {
+        let mut cleaned = vec![];
+        for (relation, entries) in self.0 {
+            let mut cleaned_entries = vec![];
+            for (vector, sum) in entries {
+                if !sum.is_zero() {
+                    cleaned_entries.push((vector, sum));
+                }
+            }
+            if !cleaned_entries.is_empty() {
+                cleaned.push((relation, cleaned_entries));
+            }
+        }
+        Self(cleaned)
+    }
 }
 impl Debug for RelationSummary {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {