cairo_lang_lowering/optimizations/
gas_redeposit.rs

1#[cfg(test)]
2#[path = "gas_redeposit_test.rs"]
3mod test;
4
5use cairo_lang_filesystem::flag::Flag;
6use cairo_lang_filesystem::ids::FlagId;
7use cairo_lang_semantic::corelib;
8use itertools::Itertools;
9
10use crate::db::LoweringGroup;
11use crate::ids::{ConcreteFunctionWithBodyId, SemanticFunctionIdEx};
12use crate::implicits::FunctionImplicitsTrait;
13use crate::{BlockId, FlatBlockEnd, FlatLowered, Statement, StatementCall};
14
15/// Adds redeposit gas actions.
16///
17/// The algorithm is as follows:
18/// Check if the function will have the `GasBuiltin` implicit after the lower_implicits stage.
19/// If so, add a `redeposit_gas` call at the beginning of every branch in the code.
20/// Otherwise, do nothing.
21///
22/// Note that for implementation simplicity this stage must be applied before `LowerImplicits`
23/// stage.
24pub fn gas_redeposit(
25    db: &dyn LoweringGroup,
26    function_id: ConcreteFunctionWithBodyId,
27    lowered: &mut FlatLowered,
28) {
29    if lowered.blocks.is_empty() {
30        return;
31    }
32    if matches!(db.get_flag(FlagId::new(db.upcast(), "add_withdraw_gas")),
33        Some(flag) if matches!(*flag, Flag::AddWithdrawGas(false)))
34    {
35        return;
36    }
37    let gb_ty = corelib::get_core_ty_by_name(db.upcast(), "GasBuiltin".into(), vec![]);
38    // Checking if the implicits of this function past lowering includes `GasBuiltin`.
39    if let Ok(implicits) = db.function_with_body_implicits(function_id) {
40        if !implicits.into_iter().contains(&gb_ty) {
41            return;
42        }
43    }
44    assert!(
45        lowered.parameters.iter().all(|p| lowered.variables[*p].ty != gb_ty),
46        "`GasRedeposit` stage must be called before `LowerImplicits` stage"
47    );
48
49    let redeposit_gas = corelib::get_function_id(
50        db.upcast(),
51        corelib::core_submodule(db.upcast(), "gas"),
52        "redeposit_gas".into(),
53        vec![],
54    )
55    .lowered(db);
56    let mut stack = vec![BlockId::root()];
57    let mut visited = vec![false; lowered.blocks.len()];
58    let mut redeposit_commands = vec![];
59    while let Some(block_id) = stack.pop() {
60        if visited[block_id.0] {
61            continue;
62        }
63        visited[block_id.0] = true;
64        let block = &lowered.blocks[block_id];
65        match &block.end {
66            FlatBlockEnd::Goto(block_id, _) => {
67                stack.push(*block_id);
68            }
69            FlatBlockEnd::Match { info } => {
70                let location = info.location().with_auto_generation_note(db, "withdraw_gas");
71                for arm in info.arms() {
72                    stack.push(arm.block_id);
73                    redeposit_commands.push((arm.block_id, location));
74                }
75            }
76            &FlatBlockEnd::Return(..) | FlatBlockEnd::Panic(_) => {}
77            FlatBlockEnd::NotSet => unreachable!("Block end not set"),
78        }
79    }
80    for (block_id, location) in redeposit_commands {
81        let block = &mut lowered.blocks[block_id];
82        block.statements.insert(
83            0,
84            Statement::Call(StatementCall {
85                function: redeposit_gas,
86                inputs: vec![],
87                with_coupon: false,
88                outputs: vec![],
89                location,
90            }),
91        );
92    }
93}
cairo_lang_lowering/optimizations/gas_redeposit.rs

cairo_lang_lowering/optimizations/
gas_redeposit.rs
