compute program_segment in the adapter and add it to ProverInput (#1671)

Author: ohad-nir-starkware

stwo_cairo_prover/crates/adapter/src/adapter.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
 use cairo_vm::vm::runners::cairo_runner::CairoRunner;
+use itertools::Itertools;
 use tracing::{info, span, Level};
 
 use super::memory::{MemoryBuilder, MemoryConfig};
@@ -41,6 +42,17 @@ pub fn adapt(runner: &CairoRunner) -> Result<ProverInput> {
     // TODO(spapini): Add output builtin to public memory.
     let (memory, inst_cache) = memory.build();
 
+    // Compute program segment.
+    let initial_pc = state_transitions.initial_state.pc.0;
+    let initial_ap = state_transitions.initial_state.ap.0;
+    let program = (initial_pc..initial_ap - 2)
+        .map(|addr| {
+            let id = memory.get_raw_id(addr);
+            let value = memory.get(addr).as_u256();
+            (id, value)
+        })
+        .collect_vec();
+
     // TODO(Ohad): take this from the input.
     let public_segment_context = PublicSegmentContext::bootloader_context();
 
@@ -51,6 +63,7 @@ pub fn adapt(runner: &CairoRunner) -> Result<ProverInput> {
         public_memory_addresses,
         builtin_segments,
         public_segment_context,
+        program,
         #[cfg(feature = "extract-mem-trace")]
         relocated_mem: relocated_memory_clone,
         #[cfg(feature = "extract-mem-trace")]

stwo_cairo_prover/crates/adapter/src/lib.rs

@@ -32,6 +32,7 @@ pub struct ProverInput {
     pub public_memory_addresses: Vec<u32>,
     pub builtin_segments: BuiltinSegments,
     pub public_segment_context: PublicSegmentContext,
+    pub program: Vec<(u32, [u32; 8])>,
     #[cfg(feature = "extract-mem-trace")]
     pub relocated_mem: Vec<MemoryEntry>,
     #[cfg(feature = "extract-mem-trace")]

stwo_cairo_prover/crates/prover/src/witness/cairo.rs

@@ -64,23 +64,17 @@ fn extract_public_segments(
 
 fn extract_sections_from_memory(
     memory: &Memory,
-    initial_pc: u32,
     initial_ap: u32,
     final_ap: u32,
+    program: Vec<(u32, [u32; 8])>,
     public_segment_context: PublicSegmentContext,
 ) -> PublicMemory {
     let public_segments =
         extract_public_segments(memory, initial_ap, final_ap, public_segment_context);
-    let program_memory_addresses = initial_pc..initial_ap - 2;
     let safe_call_addresses = initial_ap - 2..initial_ap;
     let output_memory_addresses =
         public_segments.output.start_ptr.value..public_segments.output.stop_ptr.value;
-    let [program, safe_call, output] = [
-        program_memory_addresses,
-        safe_call_addresses,
-        output_memory_addresses,
-    ]
-    .map(|range| {
+    let [safe_call, output] = [safe_call_addresses, output_memory_addresses].map(|range| {
         range
             .map(|addr| {
                 let id = memory.get_raw_id(addr);
@@ -113,6 +107,7 @@ pub fn create_cairo_claim_generator(
         public_memory_addresses,
         builtin_segments,
         public_segment_context,
+        program,
         ..
     }: ProverInput,
     preprocessed_trace: Arc<PreProcessedTrace>,
@@ -135,14 +130,13 @@ pub fn create_cairo_claim_generator(
     all_components.insert("verify_bitwise_xor_9");
 
     // Public data.
-    let initial_pc = initial_state.pc.0;
     let initial_ap = initial_state.ap.0;
     let final_ap = final_state.ap.0;
     let public_memory = extract_sections_from_memory(
         &memory,
-        initial_pc,
         initial_ap,
         final_ap,
+        program,
         public_segment_context,
     );