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, 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 {