From 6ec9a77f705ab8e5ef92eb39c896e55a290812a5 Mon Sep 17 00:00:00 2001
From: Esteve Soler Arderiu <esteve.soler@lambdaclass.com>
Date: Tue, 26 Nov 2024 18:08:05 +0100
Subject: [PATCH] Fix gas.

---
 Cargo.toml |  2 ++
 src/gas.rs | 19 +------------------
 src/vm.rs  | 31 +++++++++++++++++++++++++++++--
 3 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index cde8535..323d25f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,9 +6,11 @@ edition = "2021"
 [dependencies]
 cairo-lang-compiler = "2.9.0-dev.0"
 cairo-lang-filesystem = "2.9.0-dev.0"
+cairo-lang-runner = "2.9.0-dev.0"
 cairo-lang-sierra = "2.9.0-dev.0"
 cairo-lang-sierra-ap-change = "2.9.0-dev.0"
 cairo-lang-sierra-gas = "2.9.0-dev.0"
+cairo-lang-starknet-classes = "2.9.0-dev.0"
 cairo-lang-utils = "2.9.0-dev.0"
 clap = { version = "4.5.20", features = ["derive"] }
 k256 = "0.13.4"
diff --git a/src/gas.rs b/src/gas.rs
index 3fb01d3..eeb2f75 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -1,3 +1,4 @@
+use cairo_lang_runner::token_gas_cost;
 use cairo_lang_sierra::{
     extensions::gas::CostTokenType,
     ids::FunctionId,
@@ -195,21 +196,3 @@ fn calc_metadata(
         gas_info: pre_gas_info.combine(post_gas_info),
     })
 }
-
-pub fn token_gas_cost(token_type: CostTokenType) -> usize {
-    match token_type {
-        CostTokenType::Const => 1,
-        CostTokenType::Step
-        | CostTokenType::Hole
-        | CostTokenType::RangeCheck
-        | CostTokenType::RangeCheck96 => {
-            panic!("Token type {:?} has no gas cost.", token_type)
-        }
-        CostTokenType::Pedersen => 4130,
-        CostTokenType::Poseidon => 500,
-        CostTokenType::Bitwise => 594,
-        CostTokenType::EcOp => 4166,
-        CostTokenType::AddMod => 234,
-        CostTokenType::MulMod => 616,
-    }
-}
diff --git a/src/vm.rs b/src/vm.rs
index 4426b77..c274d4f 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -9,6 +9,7 @@ use cairo_lang_sierra::{
     extensions::{
         circuit::CircuitTypeConcrete,
         core::{CoreConcreteLibfunc, CoreLibfunc, CoreType, CoreTypeConcrete},
+        gas::CostTokenType,
         starknet::StarkNetTypeConcrete,
         ConcreteLibfunc, ConcreteType,
     },
@@ -16,6 +17,9 @@ use cairo_lang_sierra::{
     program::{GenFunction, GenStatement, Invocation, Program, StatementIdx},
     program_registry::ProgramRegistry,
 };
+use cairo_lang_starknet_classes::{
+    casm_contract_class::ENTRY_POINT_COST, contract_class::ContractEntryPoints,
+};
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use smallvec::{smallvec, SmallVec};
 use starknet_types_core::felt::Felt;
@@ -79,10 +83,33 @@ impl VirtualMachine {
 }
 
 impl<S: StarknetSyscallHandler> VirtualMachine<S> {
-    pub fn new_starknet(program: Arc<Program>, syscall_handler: S) -> Self {
+    pub fn new_starknet(
+        program: Arc<Program>,
+        entry_points: &ContractEntryPoints,
+        syscall_handler: S,
+    ) -> Self {
         let registry = ProgramRegistry::new(&program).unwrap();
         Self {
-            gas: GasMetadata::new(&program, Some(MetadataComputationConfig::default())).unwrap(),
+            gas: GasMetadata::new(
+                &program,
+                Some(MetadataComputationConfig {
+                    function_set_costs: entry_points
+                        .constructor
+                        .iter()
+                        .chain(entry_points.external.iter())
+                        .chain(entry_points.l1_handler.iter())
+                        .map(|id| {
+                            (
+                                program.funcs[id.function_idx].id.clone(),
+                                [(CostTokenType::Const, ENTRY_POINT_COST)].into(),
+                            )
+                        })
+                        .collect(),
+                    linear_gas_solver: true,
+                    linear_ap_change_solver: true,
+                }),
+            )
+            .unwrap(),
             program,
             registry,
             syscall_handler,