support builtins in the adapter

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

@@ -0,0 +1,145 @@
+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)]
+pub struct BuiltinsSegments {
+    pub range_check_bits_128_builtin: MemorySegmentAddresses,
+    pub range_check_bits_96_builtin: MemorySegmentAddresses,
+    pub bitwise_builtin: MemorySegmentAddresses,
+    pub add_mod_builtin: MemorySegmentAddresses,
+    pub ec_op_builtin: MemorySegmentAddresses,
+    pub ecdsa_builtin: MemorySegmentAddresses,
+    pub keccak_builtin: MemorySegmentAddresses,
+    pub mul_mod_builtin: MemorySegmentAddresses,
+    pub pedersen_builtin: MemorySegmentAddresses,
+    pub poseidon_builtin: MemorySegmentAddresses,
+}
+
+impl Default for BuiltinsSegments {
+    fn default() -> Self {
+        Self {
+            range_check_bits_128_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            range_check_bits_96_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            bitwise_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            add_mod_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            ec_op_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            ecdsa_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            keccak_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            mul_mod_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            pedersen_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+            poseidon_builtin: MemorySegmentAddresses {
+                begin_addr: 0,
+                stop_ptr: 0,
+            },
+        }
+    }
+}
+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 = Self::default();
+        for (name, value) in memory_segments.iter() {
+            match *name {
+                "range_check" => {
+                    res.range_check_bits_128_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "range_check96" => {
+                    res.range_check_bits_96_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "bitwise" => {
+                    res.bitwise_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "add_mod" => {
+                    res.add_mod_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "ec_op" => {
+                    res.ec_op_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "ecdsa" => {
+                    res.ecdsa_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "keccak" => {
+                    res.keccak_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "mul_mod" => {
+                    res.mul_mod_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "pedersen" => {
+                    res.pedersen_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                "poseidon" => {
+                    res.poseidon_builtin = MemorySegmentAddresses {
+                        begin_addr: value.begin_addr,
+                        stop_ptr: value.stop_ptr,
+                    }
+                }
+                // Output, executaion 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,22 +70,25 @@ pub fn input_from_finished_runner(runner: CairoRunner, dev_mode: bool) -> CairoI
                 val: bytemuck::cast(v.to_bytes_le()),
             })
         });
-    let trace = runner.relocated_trace.unwrap();
+    let trace = runner.relocated_trace.clone().unwrap();
     let trace = trace.iter().map(|t| t.clone().into());
 
     let mem_config = MemConfig::default();
     let mut mem = MemoryBuilder::from_iter(mem_config, mem);
     let state_transitions = StateTransitions::from_iter(trace, &mut mem, dev_mode);
 
+    let public_input = runner
+        .get_air_public_input()
+        .expect("Unable to get public input from the runner");
+
+    let builtins_segments = BuiltinsSegments::from_memory_segments(&public_input.memory_segments);
+
     // TODO(spapini): Add output builtin to public memory.
     let public_mem_addresses = (0..(program_len as u32)).collect_vec();
     CairoInput {
         state_transitions,
         mem: mem.build(),
         public_mem_addresses,
-        range_check_builtin: MemorySegmentAddresses {
-            begin_addr: 24,
-            stop_ptr: 64,
-        },
+        builtins_segments,
     }
 }