From dba194c7eebf735c2fd576ee0da06cbd46d281dc Mon Sep 17 00:00:00 2001
From: schaeff <thibaut@powdrlabs.com>
Date: Thu, 4 Dec 2025 23:10:23 +0100
Subject: [PATCH] pass segment idx to do_with_trace callback

---
 openvm/src/empirical_constraints.rs | 7 ++-----
 openvm/src/lib.rs                   | 2 +-
 openvm/src/trace_generation.rs      | 5 +++--
 3 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/openvm/src/empirical_constraints.rs b/openvm/src/empirical_constraints.rs
index ae2ded80e2..40bca20357 100644
--- a/openvm/src/empirical_constraints.rs
+++ b/openvm/src/empirical_constraints.rs
@@ -81,9 +81,8 @@ pub fn detect_empirical_constraints(
 fn collect_trace(program: &CompiledProgram, inputs: StdIn) -> (Trace, DebugInfo) {
     let mut trace = Trace::default();
     let mut debug_info = DebugInfo::default();
-    let mut seg_idx = 0;
 
-    do_with_trace(program, inputs, |vm, _pk, ctx| {
+    do_with_trace(program, inputs, |seg_idx, vm, _pk, ctx| {
         let global_airs = vm
             .config()
             .create_airs()
@@ -130,7 +129,7 @@ fn collect_trace(program: &CompiledProgram, inputs: StdIn) -> (Trace, DebugInfo)
                 let row = Row {
                     cells: row,
                     pc: pc_value,
-                    timestamp: (seg_idx, ts_value),
+                    timestamp: (seg_idx as u32, ts_value),
                 };
                 trace.rows.push(row);
 
@@ -149,8 +148,6 @@ fn collect_trace(program: &CompiledProgram, inputs: StdIn) -> (Trace, DebugInfo)
                 }
             }
         }
-
-        seg_idx += 1;
     })
     .unwrap();
     (trace, debug_info)
diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs
index 6a43ce2452..65952988fb 100644
--- a/openvm/src/lib.rs
+++ b/openvm/src/lib.rs
@@ -776,7 +776,7 @@ pub fn prove(
     segment_height: Option<usize>, // uses the default height if None
 ) -> Result<(), Box<dyn std::error::Error>> {
     if mock {
-        do_with_trace(program, inputs, |vm, pk, ctx| {
+        do_with_trace(program, inputs, |_segment_idx, vm, pk, ctx| {
             debug_proving_ctx(vm, pk, &ctx);
         })?;
     } else {
diff --git a/openvm/src/trace_generation.rs b/openvm/src/trace_generation.rs
index 96454e9b34..4260792c45 100644
--- a/openvm/src/trace_generation.rs
+++ b/openvm/src/trace_generation.rs
@@ -18,11 +18,12 @@ use crate::SpecializedConfigCpuBuilder as SpecializedConfigBuilder;
 use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine;
 
 /// Given a program and input, generates the trace segment by segment and calls the provided
-/// callback with the VM, proving key, and proving context (containing the trace) for each segment.
+/// callback with the segment id, VM, proving key, and proving context (containing the trace) for each segment.
 pub fn do_with_trace(
     program: &CompiledProgram,
     inputs: StdIn,
     mut callback: impl FnMut(
+        usize,
         &VirtualMachine<BabyBearPoseidon2Engine, SpecializedConfigBuilder>,
         &MultiStarkProvingKey<BabyBearSC>,
         ProvingContext<<BabyBearPoseidon2Engine as StarkEngine>::PB>,
@@ -85,7 +86,7 @@ pub fn do_with_trace(
 
         let ctx = vm.generate_proving_ctx(system_records, record_arenas)?;
 
-        callback(&vm, &pk, ctx);
+        callback(seg_idx, &vm, &pk, ctx);
     }
     Ok(())
 }