Added gas redeposit optimization stage. (#6280)

corelib/src/gas.cairo

@@ -35,5 +35,12 @@ pub extern fn withdraw_gas_all(
     costs: BuiltinCosts
 ) -> Option<()> implicits(RangeCheck, GasBuiltin) nopanic;
 
+
+/// Returns unused gas into the gas builtin.
+///
+/// Useful for cases where different branches take different amounts of gas, but gas withdrawal is
+/// the same for both.
+pub extern fn redeposit_gas() implicits(GasBuiltin) nopanic;
+
 /// Returns the `BuiltinCosts` table to be used in `withdraw_gas_all`.
 pub extern fn get_builtin_costs() -> BuiltinCosts nopanic;

corelib/src/test/coupon_test.cairo

@@ -1,5 +1,3 @@
-use crate::test::test_utils::assert_eq;
-
 extern fn coupon_buy<T>() -> T nopanic;
 
 #[feature("corelib-internal-use")]
@@ -22,6 +20,6 @@ fn test_arr_sum() {
     let available_gas = crate::testing::get_available_gas();
     let res = arr_sum(arr);
     // Check that arr_sum did not consume any gas.
-    assert_eq(@crate::testing::get_available_gas(), @available_gas, 'Gas was consumed by arr_sum');
-    assert_eq(@res, @12, 'Wrong array sum.');
+    assert_ge!(core::testing::get_available_gas(), available_gas, "Gas was consumed by arr_sum.");
+    assert_eq!(res, 12);
 }

crates/cairo-lang-lowering/src/optimizations/gas_redeposit.rs

@@ -0,0 +1,94 @@
+#[cfg(test)]
+#[path = "gas_redeposit_test.rs"]
+mod test;
+
+use cairo_lang_filesystem::flag::Flag;
+use cairo_lang_filesystem::ids::FlagId;
+use cairo_lang_semantic::corelib;
+use itertools::Itertools;
+
+use crate::db::LoweringGroup;
+use crate::ids::{ConcreteFunctionWithBodyId, SemanticFunctionIdEx};
+use crate::implicits::FunctionImplicitsTrait;
+use crate::{BlockId, FlatBlockEnd, FlatLowered, Statement, StatementCall};
+
+/// Adds redeposit gas actions.
+///
+/// The algorithm is as follows:
+/// Check if the function will have the `GasBuiltin` implicit after the lower_implicits stage.
+/// If so, add a `redeposit_gas` call at the beginning of every branch in the code.
+/// Otherwise, do nothing.
+///
+/// Note that for implementation simplicity this stage must be applied before `LowerImplicits`
+/// stage.
+pub fn gas_redeposit(
+    db: &dyn LoweringGroup,
+    function_id: ConcreteFunctionWithBodyId,
+    lowered: &mut FlatLowered,
+) {
+    if lowered.blocks.is_empty() {
+        return;
+    }
+    if !matches!(
+        db.get_flag(FlagId::new(db.upcast(), "add_redeposit_gas")),
+        Some(flag) if matches!(*flag, Flag::AddRedepositGas(true))
+    ) {
+        return;
+    }
+    let gb_ty = corelib::get_core_ty_by_name(db.upcast(), "GasBuiltin".into(), vec![]);
+    // Checking if the implicits of this function past lowering includes `GasBuiltin`.
+    if let Ok(implicits) = db.function_with_body_implicits(function_id) {
+        if !implicits.into_iter().contains(&gb_ty) {
+            return;
+        }
+    }
+    assert!(
+        lowered.parameters.iter().all(|p| lowered.variables[*p].ty != gb_ty),
+        "`GasRedeposit` stage must be called before `LowerImplicits` stage"
+    );
+
+    let redeposit_gas = corelib::get_function_id(
+        db.upcast(),
+        corelib::core_submodule(db.upcast(), "gas"),
+        "redeposit_gas".into(),
+        vec![],
+    )
+    .lowered(db);
+    let mut stack = vec![BlockId::root()];
+    let mut visited = vec![false; lowered.blocks.len()];
+    let mut redeposit_commands = vec![];
+    while let Some(block_id) = stack.pop() {
+        if visited[block_id.0] {
+            continue;
+        }
+        visited[block_id.0] = true;
+        let block = &lowered.blocks[block_id];
+        match &block.end {
+            FlatBlockEnd::Goto(block_id, _) => {
+                stack.push(*block_id);
+            }
+            FlatBlockEnd::Match { info } => {
+                let location = info.location().with_auto_generation_note(db, "withdraw_gas");
+                for arm in info.arms() {
+                    stack.push(arm.block_id);
+                    redeposit_commands.push((arm.block_id, location));
+                }
+            }
+            &FlatBlockEnd::Return(..) | FlatBlockEnd::Panic(_) => {}
+            FlatBlockEnd::NotSet => unreachable!("Block end not set"),
+        }
+    }
+    for (block_id, location) in redeposit_commands {
+        let block = &mut lowered.blocks[block_id];
+        block.statements.insert(
+            0,
+            Statement::Call(StatementCall {
+                function: redeposit_gas,
+                inputs: vec![],
+                with_coupon: false,
+                outputs: vec![],
+                location,
+            }),
+        );
+    }
+}

crates/cairo-lang-lowering/src/optimizations/gas_redeposit_test.rs

@@ -0,0 +1,77 @@
+use std::ops::Deref;
+use std::sync::Arc;
+
+use cairo_lang_debug::DebugWithDb;
+use cairo_lang_filesystem::db::FilesGroupEx;
+use cairo_lang_filesystem::flag::Flag;
+use cairo_lang_filesystem::ids::FlagId;
+use cairo_lang_semantic::test_utils::setup_test_function;
+use cairo_lang_test_utils::parse_test_file::{TestFileRunner, TestRunnerResult};
+use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
+
+use super::gas_redeposit;
+use crate::db::LoweringGroup;
+use crate::fmt::LoweredFormatter;
+use crate::ids::ConcreteFunctionWithBodyId;
+use crate::test_utils::LoweringDatabaseForTesting;
+
+cairo_lang_test_utils::test_file_test_with_runner!(
+    gas_redeposit,
+    "src/optimizations/test_data",
+    {
+        gas_redeposit: "gas_redeposit",
+    },
+    GetRedepositTestRunner
+);
+
+struct GetRedepositTestRunner {
+    db: LoweringDatabaseForTesting,
+}
+impl Default for GetRedepositTestRunner {
+    fn default() -> Self {
+        let mut db = LoweringDatabaseForTesting::new();
+        let flag = FlagId::new(&db, "add_redeposit_gas");
+        db.set_flag(flag, Some(Arc::new(Flag::AddRedepositGas(true))));
+        Self { db }
+    }
+}
+
+impl TestFileRunner for GetRedepositTestRunner {
+    fn run(
+        &mut self,
+        inputs: &OrderedHashMap<String, String>,
+        _args: &OrderedHashMap<String, String>,
+    ) -> TestRunnerResult {
+        let db = &self.db;
+        let (test_function, semantic_diagnostics) = setup_test_function(
+            db,
+            inputs["function"].as_str(),
+            inputs["function_name"].as_str(),
+            inputs["module_code"].as_str(),
+        )
+        .split();
+        let function_id =
+            ConcreteFunctionWithBodyId::from_semantic(db, test_function.concrete_function_id);
+
+        let before =
+            db.concrete_function_with_body_postpanic_lowered(function_id).unwrap().deref().clone();
+
+        let lowering_diagnostics = db.module_lowering_diagnostics(test_function.module_id).unwrap();
+
+        let mut after = before.clone();
+        gas_redeposit(db, function_id, &mut after);
+
+        TestRunnerResult::success(OrderedHashMap::from([
+            ("semantic_diagnostics".into(), semantic_diagnostics),
+            (
+                "before".into(),
+                format!("{:?}", before.debug(&LoweredFormatter::new(db, &before.variables))),
+            ),
+            (
+                "after".into(),
+                format!("{:?}", after.debug(&LoweredFormatter::new(db, &after.variables))),
+            ),
+            ("lowering_diagnostics".into(), lowering_diagnostics.format(db)),
+        ]))
+    }
+}

crates/cairo-lang-lowering/src/optimizations/mod.rs

@@ -2,6 +2,7 @@ pub mod branch_inversion;
 pub mod cancel_ops;
 pub mod config;
 pub mod const_folding;
+pub mod gas_redeposit;
 pub mod match_optimizer;
 pub mod remappings;
 pub mod reorder_statements;

crates/cairo-lang-lowering/src/optimizations/strategy.rs

@@ -1,6 +1,7 @@
 use cairo_lang_diagnostics::Maybe;
 use cairo_lang_utils::{define_short_id, Intern, LookupIntern};
 
+use super::gas_redeposit::gas_redeposit;
 use crate::db::LoweringGroup;
 use crate::ids::ConcreteFunctionWithBodyId;
 use crate::implicits::lower_implicits;
@@ -29,6 +30,7 @@ pub enum OptimizationPhase {
     ReorganizeBlocks,
     ReturnOptimization,
     SplitStructs,
+    GasRedeposit,
     /// The following is not really an optimization but we want to apply optimizations before and
     /// after it, so it is convenient to treat it as an optimization.
     LowerImplicits,
@@ -56,6 +58,7 @@ impl OptimizationPhase {
             OptimizationPhase::ReturnOptimization => return_optimization(db, lowered),
             OptimizationPhase::SplitStructs => split_structs(lowered),
             OptimizationPhase::LowerImplicits => lower_implicits(db, function, lowered),
+            OptimizationPhase::GasRedeposit => gas_redeposit(db, function, lowered),
         }
         Ok(())
     }
@@ -120,6 +123,7 @@ pub fn baseline_optimization_strategy(db: &dyn LoweringGroup) -> OptimizationStr
 /// Query implementation of [crate::db::LoweringGroup::final_optimization_strategy].
 pub fn final_optimization_strategy(db: &dyn LoweringGroup) -> OptimizationStrategyId {
     OptimizationStrategy(vec![
+        OptimizationPhase::GasRedeposit,
         OptimizationPhase::LowerImplicits,
         OptimizationPhase::ReorganizeBlocks,
         OptimizationPhase::CancelOps,