support builtins in the adapter

stwo_cairo_prover/crates/adapted_prover/src/main.rs

@@ -65,8 +65,9 @@ fn run(args: impl Iterator<Item = String>) -> Result<CairoProof<Blake2sMerkleHas
     let _span = span!(Level::INFO, "run").entered();
     let args = Args::try_parse_from(args)?;
 
+    // dev_mode = true
     let vm_output: CairoInput =
-        import_from_vm_output(args.pub_json.as_path(), args.priv_json.as_path())?;
+        import_from_vm_output(args.pub_json.as_path(), args.priv_json.as_path(), true)?;
 
     let casm_states_by_opcode_count = &vm_output.state_transitions.casm_states_by_opcode.counts();
     log::info!("Casm states by opcode count: {casm_states_by_opcode_count:?}");

stwo_cairo_prover/crates/prover/src/input/builtins_segments.rs

@@ -0,0 +1,50 @@
+use std::collections::HashMap;
+
+use cairo_vm::air_public_input::MemorySegmentAddresses;
+
+// TODO(Stav): Understand if the adapter should pass builtins that won't be supported by stwo.
+/// This struct holds the builtins used in a Cairo program.
+/// Most of them are not implemented yet by Stwo.
+#[derive(Debug, Default)]
+pub struct BuiltinsSegments {
+    pub range_check_bits_128_builtin: Option<MemorySegmentAddresses>,
+    pub range_check_bits_96_builtin: Option<MemorySegmentAddresses>,
+    pub bitwise_builtin: Option<MemorySegmentAddresses>,
+    pub add_mod_builtin: Option<MemorySegmentAddresses>,
+    pub ec_op_builtin: Option<MemorySegmentAddresses>,
+    pub ecdsa_builtin: Option<MemorySegmentAddresses>,
+    pub keccak_builtin: Option<MemorySegmentAddresses>,
+    pub mul_mod_builtin: Option<MemorySegmentAddresses>,
+    pub pedersen_builtin: Option<MemorySegmentAddresses>,
+    pub poseidon_builtin: Option<MemorySegmentAddresses>,
+}
+
+impl BuiltinsSegments {
+    /// Create a new `BuiltinsSegments` struct from a map of memory MemorySegmentAddressess.
+    pub fn from_memory_segments(
+        memory_segments: &HashMap<&str, cairo_vm::air_public_input::MemorySegmentAddresses>,
+    ) -> Self {
+        let mut res = BuiltinsSegments::default();
+        for (name, value) in memory_segments.iter() {
+            let value = (value.begin_addr, value.stop_ptr).into();
+            match *name {
+                "range_check" => res.range_check_bits_128_builtin = Some(value),
+                "range_check96" => res.range_check_bits_96_builtin = Some(value),
+                "bitwise" => res.bitwise_builtin = Some(value),
+                "add_mod" => res.add_mod_builtin = Some(value),
+                "ec_op" => res.ec_op_builtin = Some(value),
+                "ecdsa" => res.ecdsa_builtin = Some(value),
+                "keccak" => res.keccak_builtin = Some(value),
+                "mul_mod" => res.mul_mod_builtin = Some(value),
+                "pedersen" => res.pedersen_builtin = Some(value),
+                "poseidon" => res.poseidon_builtin = Some(value),
+                // Output, execution and program segments are not builtins.
+                "output" => {}
+                "execution" => {}
+                "program" => {}
+                _ => panic!("Unknown memory segment: {name}"),
+            }
+        }
+        res
+    }
+}

stwo_cairo_prover/crates/prover/src/input/mod.rs

@@ -1,7 +1,8 @@
-use cairo_vm::air_public_input::MemorySegmentAddresses;
+use builtins_segments::BuiltinsSegments;
 use mem::Memory;
 use state_transitions::StateTransitions;
 
+pub mod builtins_segments;
 mod decode;
 pub mod mem;
 pub mod plain;
@@ -19,5 +20,5 @@ pub struct CairoInput {
     pub public_mem_addresses: Vec<u32>,
 
     // Builtins.
-    pub range_check_builtin: MemorySegmentAddresses,
+    pub builtins_segments: BuiltinsSegments,
 }

stwo_cairo_prover/crates/prover/src/input/plain.rs

@@ -9,8 +9,7 @@ use itertools::Itertools;
 use super::mem::{MemConfig, MemoryBuilder};
 use super::state_transitions::StateTransitions;
 use super::vm_import::MemEntry;
-use super::CairoInput;
-use crate::input::MemorySegmentAddresses;
+use super::{BuiltinsSegments, CairoInput};
 
 // TODO(Ohad): remove dev_mode after adding the rest of the opcodes.
 /// Translates a plain casm into a CairoInput by running the program and extracting the memory and
@@ -71,6 +70,13 @@ pub fn input_from_finished_runner(runner: CairoRunner, dev_mode: bool) -> CairoI
                 val: bytemuck::cast(v.to_bytes_le()),
             })
         });
+
+    let memory_segments = &runner
+        .get_air_public_input()
+        .expect("Unable to get public input from the runner")
+        .memory_segments;
+    let builtins_segments = BuiltinsSegments::from_memory_segments(memory_segments);
+
     let trace = runner.relocated_trace.unwrap();
     let trace = trace.iter().map(|t| t.clone().into());
 
@@ -84,9 +90,6 @@ pub fn input_from_finished_runner(runner: CairoRunner, dev_mode: bool) -> CairoI
         state_transitions,
         mem: mem.build(),
         public_mem_addresses,
-        range_check_builtin: MemorySegmentAddresses {
-            begin_addr: 24,
-            stop_ptr: 64,
-        },
+        builtins_segments,
     }
 }

stwo_cairo_prover/crates/prover/src/input/vm_import/mod.rs

@@ -4,6 +4,8 @@ use std::io::Read;
 use std::path::Path;
 
 use bytemuck::{bytes_of_mut, Pod, Zeroable};
+#[cfg(test)]
+use cairo_vm::air_public_input::MemorySegmentAddresses;
 use cairo_vm::air_public_input::PublicInput;
 use cairo_vm::vm::trace::trace_entry::RelocatedTraceEntry;
 use json::PrivateInput;
@@ -14,7 +16,7 @@ use super::mem::MemConfig;
 use super::state_transitions::StateTransitions;
 use super::CairoInput;
 use crate::input::mem::MemoryBuilder;
-use crate::input::MemorySegmentAddresses;
+use crate::input::BuiltinsSegments;
 
 #[derive(Debug, Error)]
 pub enum VmImportError {
@@ -26,9 +28,11 @@ pub enum VmImportError {
     NoMemorySegments,
 }
 
+// TODO(Ohad): remove dev_mode after adding the rest of the opcodes.
 pub fn import_from_vm_output(
     pub_json: &Path,
     priv_json: &Path,
+    dev_mode: bool,
 ) -> Result<CairoInput, VmImportError> {
     let _span = span!(Level::INFO, "import_from_vm_output").entered();
     let pub_data_string = std::fs::read_to_string(pub_json)?;
@@ -50,22 +54,22 @@ pub fn import_from_vm_output(
     let mut trace_file = std::io::BufReader::new(std::fs::File::open(trace_path)?);
     let mut mem_file = std::io::BufReader::new(std::fs::File::open(mem_path)?);
     let mut mem = MemoryBuilder::from_iter(mem_config, MemEntryIter(&mut mem_file));
-    let state_transitions = StateTransitions::from_iter(TraceIter(&mut trace_file), &mut mem, true);
+    let state_transitions =
+        StateTransitions::from_iter(TraceIter(&mut trace_file), &mut mem, dev_mode);
 
     let public_mem_addresses = pub_data
         .public_memory
         .iter()
         .map(|entry| entry.address as u32)
         .collect();
 
+    let builtins_segments = BuiltinsSegments::from_memory_segments(&pub_data.memory_segments);
+
     Ok(CairoInput {
         state_transitions,
         mem: mem.build(),
         public_mem_addresses,
-        range_check_builtin: MemorySegmentAddresses {
-            begin_addr: pub_data.memory_segments["range_check"].begin_addr as usize,
-            stop_ptr: pub_data.memory_segments["range_check"].stop_ptr as usize,
-        },
+        builtins_segments,
     })
 }
 
@@ -135,7 +139,12 @@ pub mod tests {
         let mut d = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
         d.push("test_data/large_cairo_input");
 
-        import_from_vm_output(d.join("pub.json").as_path(), d.join("priv.json").as_path()).expect(
+        import_from_vm_output(
+            d.join("pub.json").as_path(),
+            d.join("priv.json").as_path(),
+            false,
+        )
+        .expect(
             "
             Failed to read test files. Maybe git-lfs is not installed? Checkout README.md.",
         )
@@ -144,19 +153,24 @@ pub mod tests {
     pub fn small_cairo_input() -> CairoInput {
         let mut d = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
         d.push("test_data/small_cairo_input");
-        import_from_vm_output(d.join("pub.json").as_path(), d.join("priv.json").as_path()).expect(
+        import_from_vm_output(
+            d.join("pub.json").as_path(),
+            d.join("priv.json").as_path(),
+            false,
+        )
+        .expect(
             "
             Failed to read test files. Maybe git-lfs is not installed? Checkout README.md.",
         )
     }
 
     // TODO (Stav): Once all the components are in, verify the proof to ensure the sort was correct.
-    // TODO (Ohad): remove the following doc after deleting dev_mod.
-    /// When not ignored, the test passes only with dev_mod = false.
     #[ignore]
     #[test]
     fn test_read_from_large_files() {
         let input = large_cairo_input();
+
+        // Test opcode components.
         let components = input.state_transitions.casm_states_by_opcode;
         assert_eq!(components.generic_opcode.len(), 0);
         assert_eq!(components.add_ap_opcode_is_imm_f_op1_base_fp_f.len(), 0);
@@ -214,13 +228,87 @@ pub mod tests {
         assert_eq!(components.mul_opcode_is_small_f_is_imm_f.len(), 4583);
         assert_eq!(components.mul_opcode_is_small_f_is_imm_t.len(), 9047);
         assert_eq!(components.ret_opcode.len(), 49472);
+
+        // Test builtins.
+        let builtins_segments = input.builtins_segments;
+        assert_eq!(
+            builtins_segments.range_check_bits_128_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 1715768,
+                stop_ptr: 1757348
+            })
+        );
+        assert_eq!(
+            builtins_segments.range_check_bits_96_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 17706552,
+                stop_ptr: 17706552
+            })
+        );
+        assert_eq!(
+            builtins_segments.bitwise_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 5942840,
+                stop_ptr: 5942840
+            })
+        );
+        assert_eq!(
+            builtins_segments.add_mod_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 21900856,
+                stop_ptr: 21900856
+            })
+        );
+        assert_eq!(
+            builtins_segments.ec_op_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 16428600,
+                stop_ptr: 16428747
+            })
+        );
+        assert_eq!(
+            builtins_segments.ecdsa_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 5910072,
+                stop_ptr: 5910072
+            })
+        );
+        assert_eq!(
+            builtins_segments.keccak_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 16657976,
+                stop_ptr: 16657976
+            })
+        );
+        assert_eq!(
+            builtins_segments.mul_mod_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 23735864,
+                stop_ptr: 23735864
+            })
+        );
+        assert_eq!(
+            builtins_segments.pedersen_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 1322552,
+                stop_ptr: 1337489
+            })
+        );
+        assert_eq!(
+            builtins_segments.poseidon_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 16920120,
+                stop_ptr: 17444532
+            })
+        );
     }
 
-    // When not ignored, the test passes only with dev_mod = false.
     #[ignore]
     #[test]
     fn test_read_from_small_files() {
         let input = small_cairo_input();
+
+        // Test opcode components.
         let components = input.state_transitions.casm_states_by_opcode;
         assert_eq!(components.generic_opcode.len(), 0);
         assert_eq!(components.add_ap_opcode_is_imm_f_op1_base_fp_f.len(), 0);
@@ -278,5 +366,78 @@ pub mod tests {
         assert_eq!(components.mul_opcode_is_small_f_is_imm_f.len(), 0);
         assert_eq!(components.mul_opcode_is_small_f_is_imm_t.len(), 0);
         assert_eq!(components.ret_opcode.len(), 462);
+
+        // Test builtins.
+        let builtins_segments = input.builtins_segments;
+        assert_eq!(
+            builtins_segments.range_check_bits_128_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 6000,
+                stop_ptr: 6050
+            })
+        );
+        assert_eq!(
+            builtins_segments.range_check_bits_96_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 68464,
+                stop_ptr: 68514
+            })
+        );
+        assert_eq!(
+            builtins_segments.bitwise_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 22512,
+                stop_ptr: 22762
+            })
+        );
+        assert_eq!(
+            builtins_segments.add_mod_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 84848,
+                stop_ptr: 84848
+            })
+        );
+        assert_eq!(
+            builtins_segments.ec_op_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 63472,
+                stop_ptr: 63822
+            })
+        );
+        assert_eq!(
+            builtins_segments.ecdsa_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 22384,
+                stop_ptr: 22484
+            })
+        );
+        assert_eq!(
+            builtins_segments.keccak_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 64368,
+                stop_ptr: 65168
+            })
+        );
+        assert_eq!(
+            builtins_segments.mul_mod_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 92016,
+                stop_ptr: 92016
+            })
+        );
+        assert_eq!(
+            builtins_segments.pedersen_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 4464,
+                stop_ptr: 4614
+            })
+        );
+        assert_eq!(
+            builtins_segments.poseidon_builtin,
+            Some(MemorySegmentAddresses {
+                begin_addr: 65392,
+                stop_ptr: 65692
+            })
+        );
     }
 }