continuations: Remove the need for dummy segments

evm_arithmetization/src/fixed_recursive_verifier.rs

@@ -49,7 +49,10 @@ use crate::proof::{
     FinalPublicValues, MemCapTarget, PublicValues, PublicValuesTarget, RegistersDataTarget,
     TrieRoots, TrieRootsTarget, TARGET_HASH_SIZE,
 };
-use crate::prover::{check_abort_signal, generate_all_data_segments, prove, GenerationSegmentData};
+use crate::prover::{
+    check_abort_signal, generate_all_data_segments, prove, GenerationSegmentData,
+    SegmentDataIterator,
+};
 use crate::recursive_verifier::{
     add_common_recursion_gates, add_virtual_public_values, get_memory_extra_looking_sum_circuit,
     recursive_stark_circuit, set_public_value_targets, PlonkWrapperCircuit, PublicInputs,
@@ -1560,15 +1563,21 @@ where
         timing: &mut TimingTree,
         abort_signal: Option<Arc<AtomicBool>>,
     ) -> anyhow::Result<Vec<ProverOutputData<F, C, D>>> {
-        let mut all_data_segments =
-            generate_all_data_segments::<F>(Some(max_cpu_len_log), &generation_inputs)?;
-        let mut proofs = Vec::with_capacity(all_data_segments.len());
-        for mut data in all_data_segments {
+        let mut it_segment_data = SegmentDataIterator {
+            inputs: generation_inputs.clone(),
+            partial_next_data: None,
+            max_cpu_len_log: Some(max_cpu_len_log),
+            nb_segments: 0,
+        };
+
+        let mut proofs = vec![];
+
+        for mut next_data in it_segment_data {
             let proof = self.prove_segment(
                 all_stark,
                 config,
                 generation_inputs.clone(),
-                &mut data,
+                &mut next_data.1,
                 timing,
                 abort_signal.clone(),
             )?;

evm_arithmetization/src/prover.rs

@@ -5,6 +5,7 @@ use anyhow::{anyhow, Result};
 use itertools::Itertools;
 use once_cell::sync::Lazy;
 use plonky2::field::extension::Extendable;
+use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::field::types::Field;
 use plonky2::fri::oracle::PolynomialBatch;
@@ -25,7 +26,7 @@ use starky::stark::Stark;
 use crate::all_stark::{AllStark, Table, NUM_TABLES};
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::interpreter::{set_registers_and_run, ExtraSegmentData, Interpreter};
-use crate::generation::state::GenerationState;
+use crate::generation::state::{GenerationState, State};
 use crate::generation::{generate_traces, GenerationInputs};
 use crate::get_challenges::observe_public_values;
 use crate::memory::segments::Segment;
@@ -522,6 +523,92 @@ fn build_segment_data<F: RichField>(
     }
 }
 
+pub struct SegmentDataIterator {
+    pub partial_next_data: Option<GenerationSegmentData>,
+    pub inputs: GenerationInputs,
+    pub max_cpu_len_log: Option<usize>,
+    pub nb_segments: usize,
+}
+
+type F = GoldilocksField;
+impl Iterator for SegmentDataIterator {
+    type Item = (GenerationInputs, GenerationSegmentData);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let cur_and_next_data = generate_next_segment::<F>(
+            self.max_cpu_len_log,
+            &self.inputs,
+            self.partial_next_data.clone(),
+        );
+
+        if cur_and_next_data.is_some() {
+            let (data, next_data) = cur_and_next_data.expect(
+                "Data cannot be
+`None`",
+            );
+            self.nb_segments += 1;
+            self.partial_next_data = next_data;
+            Some((self.inputs.clone(), data))
+        } else {
+            None
+        }
+    }
+}
+
+/// Returns the data for the current segment, as well as the data -- except
+/// registers_after -- for the next segment.
+pub(crate) fn generate_next_segment<F: RichField>(
+    max_cpu_len_log: Option<usize>,
+    inputs: &GenerationInputs,
+    partial_segment_data: Option<GenerationSegmentData>,
+) -> Option<(GenerationSegmentData, Option<GenerationSegmentData>)> {
+    let mut interpreter = Interpreter::<F>::new_with_generation_inputs(
+        KERNEL.global_labels["init"],
+        vec![],
+        inputs,
+        max_cpu_len_log,
+    );
+
+    // Get the (partial) current segment data, if it is provided. Otherwise,
+    // initialize it.
+    let mut segment_data = if let Some(partial) = partial_segment_data {
+        if partial.registers_after.program_counter == KERNEL.global_labels["halt"] {
+            return None;
+        }
+        interpreter
+            .get_mut_generation_state()
+            .set_segment_data(&partial);
+        interpreter.generation_state.memory = partial.memory.clone();
+        partial
+    } else {
+        build_segment_data(0, None, None, None, &interpreter)
+    };
+
+    let segment_index = segment_data.segment_index;
+
+    // Run the interpreter to get `registers_after` and the partial data for the
+    // next segment.
+    let run_result = set_registers_and_run(segment_data.registers_after, &mut interpreter);
+
+    if let Ok((updated_registers, mem_after)) = run_result {
+        // Set `registers_after` correctly and push the data.
+        let before_registers = segment_data.registers_after;
+
+        let partial_segment_data = Some(build_segment_data(
+            segment_index + 1,
+            Some(updated_registers),
+            Some(updated_registers),
+            mem_after,
+            &interpreter,
+        ));
+
+        segment_data.registers_after = updated_registers;
+        Some((segment_data, partial_segment_data))
+    } else {
+        None
+    }
+}
+
 /// Returns a vector containing the data required to generate all the segments
 /// of a transaction.
 pub fn generate_all_data_segments<F: RichField>(