remove the program from PublicMemory when program_in_ppt=true

stwo_cairo_prover/crates/cairo-air/src/air.rs

@@ -237,8 +237,10 @@ impl PublicData {
         for (id, _) in output {
             public_claim.push(*id);
         }
-        for (id, _) in program {
-            public_claim.push(*id);
+        if let Some(program) = program {
+            for (id, _) in program {
+                public_claim.push(*id);
+            }
         }
 
         // Collect output values.
@@ -250,9 +252,13 @@ impl PublicData {
 
         // Collect program values.
         let mut program_claim = vec![];
-        for (_, value) in program {
-            program_claim
-                .extend::<[u32; FELT252_N_WORDS]>(split(*value, (1 << FELT252_BITS_PER_WORD) - 1));
+        if let Some(program) = program {
+            for (_, value) in program {
+                program_claim.extend::<[u32; FELT252_N_WORDS]>(split(
+                    *value,
+                    (1 << FELT252_BITS_PER_WORD) - 1,
+                ));
+            }
         }
 
         (public_claim, output_claim, program_claim)
@@ -469,7 +475,7 @@ pub type MemorySection = Vec<PubMemoryValue>;
 
 #[derive(Serialize, Deserialize, CairoSerialize, CairoDeserialize, Default, Clone)]
 pub struct PublicMemory {
-    pub program: MemorySection,
+    pub program: Option<MemorySection>,
     pub public_segments: PublicSegmentRanges,
     pub output: MemorySection,
     pub safe_call_ids: [u32; 2],
@@ -483,9 +489,14 @@ impl PublicMemory {
         initial_ap: u32,
         final_ap: u32,
     ) -> impl Iterator<Item = PubMemoryEntry> {
-        let [program, output] =
-            [&self.program, &self.output].map(|section| section.clone().into_iter().enumerate());
-        let program_iter = program.map(move |(i, (id, value))| (initial_pc + i as u32, id, value));
+        let output = self.output.clone().into_iter().enumerate();
+        let program_iter = self
+            .program
+            .clone()
+            .unwrap_or_default()
+            .into_iter()
+            .enumerate()
+            .map(move |(i, (id, value))| (initial_pc + i as u32, id, value));
         let output_iter = output.map(move |(i, (id, value))| (final_ap + i as u32, id, value));
 
         let [safe_call_id0, safe_call_id1] = self.safe_call_ids;
@@ -598,7 +609,7 @@ mod tests {
         let dummy_lookup_elements = CommonLookupElements::dummy();
         let public_data = PublicData {
             public_memory: PublicMemory {
-                program,
+                program: Some(program),
                 public_segments: PublicSegmentRanges {
                     output: SegmentRange {
                         start_ptr: MemorySmallValue {

stwo_cairo_prover/crates/cairo-air/src/flat_claims.rs

@@ -20,7 +20,13 @@ impl FlatClaim {
             self.component_log_sizes.iter().cloned(),
         ));
         channel.mix_felts(&pack_into_secure_felts(
-            [self.public_data.public_memory.program.len() as u32].into_iter(),
+            [self
+                .public_data
+                .public_memory
+                .program
+                .as_ref()
+                .map_or(0, |p| p.len()) as u32]
+            .into_iter(),
         ));
         self.public_data.mix_into::<MC>(channel);
     }

stwo_cairo_prover/crates/cairo-air/src/utils.rs

@@ -166,7 +166,11 @@ pub struct VerificationOutput {
 
 /// Extract program hash (blake2s) and public output from the public memory.
 pub fn get_verification_output(public_memory: &PublicMemory) -> VerificationOutput {
-    let program_hash = construct_f252(&encode_and_hash_memory_section(&public_memory.program));
+    let program_hash = public_memory
+        .program
+        .as_ref()
+        .map(|p| construct_f252(&encode_and_hash_memory_section(p)))
+        .unwrap_or_default();
     let output = public_memory
         .output
         .iter()

stwo_cairo_prover/crates/cairo-air/src/verifier.rs

@@ -63,7 +63,9 @@ fn verify_claim(claim: &CairoClaim) {
         public_segments,
     );
 
-    verify_program(program, public_segments);
+    if let Some(program) = program {
+        verify_program(program, public_segments);
+    }
 
     assert_eq!(*initial_pc, BaseField::one());
     assert!(

stwo_cairo_prover/crates/prover/src/debug_tools/assert_constraints.rs

@@ -265,7 +265,8 @@ pub fn assert_cairo_constraints(input: ProverInput, preprocessed_trace: Arc<PreP
     tree_builder.finalize_interaction();
 
     // Base trace.
-    let cairo_claim_generator = create_cairo_claim_generator(input, preprocessed_trace.clone());
+    let cairo_claim_generator =
+        create_cairo_claim_generator(input, preprocessed_trace.clone(), false);
     let mut tree_builder = commitment_scheme.tree_builder();
     let (claim, interaction_generator) = cairo_claim_generator.write_trace(&mut tree_builder);
     tree_builder.finalize_interaction();

stwo_cairo_prover/crates/prover/src/prover.rs

@@ -165,7 +165,8 @@ where
     commitment_scheme.commit_tree(preprocessed_tree, channel);
 
     // Run Cairo.
-    let cairo_claim_generator = create_cairo_claim_generator(input, preprocessed_trace.clone());
+    let cairo_claim_generator =
+        create_cairo_claim_generator(input, preprocessed_trace.clone(), program_in_ppt);
     // Base trace.
     let mut tree_builder = commitment_scheme.tree_builder();
     let span = span!(Level::INFO, "Base trace").entered();

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

@@ -66,7 +66,7 @@ fn extract_sections_from_memory(
     memory: &Memory,
     initial_ap: u32,
     final_ap: u32,
-    program: Vec<(u32, [u32; 8])>,
+    program: Option<Vec<(u32, [u32; 8])>>,
     public_segment_context: PublicSegmentContext,
 ) -> PublicMemory {
     let public_segments =
@@ -111,6 +111,7 @@ pub fn create_cairo_claim_generator(
         ..
     }: ProverInput,
     preprocessed_trace: Arc<PreProcessedTrace>,
+    program_in_ppt: bool,
 ) -> CairoClaimGenerator {
     let initial_state = state_transitions.initial_state;
     let final_state = state_transitions.final_state;
@@ -130,13 +131,15 @@ 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 program_len = program.len() as u32;
     let public_memory = extract_sections_from_memory(
         &memory,
         initial_ap,
         final_ap,
-        program,
+        if program_in_ppt { None } else { Some(program) },
         public_segment_context,
     );
 
@@ -163,9 +166,12 @@ pub fn create_cairo_claim_generator(
     let memory_id_to_value_trace_generator =
         cairo_claim_generator.memory_id_to_big.as_ref().unwrap();
     // Yield public memory.
+    // When program_in_ppt is true, program entries are handled by the preprocessed trace,
+    // not by public memory. Exclude program addresses to keep the lookup argument balanced.
     for addr in public_memory_addresses
         .iter()
         .copied()
+        .filter(|&addr| !program_in_ppt || addr < initial_pc || addr >= initial_pc + program_len)
         .map(M31::from_u32_unchecked)
     {
         let id = memory_address_to_id_trace_generator.get_id(addr);

stwo_cairo_prover/test_data/test_prove_verify_all_opcode_components/proof.json

@@ -1,4 +1,5 @@
 [
+  "0x0",
   "0x535",
   "0x0",
   "0x7fff7fff",

stwo_cairo_prover/test_data/test_prove_verify_ret_opcode/proof.json

@@ -1,4 +1,5 @@
 [
+  "0x0",
   "0x12",
   "0x0",
   "0x7fff7fff",

stwo_cairo_verifier/crates/cairo_air/src/claim.cairo

@@ -34,10 +34,11 @@ pub impl FlatClaimImpl of FlatClaimTrait {
         channel.mix_felts(pack_into_qm31s(array![self.component_enable_bits.len()].span()));
         channel.mix_felts(pack_into_qm31s(enable_bits_to_u32s(*self.component_enable_bits)));
         channel.mix_felts(pack_into_qm31s(*self.component_log_sizes));
-        channel
-            .mix_felts(
-                pack_into_qm31s(array![self.public_data.public_memory.program.len()].span()),
-            );
+        let program_len: u32 = match self.public_data.public_memory.program {
+            Some(program) => program.len(),
+            None => 0,
+        };
+        channel.mix_felts(pack_into_qm31s(array![program_len].span()));
         self.public_data.mix_into(ref channel);
     }
 }

stwo_cairo_verifier/crates/cairo_air/src/lib.cairo

@@ -159,9 +159,10 @@ pub struct VerificationOutput {
 pub fn get_verification_output(proof: @CairoProof) -> VerificationOutput {
     // Note: the blake hash yields a 256-bit integer, the given program hash is taken modulo the
     // f252 prime to yield a felt.
-    let program_hash = construct_f252(
-        encode_and_hash_program_memory_section(*proof.claim.public_data.public_memory.program),
-    );
+    let program_hash = match proof.claim.public_data.public_memory.program {
+        Some(program) => construct_f252(encode_and_hash_program_memory_section(*program)),
+        None => 0,
+    };
 
     let mut output = array![];
     for entry in proof.claim.public_data.public_memory.output {
@@ -176,9 +177,10 @@ pub fn get_verification_output(proof: @CairoProof) -> VerificationOutput {
 pub fn get_verification_output(proof: @CairoProof) -> VerificationOutput {
     // Note: the blake hash yields a 256-bit integer, the given program hash is taken modulo the
     // f252 prime to yield a felt.
-    let program_hash = construct_f252(
-        encode_and_hash_program_memory_section(*proof.claim.public_data.public_memory.program),
-    );
+    let program_hash = match proof.claim.public_data.public_memory.program {
+        Some(program) => construct_f252(encode_and_hash_program_memory_section(*program)),
+        None => 0,
+    };
 
     let output_hash = construct_f252(
         encode_and_hash_outputs_memory_section(*proof.claim.public_data.public_memory.output),
@@ -300,7 +302,9 @@ fn verify_claim(claim: @CairoClaim) {
         claim.ec_op_builtin,
         public_segments,
     );
-    verify_program(*program, public_segments);
+    if let Some(program) = program {
+        verify_program(*program, public_segments);
+    }
 
     let initial_pc: u32 = (*initial_pc).into();
     let initial_ap: u32 = (*initial_ap).into();
@@ -651,7 +655,7 @@ impl PublicSegmentRangesImpl of PublicSegmentRangesTrait {
 
 #[derive(Copy, Serde, Drop)]
 pub struct PublicMemory {
-    pub program: MemorySection,
+    pub program: Option<MemorySection>,
     pub public_segments: PublicSegmentRanges,
     pub output: MemorySection,
     pub safe_call_ids: [u32; 2],
@@ -666,7 +670,9 @@ pub impl PublicMemoryImpl of PublicMemoryTrait {
         let mut pub_memory_entries = PublicMemoryEntriesTrait::empty();
 
         // The program is loaded to `initial_pc`.
-        pub_memory_entries.add_memory_section(self.program, initial_pc);
+        if let Some(program) = self.program {
+            pub_memory_entries.add_memory_section(program, initial_pc);
+        }
         // Output was written to `final_ap`.
         pub_memory_entries.add_memory_section(self.output, final_ap);
 
@@ -721,7 +727,9 @@ pub impl PublicMemoryImpl of PublicMemoryTrait {
         let PublicMemory { program, public_segments, output, safe_call_ids } = self;
 
         // Mix program memory section.
-        channel.mix_memory_section(*program);
+        if let Some(program) = program {
+            channel.mix_memory_section(*program);
+        }
 
         // Mix public segments.
         public_segments.mix_into(ref channel);
@@ -896,8 +904,10 @@ impl PublicDataImpl of PublicDataTrait {
         for (id, _) in output {
             public_claim.append(*id);
         }
-        for (id, _) in program {
-            public_claim.append(*id);
+        if let Some(program) = program {
+            for (id, _) in program {
+                public_claim.append(*id);
+            }
         }
 
         // Collect output values.
@@ -911,11 +921,13 @@ impl PublicDataImpl of PublicDataTrait {
 
         // Collect program values.
         let mut program_claim = array![];
-        for (_, value) in program {
-            let fixed_arr: [u32; 8] = *value;
-            let new_value: [u32; N_M31_IN_FELT252] = split(fixed_arr);
-            let arr: Array<u32> = new_value.span().into_iter().map(|x| *x).collect();
-            program_claim.extend(arr);
+        if let Some(program) = program {
+            for (_, value) in program {
+                let fixed_arr: [u32; 8] = *value;
+                let new_value: [u32; N_M31_IN_FELT252] = split(fixed_arr);
+                let arr: Array<u32> = new_value.span().into_iter().map(|x| *x).collect();
+                program_claim.extend(arr);
+            }
         }
 
         (public_claim.span(), output_claim.span(), program_claim.span())

stwo_cairo_verifier/crates/cairo_air/src/test.cairo

@@ -35,7 +35,7 @@ fn test_public_data_logup_sum() {
 
     let public_data = PublicData {
         public_memory: PublicMemory {
-            program: program.span(),
+            program: Option::Some(program.span()),
             public_segments: PublicSegmentRanges {
                 output: SegmentRange {
                     start_ptr: MemorySmallValue { id: 228, value: 2520 },

stwo_cairo_verifier/crates/cairo_air/src/test_utils.cairo

@@ -31,7 +31,7 @@ pub fn mock_public_memory_with_outputs(output_len: u32) -> PublicMemory {
     };
 
     PublicMemory {
-        program: [].span(),
+        program: Option::Some([].span()),
         public_segments: PublicSegmentRanges {
             output: empty_segment,
             pedersen: empty_segment,