Change end of program relocation to point after the const segment. (#…

…4714)

crates/cairo-lang-casm/src/assembler.rs

@@ -1,5 +1,6 @@
 use num_bigint::{BigInt, ToBigInt};
 
+use crate::hints::Hint;
 use crate::instructions::{Instruction, InstructionBody};
 use crate::operand::{DerefOrImmediate, Operation, Register, ResOperand};
 
@@ -71,6 +72,14 @@ pub struct InstructionRepr {
     pub opcode: Opcode,
 }
 
+/// An assembled representation of a cairo program.
+pub struct AssembledCairoProgram {
+    /// The bytecode of the program.
+    pub bytecode: Vec<BigInt>,
+    /// The list of hints, and the instruction index they refer to.
+    pub hints: Vec<(usize, Vec<Hint>)>,
+}
+
 impl Instruction {
     pub fn assemble(&self) -> InstructionRepr {
         match &self.body {

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

@@ -1,5 +1,6 @@
 use std::fmt::Display;
 
+use cairo_lang_casm::assembler::AssembledCairoProgram;
 use cairo_lang_casm::instructions::{Instruction, InstructionBody, RetInstruction};
 use cairo_lang_sierra::extensions::core::{CoreConcreteLibfunc, CoreLibfunc, CoreType};
 use cairo_lang_sierra::extensions::lib_func::SierraApChange;
@@ -79,6 +80,31 @@ impl Display for CairoProgram {
     }
 }
 
+impl CairoProgram {
+    /// Creates an assembled representation of the program.
+    pub fn assemble(&self) -> AssembledCairoProgram {
+        let mut bytecode = vec![];
+        let mut hints = vec![];
+        for instruction in self.instructions.iter() {
+            if !instruction.hints.is_empty() {
+                hints.push((bytecode.len(), instruction.hints.clone()))
+            }
+            bytecode.extend(instruction.assemble().encode().into_iter())
+        }
+        let [ref ret_bytecode] = Instruction::new(InstructionBody::Ret(RetInstruction {}), false)
+            .assemble()
+            .encode()[..]
+        else {
+            panic!("Ret instruction is a single word")
+        };
+        for const_allocation in self.const_segment_info.const_allocations.values() {
+            bytecode.push(ret_bytecode.clone());
+            bytecode.extend(const_allocation.values.clone());
+        }
+        AssembledCairoProgram { bytecode, hints }
+    }
+}
+
 /// The debug information of a compilation from Sierra to casm.
 #[derive(Debug, Eq, PartialEq)]
 pub struct SierraStatementDebugInfo {
@@ -121,7 +147,7 @@ impl ConstSegmentInfoBuilder {
             const_segment_size += const_allocation.values.len() + 1;
             const_allocations.insert(const_type.clone(), const_allocation);
         }
-        Ok(ConstSegmentInfo { const_allocations })
+        Ok(ConstSegmentInfo { const_allocations, const_segment_size })
     }
 }
 
@@ -139,6 +165,8 @@ pub struct ConstAllocation {
 pub struct ConstSegmentInfo {
     /// A map between the const type and the data of the const.
     pub const_allocations: OrderedHashMap<ConcreteTypeId, ConstAllocation>,
+    /// The size of the constants segment.
+    pub const_segment_size: usize,
 }
 
 /// Gets a concrete type, if it is a const type returns a vector of the values to be stored in the

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

@@ -697,6 +697,67 @@ indoc! {"
         dw 5;
     "};
     "Constant enums.")]
+#[test_case(indoc! {"
+
+        type BuiltinCosts = BuiltinCosts;
+
+        libfunc get_builtin_costs = get_builtin_costs;
+        libfunc store_temp<BuiltinCosts> = store_temp<BuiltinCosts>;
+        libfunc drop<BuiltinCosts> = drop<BuiltinCosts>;
+
+        get_builtin_costs() -> ([1]);
+        store_temp<BuiltinCosts>([1]) -> ([1]);
+        drop<BuiltinCosts>([1]) -> ();
+        return();
+
+        test_program@0() -> ();
+    "},
+    false,
+    indoc! {"
+        call rel 6;
+        [ap + 0] = [ap + -1] + 5, ap++;
+        [ap + 0] = [[ap + -1] + 0], ap++;
+        ret;
+    "};
+    "Get builtin costs.")]
+#[test_case(indoc! {"
+        type BuiltinCosts = BuiltinCosts;
+        type felt252 = felt252;
+        type const<felt252, 5> = const<felt252, 5>;
+        type Box<felt252> = Box<felt252>;
+
+        libfunc get_builtin_costs = get_builtin_costs;
+        libfunc store_temp<BuiltinCosts> = store_temp<BuiltinCosts>;
+        libfunc drop<BuiltinCosts> = drop<BuiltinCosts>;
+        libfunc const_as_box<const<felt252, 5>> = const_as_box<const<felt252, 5>>;
+        libfunc unbox<felt252> = unbox<felt252>;
+        libfunc store_temp<felt252> = store_temp<felt252>;
+        libfunc drop<felt252> = drop<felt252>;
+
+        get_builtin_costs() -> ([1]);
+        store_temp<BuiltinCosts>([1]) -> ([1]);
+        drop<BuiltinCosts>([1]) -> ();
+        const_as_box<const<felt252, 5>>() -> ([2]);
+        unbox<felt252>([2]) -> ([2]);
+        store_temp<felt252>([2]) -> ([2]);
+        drop<felt252>([2]) -> ();
+        return();
+
+        test_program@0() -> ();
+    "},
+    false,
+    indoc! {"
+        call rel 13;
+        [ap + 0] = [ap + -1] + 12, ap++;
+        [ap + 0] = [[ap + -1] + 0], ap++;
+        call rel 6;
+        [ap + 0] = [ap + -1] + 5, ap++;
+        [ap + 0] = [[ap + -1] + 0], ap++;
+        ret;
+        ret;
+        dw 5;
+    "};
+    "Get builtin costs with a const segment.")]
 fn sierra_to_casm(sierra_code: &str, check_gas_usage: bool, expected_casm: &str) {
     let program = ProgramParser::new().parse(sierra_code).unwrap();
     pretty_assertions::assert_eq!(

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

@@ -34,7 +34,9 @@ impl Relocation {
     ) {
         let target_pc = match self {
             Relocation::RelativeStatementId(statement_idx) => statement_offsets[statement_idx.0],
-            Relocation::EndOfProgram => *statement_offsets.last().unwrap(),
+            Relocation::EndOfProgram => {
+                *statement_offsets.last().unwrap() + const_segment_info.const_segment_size
+            }
             Relocation::Const(ty) => {
                 *statement_offsets.last().unwrap()
                     + const_segment_info

crates/cairo-lang-starknet/src/casm_contract_class.rs

@@ -3,6 +3,7 @@
 mod test;
 
 use cairo_felt::Felt252;
+use cairo_lang_casm::assembler::AssembledCairoProgram;
 use cairo_lang_casm::hints::{Hint, PythonicHint};
 use cairo_lang_sierra::extensions::array::ArrayType;
 use cairo_lang_sierra::extensions::bitwise::BitwiseType;
@@ -345,20 +346,16 @@ impl CasmContractClass {
         let cairo_program =
             cairo_lang_sierra_to_casm::compiler::compile(&program, &metadata, gas_usage_check)?;
 
-        let mut bytecode = vec![];
-        let mut hints = vec![];
-        for instruction in cairo_program.instructions {
-            if !instruction.hints.is_empty() {
-                hints.push((bytecode.len(), instruction.hints.clone()))
-            }
-            bytecode.extend(instruction.assemble().encode().iter().map(|big_int| {
+        let AssembledCairoProgram { bytecode, hints } = cairo_program.assemble();
+        let bytecode = bytecode
+            .iter()
+            .map(|big_int| {
                 let (_q, reminder) = big_int.magnitude().div_rem(&prime);
-
                 BigUintAsHex {
                     value: if big_int.is_negative() { &prime - reminder } else { reminder },
                 }
-            }))
-        }
+            })
+            .collect();
 
         let builtin_types = UnorderedHashSet::<GenericTypeId>::from_iter([
             RangeCheckType::id(),