From 45078310b41768291e97ca7fc7dff3456d77c739 Mon Sep 17 00:00:00 2001
From: Gil Ben-Shachar <gil@starkware.co>
Date: Tue, 2 Jan 2024 12:50:29 +0200
Subject: [PATCH] Change end of program relocation to point after the const
 segment.

commit-id:ef16e798
---
 .../cairo-lang-sierra-to-casm/src/compiler.rs |  4 +-
 .../src/compiler_test.rs                      | 61 +++++++++++++++++++
 .../src/relocations.rs                        |  4 +-
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/crates/cairo-lang-sierra-to-casm/src/compiler.rs b/crates/cairo-lang-sierra-to-casm/src/compiler.rs
index ddf1c2616b0..6b5f641001d 100644
--- a/crates/cairo-lang-sierra-to-casm/src/compiler.rs
+++ b/crates/cairo-lang-sierra-to-casm/src/compiler.rs
@@ -147,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 })
     }
 }
 
@@ -165,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
diff --git a/crates/cairo-lang-sierra-to-casm/src/compiler_test.rs b/crates/cairo-lang-sierra-to-casm/src/compiler_test.rs
index dd841dd3c3f..644b6be993b 100644
--- a/crates/cairo-lang-sierra-to-casm/src/compiler_test.rs
+++ b/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!(
diff --git a/crates/cairo-lang-sierra-to-casm/src/relocations.rs b/crates/cairo-lang-sierra-to-casm/src/relocations.rs
index ce1ee3c79cf..a7870f53360 100644
--- a/crates/cairo-lang-sierra-to-casm/src/relocations.rs
+++ b/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