From 61b1371d15f8709c7a6ae99f41546d3cb21489aa Mon Sep 17 00:00:00 2001
From: liorgold2 <38202661+liorgold2@users.noreply.github.com>
Date: Fri, 24 Nov 2023 20:53:23 +0200
Subject: [PATCH] Add an error if withdraw_gas/redeposit contain pre-cost
 metadata. (#4468)

---
 .../src/invocations/gas.rs                    | 21 +++++++++++++++++++
 .../src/invocations/mod.rs                    |  3 +++
 2 files changed, 24 insertions(+)

diff --git a/crates/cairo-lang-sierra-to-casm/src/invocations/gas.rs b/crates/cairo-lang-sierra-to-casm/src/invocations/gas.rs
index 10c66fddf18..0a714ded774 100644
--- a/crates/cairo-lang-sierra-to-casm/src/invocations/gas.rs
+++ b/crates/cairo-lang-sierra-to-casm/src/invocations/gas.rs
@@ -3,6 +3,8 @@ use cairo_lang_casm::casm_build_extend;
 use cairo_lang_casm::cell_expression::{CellExpression, CellOperator};
 use cairo_lang_casm::operand::{CellRef, DerefOrImmediate, Register};
 use cairo_lang_sierra::extensions::gas::{CostTokenType, GasConcreteLibfunc};
+use cairo_lang_sierra::program::StatementIdx;
+use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use num_bigint::BigInt;
 
 use super::misc::get_pointer_after_program_code;
@@ -27,6 +29,19 @@ pub fn build(
     }
 }
 
+/// Checks that all the pre-cost variables of the given statement are zero.
+fn check_zero_precost_values(
+    variable_values: &OrderedHashMap<(StatementIdx, CostTokenType), i64>,
+    statement_idx: StatementIdx,
+) -> Result<(), InvocationError> {
+    for token_type in CostTokenType::iter_precost() {
+        if variable_values.get(&(statement_idx, *token_type)).copied().unwrap_or_default() != 0 {
+            return Err(InvocationError::PreCostMetadataNotSupported);
+        }
+    }
+    Ok(())
+}
+
 /// Handles the withdraw_gas invocation.
 fn build_withdraw_gas(
     builder: CompiledInvocationBuilder<'_>,
@@ -36,6 +51,9 @@ fn build_withdraw_gas(
         .get(&(builder.idx, CostTokenType::Const))
         .copied()
         .ok_or(InvocationError::UnknownVariableData)?;
+
+    check_zero_precost_values(variable_values, builder.idx)?;
+
     let [range_check, gas_counter] = builder.try_get_single_cells()?;
 
     let failure_handle_statement_id = get_non_fallthrough_statement_id(&builder);
@@ -86,6 +104,9 @@ fn build_redeposit_gas(
         .get(&(builder.idx, CostTokenType::Const))
         .copied()
         .ok_or(InvocationError::UnknownVariableData)?;
+
+    check_zero_precost_values(variable_values, builder.idx)?;
+
     let gas_counter_value = builder.try_get_single_cells::<1>()?[0]
         .to_deref()
         .ok_or(InvocationError::InvalidReferenceExpressionForArgument)?;
diff --git a/crates/cairo-lang-sierra-to-casm/src/invocations/mod.rs b/crates/cairo-lang-sierra-to-casm/src/invocations/mod.rs
index 4bc2bb5bd79..4a2a5246af0 100644
--- a/crates/cairo-lang-sierra-to-casm/src/invocations/mod.rs
+++ b/crates/cairo-lang-sierra-to-casm/src/invocations/mod.rs
@@ -79,6 +79,9 @@ pub enum InvocationError {
     IntegerOverflow,
     #[error(transparent)]
     FrameStateError(#[from] FrameStateError),
+    // TODO(lior): Remove this error once not used.
+    #[error("This libfunc does not support pre-cost metadata yet.")]
+    PreCostMetadataNotSupported,
 }
 
 /// Describes a simple change in the ap tracking itself.