From 4178749a759a4a7885b510d3ee4cbd1a1fb0f3f7 Mon Sep 17 00:00:00 2001
From: ohad-starkware <ohad.a@starkware.co>
Date: Wed, 29 Jan 2025 18:15:51 +0200
Subject: [PATCH] relation tracker preprocessed bug fix

---
 .../src/constraint_framework/component.rs     |  2 +-
 .../constraint_framework/relation_tracker.rs  | 26 +++++++++++++++++--
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/crates/prover/src/constraint_framework/component.rs b/crates/prover/src/constraint_framework/component.rs
index bb339c381..fffdab467 100644
--- a/crates/prover/src/constraint_framework/component.rs
+++ b/crates/prover/src/constraint_framework/component.rs
@@ -47,7 +47,7 @@ pub struct TraceLocationAllocator {
     /// Mapping of tree index to next available column offset.
     next_tree_offsets: TreeVec<usize>,
     /// Mapping of preprocessed columns to their index.
-    preprocessed_columns: Vec<PreProcessedColumnId>,
+    pub(super) preprocessed_columns: Vec<PreProcessedColumnId>,
     /// Controls whether the preprocessed columns are dynamic or static (default=Dynamic).
     preprocessed_columns_allocation_mode: PreprocessedColumnsAllocationMode,
 }
diff --git a/crates/prover/src/constraint_framework/relation_tracker.rs b/crates/prover/src/constraint_framework/relation_tracker.rs
index 9269a6118..174f52991 100644
--- a/crates/prover/src/constraint_framework/relation_tracker.rs
+++ b/crates/prover/src/constraint_framework/relation_tracker.rs
@@ -6,7 +6,7 @@ use num_traits::Zero;
 
 use super::{
     Batching, EvalAtRow, FrameworkEval, InfoEvaluator, Relation, RelationEntry,
-    TraceLocationAllocator, INTERACTION_TRACE_IDX,
+    TraceLocationAllocator, INTERACTION_TRACE_IDX, PREPROCESSED_TRACE_IDX,
 };
 use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES, N_LANES};
 use crate::core::backend::simd::qm31::PackedSecureField;
@@ -34,6 +34,7 @@ pub struct RelationTrackerEntry {
 pub struct RelationTrackerComponent<E: FrameworkEval> {
     eval: E,
     trace_locations: TreeVec<TreeSubspan>,
+    preprocessed_column_indices: Vec<usize>,
     n_rows: usize,
 }
 impl<E: FrameworkEval> RelationTrackerComponent<E> {
@@ -44,11 +45,27 @@ impl<E: FrameworkEval> RelationTrackerComponent<E> {
             SecureField::default(),
         ));
         let mut mask_offsets = info.mask_offsets;
+        let preprocessed_column_indices = info
+            .preprocessed_columns
+            .iter()
+            .map(|col| {
+                if let Some(pos) = location_allocator
+                    .preprocessed_columns
+                    .iter()
+                    .position(|x| x.id == col.id)
+                {
+                    pos
+                } else {
+                    panic!()
+                }
+            })
+            .collect();
         mask_offsets.drain(INTERACTION_TRACE_IDX..);
         let trace_locations = location_allocator.next_for_structure(&mask_offsets);
         Self {
             eval,
             trace_locations,
+            preprocessed_column_indices,
             n_rows,
         }
     }
@@ -60,9 +77,14 @@ impl<E: FrameworkEval> RelationTrackerComponent<E> {
         let log_size = self.eval.log_size();
 
         // Deref the sub-tree. Only copies the references.
-        let sub_tree = trace
+        let mut sub_tree = trace
             .sub_tree(&self.trace_locations)
             .map(|vec| vec.into_iter().copied().collect_vec());
+        sub_tree[PREPROCESSED_TRACE_IDX] = self
+            .preprocessed_column_indices
+            .iter()
+            .map(|idx| trace[PREPROCESSED_TRACE_IDX][*idx])
+            .collect();
         let mut entries = vec![];
 
         for vec_row in 0..(1 << (log_size - LOG_N_LANES)) {