support builtins in the adapter

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

@@ -0,0 +1,98 @@
+use std::collections::HashMap;
+
+use cairo_vm::air_public_input::MemorySegmentAddresses;
+use cairo_vm::types::builtin_name::BuiltinName;
+
+/// This struct holds the builtins used in a Cairo program.
+#[derive(Debug, Default)]
+pub struct BuiltinSegments {
+    pub add_mod: Option<MemorySegmentAddresses>,
+    pub bitwise: Option<MemorySegmentAddresses>,
+    pub ec_op: Option<MemorySegmentAddresses>,
+    pub ecdsa: Option<MemorySegmentAddresses>,
+    pub keccak: Option<MemorySegmentAddresses>,
+    pub mul_mod: Option<MemorySegmentAddresses>,
+    pub pedersen: Option<MemorySegmentAddresses>,
+    pub poseidon: Option<MemorySegmentAddresses>,
+    pub range_check_bits_96: Option<MemorySegmentAddresses>,
+    pub range_check_bits_128: Option<MemorySegmentAddresses>,
+}
+
+impl BuiltinSegments {
+    pub fn add_segment(
+        &mut self,
+        builtin_name: BuiltinName,
+        segment: Option<MemorySegmentAddresses>,
+    ) {
+        match builtin_name {
+            BuiltinName::range_check => self.range_check_bits_128 = segment,
+            BuiltinName::pedersen => self.pedersen = segment,
+            BuiltinName::ecdsa => self.ecdsa = segment,
+            BuiltinName::keccak => self.keccak = segment,
+            BuiltinName::bitwise => self.bitwise = segment,
+            BuiltinName::ec_op => self.ec_op = segment,
+            BuiltinName::poseidon => self.poseidon = segment,
+            BuiltinName::range_check96 => self.range_check_bits_96 = segment,
+            BuiltinName::add_mod => self.add_mod = segment,
+            BuiltinName::mul_mod => self.mul_mod = segment,
+            // Not builtins.
+            BuiltinName::output | BuiltinName::segment_arena => {}
+        }
+    }
+
+    /// Creates a new `BuiltinSegments` struct from a map of memory segment names to addresses.
+    pub fn from_memory_segments(memory_segments: &HashMap<&str, MemorySegmentAddresses>) -> Self {
+        let mut res = BuiltinSegments::default();
+        for (name, value) in memory_segments.iter() {
+            if let Some(builtin_name) = BuiltinName::from_str(name) {
+                // Filter empty segments.
+                let segment = if value.begin_addr == value.stop_ptr {
+                    None
+                } else {
+                    assert!(
+                        value.begin_addr < value.stop_ptr,
+                        "Invalid segment addresses for builtin: {}",
+                        name
+                    );
+                    Some((value.begin_addr, value.stop_ptr).into())
+                };
+                res.add_segment(builtin_name, segment);
+            };
+        }
+        res
+    }
+}
+
+#[cfg(test)]
+mod test_builtin_segments {
+    use std::path::PathBuf;
+
+    use cairo_vm::air_public_input::PublicInput;
+
+    use crate::input::BuiltinSegments;
+
+    #[test]
+    fn test_builtin_segments() {
+        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+            .join("src/input/test_builtins_segments/test_public_input.json");
+        let pub_data_string = std::fs::read_to_string(&path)
+            .unwrap_or_else(|_| panic!("Unable to read file: {}", path.display()));
+        let pub_data: PublicInput<'_> =
+            sonic_rs::from_str(&pub_data_string).expect("Unable to parse JSON");
+
+        let builtins_segments = BuiltinSegments::from_memory_segments(&pub_data.memory_segments);
+        assert_eq!(builtins_segments.add_mod, None);
+        assert_eq!(builtins_segments.bitwise, Some((23581, 23901).into()));
+        assert_eq!(builtins_segments.ec_op, None);
+        assert_eq!(builtins_segments.ecdsa, None);
+        assert_eq!(builtins_segments.keccak, None);
+        assert_eq!(builtins_segments.mul_mod, None);
+        assert_eq!(builtins_segments.pedersen, None);
+        assert_eq!(builtins_segments.poseidon, None);
+        assert_eq!(builtins_segments.range_check_bits_96, None);
+        assert_eq!(
+            builtins_segments.range_check_bits_128,
+            Some((7069, 7187).into())
+        );
+    }
+}

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

@@ -1,7 +1,8 @@
-use cairo_vm::air_public_input::MemorySegmentAddresses;
+use builtin_segments::BuiltinSegments;
 use memory::Memory;
 use state_transitions::StateTransitions;
 
+pub mod builtin_segments;
 mod decode;
 pub mod memory;
 pub mod plain;
@@ -17,7 +18,5 @@ pub struct CairoInput {
     pub state_transitions: StateTransitions,
     pub memory: Memory,
     pub public_memory_addresses: Vec<u32>,
-
-    // Builtins.
-    pub range_check_builtin: MemorySegmentAddresses,
+    pub builtins_segments: BuiltinSegments,
 }

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

@@ -9,8 +9,7 @@ use itertools::Itertools;
 use super::memory::{MemoryBuilder, MemoryConfig};
 use super::state_transitions::StateTransitions;
 use super::vm_import::MemoryEntry;
-use super::CairoInput;
-use crate::input::MemorySegmentAddresses;
+use super::{BuiltinSegments, 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
@@ -72,6 +71,13 @@ pub fn input_from_finished_runner(runner: CairoRunner, dev_mode: bool) -> CairoI
                 value: 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 = BuiltinSegments::from_memory_segments(memory_segments);
+
     let trace = runner.relocated_trace.unwrap();
     let trace = trace.iter().map(|t| t.clone().into());
 
@@ -85,9 +91,6 @@ pub fn input_from_finished_runner(runner: CairoRunner, dev_mode: bool) -> CairoI
         state_transitions,
         memory: memory.build(),
         public_memory_addresses,
-        range_check_builtin: MemorySegmentAddresses {
-            begin_addr: 24,
-            stop_ptr: 64,
-        },
+        builtins_segments,
     }
 }