Trace to Prover Refactor

src/air/example/cairo.rs

@@ -162,13 +162,46 @@ pub struct PublicInputs {
     pub num_steps: usize, // number of execution steps
 }
 
+impl PublicInputs {
+    /// Creates a Public Input from register states and memory
+    /// - In the future we should use the output of the Cairo Runner. This is not currently supported in Cairo RS
+    /// - RangeChecks are not filled, and the prover mutates them inside the prove function. This works but also should be loaded from the Cairo RS output
+    pub fn from_regs_and_mem(
+        register_states: &CairoTrace,
+        memory: &CairoMemory,
+        program_size: usize,
+    ) -> Self {
+        let mut program = vec![];
+
+        for i in 1..=program_size as u64 {
+            program.push(memory.get(&i).unwrap().clone());
+        }
+
+        let last_step = &register_states.rows[register_states.steps() - 1];
+
+        PublicInputs {
+            pc_init: FE::from(register_states.rows[0].pc),
+            ap_init: FE::from(register_states.rows[0].ap),
+            fp_init: FE::from(register_states.rows[0].fp),
+            pc_final: FieldElement::from(last_step.pc),
+            ap_final: FieldElement::from(last_step.ap),
+            range_check_min: None,
+            range_check_max: None,
+            program,
+            num_steps: register_states.steps(),
+        }
+    }
+}
 #[derive(Clone)]
 pub struct CairoAIR {
     pub context: AirContext,
     pub number_steps: usize,
 }
 
 impl CairoAIR {
+    /// Creates a new CairoAIR from proof_options
+    /// full_trace_length: Padding to 2^n
+    /// number_steps: Number of steps of the execution / register steps / rows in cairo runner trace
     pub fn new(proof_options: ProofOptions, full_trace_length: usize, number_steps: usize) -> Self {
         let context = AirContext {
             options: proof_options,

src/cairo_run/cairo_layout.rs

@@ -0,0 +1,29 @@
+pub enum CairoLayout {
+    Plain,
+    Small,
+    Dex,
+    Recursive,
+    Starknet,
+    StarknetWithKeccak,
+    RecursiveLargeOutput,
+    AllCairo,
+    AllSolidity,
+    Dynamic,
+}
+
+impl CairoLayout {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            CairoLayout::Plain => "plain",
+            CairoLayout::Small => "small",
+            CairoLayout::Dex => "dex",
+            CairoLayout::Recursive => "recursive",
+            CairoLayout::Starknet => "starknet",
+            CairoLayout::StarknetWithKeccak => "starknet_with_keccak",
+            CairoLayout::RecursiveLargeOutput => "recursive_large_output",
+            CairoLayout::AllCairo => "all_cairo",
+            CairoLayout::AllSolidity => "all_solidity",
+            CairoLayout::Dynamic => "dynamic",
+        }
+    }
+}

src/cairo_run/mod.rs

@@ -1,2 +1,4 @@
+pub mod cairo_layout;
 pub mod file_writer;
 pub mod run;
+pub mod vec_writer;

src/cairo_run/run.rs

@@ -1,10 +1,13 @@
-use super::file_writer::FileWriter;
+use crate::cairo_vm::cairo_mem::CairoMemory;
+use crate::cairo_vm::cairo_trace::CairoTrace;
+
+use super::cairo_layout::CairoLayout;
+use super::vec_writer::VecWriter;
 use cairo_vm::cairo_run::{self, EncodeTraceError};
 use cairo_vm::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor;
 use cairo_vm::vm::errors::cairo_run_errors::CairoRunError;
 use cairo_vm::vm::errors::trace_errors::TraceError;
 use cairo_vm::vm::errors::vm_errors::VirtualMachineError;
-use std::io;
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -21,13 +24,13 @@ pub enum Error {
     Trace(#[from] TraceError),
 }
 
-/// Runs a cairo program in JSON format and returns trace and memory files.
+/// Runs a cairo program in JSON format and returns trace, memory and program length.
 /// Uses [cairo-rs](https://github.com/lambdaclass/cairo-rs/) project to run the program.
 ///
 ///  # Params
 ///
 /// `entrypoint_function` - the name of the entrypoint function tu run. If `None` is provided, the default value is `main`.
-/// `layout` - type of layout of Cairo.
+/// `layout` - type of layout of Cairo. ``
 /// `filename` - path to the input file.
 /// `trace_path` - path where to store the generated trace file.
 /// `memory_path` - path where to store the generated memory file.
@@ -38,11 +41,9 @@ pub enum Error {
 /// `Error` indicating the type of error.
 pub fn run_program(
     entrypoint_function: Option<&str>,
-    layout: &str,
+    layout: CairoLayout,
     filename: &str,
-    trace_path: &str,
-    memory_path: &str,
-) -> Result<(), Error> {
+) -> Result<(CairoTrace, CairoMemory, usize), Error> {
     // default value for entrypoint is "main"
     let entrypoint = entrypoint_function.unwrap_or("main");
 
@@ -52,7 +53,7 @@ pub fn run_program(
         entrypoint,
         trace_enabled,
         relocate_mem: true,
-        layout,
+        layout: layout.as_str(),
         proof_mode: false,
         secure_run: None,
     };
@@ -70,26 +71,32 @@ pub fn run_program(
 
     let relocated_trace = vm.get_relocated_trace()?;
 
-    let trace_file = std::fs::File::create(trace_path)?;
-    let mut trace_writer =
-        FileWriter::new(io::BufWriter::with_capacity(3 * 1024 * 1024, trace_file));
-
+    let mut trace_vec = Vec::<u8>::new();
+    let mut trace_writer = VecWriter::new(&mut trace_vec);
     cairo_run::write_encoded_trace(relocated_trace, &mut trace_writer)?;
-    trace_writer.flush()?;
-
-    let memory_file = std::fs::File::create(memory_path)?;
-    let mut memory_writer =
-        FileWriter::new(io::BufWriter::with_capacity(5 * 1024 * 1024, memory_file));
 
+    let mut memory_vec = Vec::<u8>::new();
+    let mut memory_writer = VecWriter::new(&mut memory_vec);
     cairo_run::write_encoded_memory(&cairo_runner.relocated_memory, &mut memory_writer)?;
+
+    trace_writer.flush()?;
     memory_writer.flush()?;
 
-    Ok(())
+    println!("From bytes: {:?}", memory_vec.len());
+
+    //TO DO: Better error handling
+    let cairo_mem = CairoMemory::from_bytes_le(&memory_vec).unwrap();
+    let cairo_trace = CairoTrace::from_bytes_le(&trace_vec).unwrap();
+
+    let data_len = cairo_runner.get_program().data_len();
+
+    Ok((cairo_trace, cairo_mem, data_len))
 }
 
 #[cfg(test)]
 mod tests {
     use crate::air::trace::TraceTable;
+    use crate::cairo_run::cairo_layout::CairoLayout;
     use crate::cairo_vm::cairo_mem::CairoMemory;
     use crate::cairo_vm::cairo_trace::CairoTrace;
     use crate::cairo_vm::execution_trace::build_cairo_execution_trace;
@@ -111,7 +118,7 @@ mod tests {
 
         println!("{}", json_filename);
 
-        super::run_program(None, "all_cairo", &json_filename, &dir_trace, &dir_memory).unwrap();
+        super::run_program(None, CairoLayout::AllCairo, &json_filename).unwrap();
 
         // read trace from file
         let raw_trace = CairoTrace::from_file(&dir_trace).unwrap();

src/cairo_run/vec_writer.rs

@@ -0,0 +1,26 @@
+use bincode::enc::write::Writer;
+use std::io::{self, Write};
+
+pub struct VecWriter<'a> {
+    buf_writer: &'a mut Vec<u8>,
+}
+
+impl Writer for VecWriter<'_> {
+    fn write(&mut self, bytes: &[u8]) -> Result<(), bincode::error::EncodeError> {
+        self.buf_writer
+            .write_all(bytes)
+            .expect("Shouldn't fail in memory vector");
+
+        Ok(())
+    }
+}
+
+impl<'a> VecWriter<'a> {
+    pub fn new(vec: &'a mut Vec<u8>) -> Self {
+        Self { buf_writer: vec }
+    }
+
+    pub fn flush(&mut self) -> io::Result<()> {
+        self.buf_writer.flush()
+    }
+}

src/cairo_vm/cairo_mem.rs

@@ -22,6 +22,14 @@ impl CairoMemory {
         self.data.get(addr)
     }
 
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.data.is_empty()
+    }
+
     pub fn from_bytes_le(bytes: &[u8]) -> Result<Self, CairoImportError> {
         // Each row is an 8 bytes address
         // and a value of 32 bytes (which is a field)