initial commit

crates/prover/src/input/mem.rs

@@ -118,7 +118,7 @@ impl MemoryBuilder {
     pub fn get_inst(&mut self, addr: u32) -> u64 {
         let mut inst_cache = std::mem::take(&mut self.inst_cache);
         let res = *inst_cache.entry(addr).or_insert_with(|| {
-            let inst: M31 = self.mem.get(addr).0.try_into().unwrap();
+            let inst: M31 = self.mem.get(addr).0 .0 .0;
             inst.0 as u64
         });
         self.inst_cache = inst_cache;

crates/runnair/Cargo.toml

@@ -16,6 +16,7 @@ sonic-rs.workspace = true
 stwo-prover.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
+stwo_cairo_prover = { path = "../prover" }
 
 [dev-dependencies]
 cairo-lang-casm.workspace = true

crates/runnair/src/main.rs

@@ -1,9 +1,10 @@
-use vm::run_fibonacci;
+use vm::{prove_fibonacci, run_fibonacci};
 
 pub mod memory;
 pub mod utils;
 pub mod vm;
 
 fn main() {
     run_fibonacci();
+    let _ = prove_fibonacci();
 }

crates/runnair/src/memory/mod.rs

@@ -20,9 +20,9 @@ const MAX_MEMORY_SIZE_BITS: u8 = 30;
 pub struct Memory {
     // TODO(alont) Consdier changing the implementation to segment -> (offset -> value) for memory
     // locality.
-    relocatable_data: Vec<Vec<Option<MaybeRelocatable<QM31>>>>,
+    pub relocatable_data: Vec<Vec<Option<MaybeRelocatable<QM31>>>>,
     // TODO: convert to a vector.
-    absolute_data: HashMap<M31, MaybeRelocatableValue>,
+    pub absolute_data: HashMap<M31, MaybeRelocatableValue>,
 }
 
 impl<T: Into<MaybeRelocatableAddr>> Index<T> for Memory {
@@ -76,6 +76,23 @@ impl Memory {
         self.relocatable_data.clear();
     }
 
+    pub fn relocation_table(&self) -> RelocationTable {
+        self.relocatable_data
+            .iter()
+            .map(|segment| segment.len() as u32)
+            .enumerate()
+            .scan(1u32, |cumulative_sum_state, (i, current_len)| {
+                let first_address_at_segment = cumulative_sum_state.clone();
+                *cumulative_sum_state += current_len;
+                Some((i, M31(first_address_at_segment)))
+            })
+            .collect()
+    }
+
+    pub fn relocatable_data(&self) -> &Vec<Vec<Option<MaybeRelocatable<QM31>>>> {
+        &self.relocatable_data
+    }
+
     pub fn insert<T: Into<MaybeRelocatableAddr>, S: Into<MaybeRelocatableValue>>(
         &mut self,
         key: T,

crates/runnair/src/vm/mod.rs

@@ -11,9 +11,15 @@ use std::fs::File;
 use std::io::BufReader;
 use std::path::PathBuf;
 
+use itertools::Itertools;
 use num_traits::Zero;
 use serde::Deserialize;
 use serde_json;
+use stwo_cairo_prover::cairo_air::CairoProof;
+use stwo_cairo_prover::input::instructions::Instructions;
+use stwo_cairo_prover::input::mem::{MemConfig, MemoryBuilder};
+use stwo_cairo_prover::input::vm_import::{MemEntry, TraceEntry};
+use stwo_cairo_prover::input::CairoInput;
 use stwo_prover::core::fields::m31::M31;
 use stwo_prover::core::fields::qm31::QM31;
 
@@ -25,7 +31,7 @@ use self::hints::*;
 use self::jmp::*;
 use self::jnz::*;
 use self::range_check::*;
-use crate::memory::relocatable::{MaybeRelocatable, Relocatable, Segment};
+use crate::memory::relocatable::{MaybeRelocatable, Relocatable, RelocationTable, Segment};
 use crate::memory::{MaybeRelocatableAddr, Memory};
 use crate::utils::{get_tests_data_dir, m31_from_hex_str, maybe_resize, u32_from_usize};
 
@@ -39,6 +45,13 @@ pub struct State {
     pc: MaybeRelocatableAddr,
 }
 
+#[derive(Clone, Copy, Debug)]
+pub struct RelocatedState {
+    ap: M31,
+    fp: M31,
+    pc: M31,
+}
+
 impl State {
     pub fn advance(self) -> Self {
         Self {
@@ -144,7 +157,7 @@ impl Program {
 #[derive(Debug)]
 pub struct VM {
     memory: Memory,
-    state: State,
+    trace: Vec<State>,
     hint_runner: HintRunner,
 }
 
@@ -153,8 +166,16 @@ impl VM {
         &self.memory
     }
 
+    pub fn memory_mut(&mut self) -> &mut Memory {
+        &mut self.memory
+    }
+
     pub fn state(&self) -> &State {
-        &self.state
+        self.trace.last().unwrap()
+    }
+
+    pub fn trace(&self) -> &Vec<State> {
+        &self.trace
     }
 }
 
@@ -220,40 +241,55 @@ impl VM {
 
         Self {
             memory,
-            state,
+            trace: vec![state],
             hint_runner,
         }
     }
 
     fn step(&mut self) {
+        let current_state = self.state().clone();
         self.hint_runner
-            .maybe_execute_hint(&mut self.memory, &self.state);
+            .maybe_execute_hint(&mut self.memory, &current_state);
         self.execute_instruction();
     }
 
     fn execute_instruction(&mut self) {
-        let MaybeRelocatable::Absolute(instruction) = self.memory[self.state.pc] else {
+        let MaybeRelocatable::Absolute(instruction) = self.memory[self.state().pc] else {
             panic!("Instruction must be an absolute value.");
         };
         let Instruction { op, args } = instruction.into();
         let instruction_fn = opcode_to_instruction(op);
 
-        self.state = instruction_fn(&mut self.memory, self.state, args);
+        let current_state = self.state().clone();
+        self.trace
+            .push(instruction_fn(&mut self.memory, current_state, args));
     }
 
     pub fn execute(&mut self) {
         let [final_fp, final_pc] =
             [Self::FINAL_FP, Self::FINAL_PC].map(|x| MaybeRelocatableAddr::Relocatable(x.into()));
 
-        while self.state.pc != final_pc {
+        while self.state().pc != final_pc {
             self.step();
         }
 
         assert_eq!(
-            self.state.fp, final_fp,
+            self.state().fp,
+            final_fp,
             "Only final `fp` is allowed when at final `pc`."
         );
     }
+
+    pub fn relocated_trace(&self, relocation_table: &RelocationTable) -> Vec<RelocatedState> {
+        self.trace
+            .iter()
+            .map(|state| RelocatedState {
+                ap: state.ap.relocate(relocation_table),
+                fp: state.fp.relocate(relocation_table),
+                pc: state.pc.relocate(relocation_table),
+            })
+            .collect()
+    }
 }
 
 // Utils.
@@ -472,12 +508,60 @@ pub(crate) fn run_fibonacci() {
     vm.execute();
 }
 
+pub(crate) fn prove_fibonacci() -> CairoProof {
+    let program_path = get_tests_data_dir().join("fibonacci_compiled.json");
+    let program = Program::from_compiled_file(program_path);
+    let public_mem_addresses = (0..(program.instructions.len() as u32)).collect_vec();
+    let input = serde_json::json!({ "fibonacci_claim_index": ["0x64", "0x0", "0x0", "0x0"]});
+    let mut vm = VM::create_for_main_entry_point(program, input);
+
+    vm.execute();
+    let mem = vm.memory_mut();
+    let relocation_table = mem.relocation_table();
+    mem.relocate(&relocation_table);
+    let relocated_trace = vm.relocated_trace(&relocation_table);
+
+    let trace = relocated_trace.iter().map(|t| TraceEntry {
+        ap: t.ap.0 as u64,
+        fp: t.fp.0 as u64,
+        pc: t.pc.0 as u64,
+    });
+
+    let mem_config = MemConfig::default();
+    let mut mem = MemoryBuilder::from_iter(
+        mem_config,
+        vm.memory()
+            .absolute_data
+            .iter()
+            .map(|(addr, value)| MemEntry {
+                addr: addr.0,
+                val: match value {
+                    MaybeRelocatable::Absolute(qm31) => qm31.to_m31_array().map(|m31| m31.0),
+                    MaybeRelocatable::Relocatable(_) => panic!("memory was relocated"),
+                },
+            }),
+    );
+    let instructions = Instructions::from_iter(trace, &mut mem);
+
+    let input = CairoInput {
+        instructions,
+        memory: mem.build(),
+        public_mem_addresses,
+    };
+
+    stwo_cairo_prover::cairo_air::prove_cairo(input).unwrap()
+}
+
 #[cfg(test)]
 mod test {
-    use crate::vm::run_fibonacci;
-
+    use crate::vm::{prove_fibonacci, run_fibonacci};
     #[test]
     fn test_runner() {
         run_fibonacci()
     }
+
+    #[test]
+    fn test_prover() {
+        prove_fibonacci();
+    }
 }