Enabled manual gas handling. (#4880)

crates/cairo-lang-lowering/src/db.rs

@@ -10,13 +10,15 @@ use cairo_lang_utils::ordered_hash_set::OrderedHashSet;
 use cairo_lang_utils::unordered_hash_set::UnorderedHashSet;
 use cairo_lang_utils::Upcast;
 use itertools::Itertools;
+use semantic::corelib;
 
 use crate::add_withdraw_gas::add_withdraw_gas;
 use crate::borrow_check::borrow_check;
 use crate::concretize::concretize_lowered;
 use crate::destructs::add_destructs;
 use crate::diagnostic::{LoweringDiagnostic, LoweringDiagnosticKind};
 use crate::graph_algorithms::feedback_set::flag_add_withdraw_gas;
+use crate::ids::FunctionLongId;
 use crate::implicits::lower_implicits;
 use crate::inline::{apply_inlining, PrivInlineData};
 use crate::lower::{lower_semantic_function, MultiLowering};
@@ -29,7 +31,9 @@ use crate::optimizations::reorder_statements::reorder_statements;
 use crate::optimizations::return_optimization::return_optimization;
 use crate::panic::lower_panics;
 use crate::reorganize_blocks::reorganize_blocks;
-use crate::{ids, DependencyType, FlatBlockEnd, FlatLowered, Location, MatchInfo, Statement};
+use crate::{
+    ids, BlockId, DependencyType, FlatBlockEnd, FlatLowered, Location, MatchInfo, Statement,
+};
 
 // Salsa database interface.
 #[salsa::query_group(LoweringDatabase)]
@@ -433,11 +437,25 @@ fn concrete_function_with_body_lowered(
 /// according to the given [DependencyType]. See [DependencyType] for more information about
 /// what is considered a dependency.
 pub(crate) fn get_direct_callees(
+    db: &dyn LoweringGroup,
     lowered_function: &FlatLowered,
     dependency_type: DependencyType,
 ) -> Vec<ids::FunctionId> {
+    // TODO(orizi): Follow calls for destructors as well.
     let mut direct_callees = Vec::new();
-    for (_, block) in &lowered_function.blocks {
+    if lowered_function.blocks.is_empty() {
+        return direct_callees;
+    }
+    let withdraw_gas_fns = corelib::core_withdraw_gas_fns(db.upcast())
+        .map(|id| db.intern_lowering_function(FunctionLongId::Semantic(id)));
+    let mut visited = vec![false; lowered_function.blocks.len()];
+    let mut stack = vec![BlockId(0)];
+    while let Some(block_id) = stack.pop() {
+        if visited[block_id.0] {
+            continue;
+        }
+        visited[block_id.0] = true;
+        let block = &lowered_function.blocks[block_id];
         for statement in &block.statements {
             if let Statement::Call(statement_call) = statement {
                 // If the dependency_type is DependencyType::Cost and this call has a coupon input,
@@ -449,8 +467,22 @@ pub(crate) fn get_direct_callees(
                 }
             }
         }
-        if let FlatBlockEnd::Match { info: MatchInfo::Extern(s) } = &block.end {
-            direct_callees.push(s.function);
+        match &block.end {
+            FlatBlockEnd::NotSet | FlatBlockEnd::Return(_) | FlatBlockEnd::Panic(_) => {}
+            FlatBlockEnd::Goto(next, _) => stack.push(*next),
+            FlatBlockEnd::Match { info } => {
+                let mut arms = info.arms().iter();
+                if let MatchInfo::Extern(s) = info {
+                    direct_callees.push(s.function);
+                    if DependencyType::Cost == dependency_type
+                        && withdraw_gas_fns.contains(&s.function)
+                    {
+                        // Not following the option when successfully fetched gas.
+                        arms.next();
+                    }
+                }
+                stack.extend(arms.map(|arm| arm.block_id));
+            }
         }
     }
     direct_callees
@@ -462,7 +494,7 @@ fn concrete_function_with_body_postinline_direct_callees(
     dependency_type: DependencyType,
 ) -> Maybe<Vec<ids::FunctionId>> {
     let lowered_function = db.priv_concrete_function_with_body_postinline_lowered(function_id)?;
-    Ok(get_direct_callees(&lowered_function, dependency_type))
+    Ok(get_direct_callees(db, &lowered_function, dependency_type))
 }
 
 fn concrete_function_with_body_postpanic_direct_callees(
@@ -471,7 +503,7 @@ fn concrete_function_with_body_postpanic_direct_callees(
     dependency_type: DependencyType,
 ) -> Maybe<Vec<ids::FunctionId>> {
     let lowered_function = db.concrete_function_with_body_postpanic_lowered(function_id)?;
-    Ok(get_direct_callees(&lowered_function, dependency_type))
+    Ok(get_direct_callees(db, &lowered_function, dependency_type))
 }
 
 /// Given a vector of FunctionIds returns the vector of FunctionWithBodyIds of the

crates/cairo-lang-lowering/src/graph_algorithms/cycles.rs

@@ -13,7 +13,7 @@ pub fn function_with_body_direct_callees(
     dependency_type: DependencyType,
 ) -> Maybe<OrderedHashSet<FunctionId>> {
     let lowered = db.function_with_body_lowering(function_id)?;
-    Ok(get_direct_callees(&lowered, dependency_type).into_iter().collect())
+    Ok(get_direct_callees(db, &lowered, dependency_type).into_iter().collect())
 }
 
 /// Query implementation of

crates/cairo-lang-lowering/src/test_data/cycles

@@ -47,3 +47,95 @@ Statements:
   (v12: core::panics::PanicResult::<((),)>) <- PanicResult::Err(v7)
 End:
   Return(v20, v21, v12)
+
+//! > ==========================================================================
+
+//! > Test explicit gas handling.
+
+//! > test_runner_name
+test_function_lowering
+
+//! > function
+fn foo(x: felt252) {  
+  match core::gas::withdraw_gas() {
+    Option::Some(_) => foo(x),
+    Option::None => {}
+  }
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > lowering_diagnostics
+
+//! > lowering_flat
+Parameters: v4: core::RangeCheck, v5: core::gas::GasBuiltin, v0: core::felt252
+blk0 (root):
+Statements:
+End:
+  Match(match core::gas::withdraw_gas(v4, v5) {
+    Option::Some(v6, v7) => blk1,
+    Option::None(v8, v9) => blk2,
+  })
+
+blk1:
+Statements:
+  (v10: core::RangeCheck, v11: core::gas::GasBuiltin, v1: ()) <- test::foo(v6, v7, v0)
+End:
+  Return(v10, v11, v1)
+
+blk2:
+Statements:
+  (v3: ()) <- struct_construct()
+End:
+  Return(v8, v9, v3)
+
+//! > ==========================================================================
+
+//! > Test nopanic function cycle with withdraw gas.
+
+//! > test_runner_name
+test_function_lowering
+
+//! > function
+fn foo(x: felt252, costs: core::gas::BuiltinCosts) nopanic {
+  match core::gas::withdraw_gas_all(costs) {
+    Option::Some(_) => foo(x, costs),
+    Option::None => {}
+  }
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > lowering_diagnostics
+
+//! > lowering_flat
+Parameters: v5: core::RangeCheck, v6: core::gas::GasBuiltin, v0: core::felt252, v1: core::gas::BuiltinCosts
+blk0 (root):
+Statements:
+End:
+  Match(match core::gas::withdraw_gas_all(v5, v6, v1) {
+    Option::Some(v7, v8) => blk1,
+    Option::None(v9, v10) => blk2,
+  })
+
+blk1:
+Statements:
+  (v11: core::RangeCheck, v12: core::gas::GasBuiltin, v2: ()) <- test::foo(v7, v8, v0, v1)
+End:
+  Return(v11, v12, v2)
+
+blk2:
+Statements:
+  (v4: ()) <- struct_construct()
+End:
+  Return(v9, v10, v4)

crates/cairo-lang-semantic/src/corelib.rs

@@ -476,8 +476,13 @@ pub fn core_felt252_is_zero(db: &dyn SemanticGroup) -> FunctionId {
     get_core_function_id(db, "felt252_is_zero".into(), vec![])
 }
 
-pub fn core_withdraw_gas(db: &dyn SemanticGroup) -> FunctionId {
-    get_function_id(db, core_submodule(db, "gas"), "withdraw_gas".into(), vec![])
+/// The gas withdrawal functions from the `gas` submodule.
+pub fn core_withdraw_gas_fns(db: &dyn SemanticGroup) -> [FunctionId; 2] {
+    let gas = core_submodule(db, "gas");
+    [
+        get_function_id(db, gas, "withdraw_gas".into(), vec![]),
+        get_function_id(db, gas, "withdraw_gas_all".into(), vec![]),
+    ]
 }
 
 pub fn internal_require_implicit(db: &dyn SemanticGroup) -> GenericFunctionId {