(optimization) A new libfunc for calculating the pointer to temporary storage using a double call

SIERRA_UPDATE_MINOR_CHANGE_TAG=adding a libfunc for optimizing creating boxes from temp storage

Author: eytan-starkware

crates/cairo-lang-runner/src/casm_run/test.rs

@@ -20,7 +20,7 @@ use crate::{CairoHintProcessor, StarknetState, build_hints_dict};
 fn assembled(casm: CasmContext) -> AssembledCairoProgram {
     CairoProgram {
         instructions: casm.instructions,
-        consts_info: Default::default(),
+        segments_info: Default::default(),
         debug_info: CairoProgramDebugInfo { sierra_statement_info: vec![] },
     }
     .assemble()

crates/cairo-lang-sierra-ap-change/src/core_libfunc_ap_change.rs

@@ -113,6 +113,10 @@ pub fn core_libfunc_ap_change<InfoProvider: InvocationApChangeInfoProvider>(
             BoxConcreteLibfunc::Into(_) => vec![ApChange::Known(1)],
             BoxConcreteLibfunc::Unbox(_) => vec![ApChange::Known(0)],
             BoxConcreteLibfunc::ForwardSnapshot(_) => vec![ApChange::Known(0)],
+            BoxConcreteLibfunc::FromTempStore(_) => {
+                // Total: 2 calls * 2
+                vec![ApChange::Known(4)]
+            }
         },
         Cast(libfunc) => match libfunc {
             CastConcreteLibfunc::Downcast(libfunc) => {

crates/cairo-lang-sierra-gas/src/core_libfunc_cost_base.rs

@@ -371,6 +371,7 @@ pub fn core_libfunc_cost(
             BoxConcreteLibfunc::Unbox(_) | BoxConcreteLibfunc::ForwardSnapshot(_) => {
                 vec![ConstCost::default().into()]
             }
+            BoxConcreteLibfunc::FromTempStore(_) => vec![ConstCost::steps(4).into()],
         },
         Mem(libfunc) => match libfunc {
             StoreTemp(libfunc) => {

crates/cairo-lang-sierra-to-casm/src/compiler.rs

@@ -1,8 +1,10 @@
 use std::fmt::Display;
 
 use cairo_lang_casm::assembler::AssembledCairoProgram;
+use cairo_lang_casm::casm;
 use cairo_lang_casm::instructions::{Instruction, InstructionBody, RetInstruction};
 use cairo_lang_sierra::extensions::ConcreteLibfunc;
+use cairo_lang_sierra::extensions::boxing::BoxConcreteLibfunc;
 use cairo_lang_sierra::extensions::circuit::{CircuitConcreteLibfunc, CircuitInfo, VALUE_SIZE};
 use cairo_lang_sierra::extensions::const_type::ConstConcreteLibfunc;
 use cairo_lang_sierra::extensions::core::{
@@ -106,7 +108,7 @@ pub struct SierraToCasmConfig {
 pub struct CairoProgram {
     pub instructions: Vec<Instruction>,
     pub debug_info: CairoProgramDebugInfo,
-    pub consts_info: ConstsInfo,
+    pub segments_info: SegmentsInfo,
 }
 impl Display for CairoProgram {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -116,24 +118,35 @@ impl Display for CairoProgram {
                 writeln!(f, "{instruction}; // {bytecode_offset}")?;
                 bytecode_offset += instruction.body.op_size();
             }
-            for segment in self.consts_info.segments.values() {
+            for segment in self.segments_info.const_segments.values() {
                 writeln!(f, "ret; // {bytecode_offset}")?;
                 bytecode_offset += 1;
                 for value in &segment.values {
                     writeln!(f, "dw {value}; // {bytecode_offset}")?;
                     bytecode_offset += 1;
                 }
             }
+            for segment in self.segments_info.utility_segments.values() {
+                for instruction in &segment.instructions {
+                    writeln!(f, "{instruction}; // {bytecode_offset}")?;
+                    bytecode_offset += instruction.body.op_size();
+                }
+            }
         } else {
             for instruction in &self.instructions {
                 writeln!(f, "{instruction};")?;
             }
-            for segment in self.consts_info.segments.values() {
+            for segment in self.segments_info.const_segments.values() {
                 writeln!(f, "ret;")?;
                 for value in &segment.values {
                     writeln!(f, "dw {value};")?;
                 }
             }
+            for segment in self.segments_info.utility_segments.values() {
+                for instruction in &segment.instructions {
+                    writeln!(f, "{instruction};")?;
+                }
+            }
         }
         Ok(())
     }
@@ -166,10 +179,13 @@ impl CairoProgram {
         else {
             panic!("`ret` instruction should be a single word.")
         };
-        for segment in self.consts_info.segments.values() {
+        for segment in self.segments_info.const_segments.values() {
             bytecode.push(ret_bytecode.clone());
             bytecode.extend(segment.values.clone());
         }
+        for segment in self.segments_info.utility_segments.values() {
+            bytecode.extend(segment.instructions.iter().flat_map(|inst| inst.assemble().encode()));
+        }
         for instruction in footer {
             assert!(
                 instruction.hints.is_empty(),
@@ -238,14 +254,18 @@ pub struct CairoProgramDebugInfo {
 
 /// The information about the constants used in the program.
 #[derive(Debug, Eq, PartialEq, Default, Clone)]
-pub struct ConstsInfo {
-    pub segments: OrderedHashMap<u32, ConstSegment>,
+pub struct SegmentsInfo {
+    pub const_segments: OrderedHashMap<u32, ConstSegment>,
     pub total_segments_size: usize,
 
     /// Maps a circuit to its segment id.
     pub circuit_segments: OrderedHashMap<ConcreteTypeId, u32>,
+
+    /// Utility segments for code reuse.
+    pub utility_segments: OrderedHashMap<UtilityId, UtilitySegment>,
 }
-impl ConstsInfo {
+
+impl SegmentsInfo {
     /// Creates a new `ConstSegmentsInfo` from the given libfuncs.
     pub fn new<'a>(
         program_info: &ProgramRegistryInfo,
@@ -272,14 +292,14 @@ impl ConstsInfo {
             Ok(())
         };
 
-        let mut segments = OrderedHashMap::default();
+        let mut const_segments = OrderedHashMap::default();
 
         for id in libfunc_ids.clone() {
             if let CoreConcreteLibfunc::Const(ConstConcreteLibfunc::AsBox(as_box)) =
                 program_info.registry.get_libfunc(id).unwrap()
             {
                 add_const(
-                    &mut segments,
+                    &mut const_segments,
                     as_box.segment_id,
                     as_box.const_type.clone(),
                     extract_const_value(program_info, &as_box.const_type).unwrap(),
@@ -288,18 +308,18 @@ impl ConstsInfo {
         }
 
         // Check that the segments were declared in order and without holes.
-        if segments
+        if const_segments
             .keys()
             .enumerate()
             .any(|(i, segment_id)| i != segment_id.into_or_panic::<usize>())
         {
             return Err(CompilationError::ConstSegmentsOutOfOrder);
         }
 
-        let mut next_segment = segments.len() as u32;
+        let mut next_segment = const_segments.len() as u32;
         let mut circuit_segments = OrderedHashMap::default();
 
-        for id in libfunc_ids {
+        for id in libfunc_ids.clone() {
             if let CoreConcreteLibfunc::Circuit(CircuitConcreteLibfunc::GetDescriptor(libfunc)) =
                 program_info.registry.get_libfunc(id).unwrap()
             {
@@ -314,19 +334,50 @@ impl ConstsInfo {
                     push_offset(gate_offsets.output);
                 }
 
-                add_const(&mut segments, next_segment, circ_ty.clone(), const_value)?;
+                add_const(&mut const_segments, next_segment, circ_ty.clone(), const_value)?;
                 circuit_segments.insert(circ_ty.clone(), next_segment);
                 next_segment += 1;
             }
         }
 
+        let mut utility_instructions = Vec::default();
+
+        // Check for BoxFromTempStore usage
+        let has_box_from_temp_store = libfunc_ids.into_iter().any(|id| {
+            matches!(
+                program_info.registry.get_libfunc(id),
+                Ok(CoreConcreteLibfunc::Box(BoxConcreteLibfunc::FromTempStore(_)))
+            )
+        });
+
+        if has_box_from_temp_store {
+            utility_instructions.push((
+                UtilityId::BoxFromTempStore,
+                casm! {
+                    call rel 2;
+                    ret;
+                }
+                .instructions,
+            ));
+        }
+
         let mut total_segments_size = 0;
-        for (_, segment) in segments.iter_mut() {
+        for (_, segment) in const_segments.iter_mut() {
             segment.segment_offset = total_segments_size;
             // Add 1 for the `ret` instruction.
             total_segments_size += 1 + segment.values.len();
         }
-        Ok(Self { segments, total_segments_size, circuit_segments })
+
+        let utility_segments = utility_instructions
+            .into_iter()
+            .map(|(id, instructions)| {
+                let segment = UtilitySegment { instructions, offset: total_segments_size };
+                total_segments_size +=
+                    segment.instructions.iter().map(|i| i.body.op_size()).sum::<usize>();
+                (id, segment)
+            })
+            .collect();
+        Ok(Self { const_segments, total_segments_size, circuit_segments, utility_segments })
     }
 }
 
@@ -341,6 +392,20 @@ pub struct ConstSegment {
     pub segment_offset: usize,
 }
 
+/// The data for a single utility segment (code).
+#[derive(Debug, Eq, PartialEq, Default, Clone)]
+pub struct UtilitySegment {
+    /// The instructions in the segment.
+    pub instructions: Vec<Instruction>,
+    /// Offset from start of segments
+    pub offset: usize,
+}
+
+#[derive(Debug, Eq, PartialEq, Hash, Clone)]
+pub enum UtilityId {
+    BoxFromTempStore,
+}
+
 /// Gets a concrete type, if it is a const type returns a vector of the values to be stored in the
 /// const segment.
 fn extract_const_value(
@@ -639,17 +704,17 @@ pub fn compile(
         .max_bytecode_size
         .checked_sub(program_offset)
         .ok_or_else(|| Box::new(CompilationError::CodeSizeLimitExceeded))?;
-    let consts_info = ConstsInfo::new(
+    let segments_info = SegmentsInfo::new(
         program_info,
         program.libfunc_declarations.iter().map(|ld| &ld.id),
         &circuits_info.circuits,
         const_segments_max_size,
     )?;
-    relocate_instructions(&relocations, &statement_offsets, &consts_info, &mut instructions);
+    relocate_instructions(&relocations, &statement_offsets, &segments_info, &mut instructions);
 
     Ok(CairoProgram {
         instructions,
-        consts_info,
+        segments_info,
         debug_info: CairoProgramDebugInfo { sierra_statement_info },
     })
 }

crates/cairo-lang-sierra-to-casm/src/compiler_test.rs

@@ -760,6 +760,93 @@ indoc! {"
         dw 5;
     "};
     "Get builtin costs with a const segment.")]
+#[test_case(indoc! {"
+        type felt252 = felt252;
+        type Box<felt252> = Box<felt252>;
+
+        libfunc felt252_const<42> = felt252_const<42>;
+        libfunc store_temp<felt252> = store_temp<felt252>;
+        libfunc drop<felt252> = drop<felt252>;
+        libfunc box_from_temp_store<felt252> = box_from_temp_store<felt252>;
+        libfunc store_temp<Box<felt252>> = store_temp<Box<felt252>>;
+
+        felt252_const<42>() -> ([0]);
+        store_temp<felt252>([0]) -> ([0]);  // Allocate and store the constant felt252
+        drop<felt252>([0]) -> ();  // Drop the stored value
+        box_from_temp_store<felt252>() -> ([1]);  // Compute address for a felt252
+        store_temp<Box<felt252>>([1]) -> ([1]);
+        return([1]);
+
+        test_program@0() -> (Box<felt252>);
+    "},
+    false,
+    indoc! {"
+        [ap + 0] = 42, ap++;
+        call rel 5;
+        [ap + 0] = [ap + -2] + -1, ap++;
+        ret;
+        call rel 2;
+        ret;
+    "};
+    "box_from_temp_store with felt252 (size 1), after storing constant at [0].")]
+#[test_case(indoc! {"
+        type u128 = u128;
+        type felt252 = felt252;
+        type core::integer::u256 = Struct<ut@core::integer::u256, u128, u128>;
+        type Box<core::integer::u256> = Box<core::integer::u256>;
+
+        libfunc u128_const<0> = u128_const<0>;
+        libfunc struct_construct<core::integer::u256> = struct_construct<core::integer::u256>;
+        libfunc store_temp<core::integer::u256> = store_temp<core::integer::u256>;
+        libfunc drop<core::integer::u256> = drop<core::integer::u256>;
+        libfunc box_from_temp_store<core::integer::u256> = box_from_temp_store<core::integer::u256>;
+        libfunc store_temp<Box<core::integer::u256>> = store_temp<Box<core::integer::u256>>;
+
+        u128_const<0>() -> ([0]);
+        u128_const<0>() -> ([1]);
+        struct_construct<core::integer::u256>([0], [1]) -> ([2]);
+        store_temp<core::integer::u256>([2]) -> ([2]);  // Store u256
+        drop<core::integer::u256>([2]) -> ();  // Drop the stored value
+        box_from_temp_store<core::integer::u256>() -> ([3]);  // Compute address for a u256
+        store_temp<Box<core::integer::u256>>([3]) -> ([3]);
+        return([3]);
+
+        test_program@0() -> (Box<core::integer::u256>);
+    "},
+    false,
+    indoc! {"
+        [ap + 0] = 0, ap++;
+        [ap + 0] = 0, ap++;
+        call rel 5;
+        [ap + 0] = [ap + -2] + -2, ap++;
+        ret;
+        call rel 2;
+        ret;
+    "};
+    "box_from_temp_store with u256 (size 2), after storing u256 at [0].")]
+#[test_case(indoc! {"
+        type Unit = Struct<ut@Tuple>;
+        type felt252 = felt252;
+        type Box<Unit> = Box<Unit>;
+
+        libfunc box_from_temp_store<Unit> = box_from_temp_store<Unit>;
+        libfunc store_temp<Box<Unit>> = store_temp<Box<Unit>>;
+
+        box_from_temp_store<Unit>() -> ([0]);  // No prior allocation for size 0 unit
+        store_temp<Box<Unit>>([0]) -> ([0]);
+        return([0]);
+
+        test_program@0() -> (Box<Unit>);
+    "},
+    false,
+    indoc! {"
+        call rel 5;
+        [ap + 0] = [ap + -2] + 0, ap++;
+        ret;
+        call rel 2;
+        ret;
+    "};
+    "box_from_temp_store with unit type (size 0), no prior allocation.")]
 fn sierra_to_casm(sierra_code: &str, gas_usage_check: bool, expected_casm: &str) {
     let program = ProgramParser::new().parse(sierra_code).unwrap();
     let program_info = ProgramRegistryInfo::new(&program).unwrap();