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!(
33        db.get_flag(FlagId::new(db.upcast(), "add_redeposit_gas")),
34        Some(flag) if matches!(*flag, Flag::AddRedepositGas(true))
35    ) {
36        return;
37    }
38    let gb_ty = corelib::get_core_ty_by_name(db.upcast(), "GasBuiltin".into(), vec![]);
39    // Checking if the implicits of this function past lowering includes `GasBuiltin`.
40    if let Ok(implicits) = db.function_with_body_implicits(function_id) {
41        if !implicits.into_iter().contains(&gb_ty) {
42            return;
43        }
44    }
45    assert!(
46        lowered.parameters.iter().all(|p| lowered.variables[*p].ty != gb_ty),
47        "`GasRedeposit` stage must be called before `LowerImplicits` stage"
48    );
49
50    let redeposit_gas = corelib::get_function_id(
51        db.upcast(),
52        corelib::core_submodule(db.upcast(), "gas"),
53        "redeposit_gas".into(),
54        vec![],
55    )
56    .lowered(db);
57    let mut stack = vec![BlockId::root()];
58    let mut visited = vec![false; lowered.blocks.len()];
59    let mut redeposit_commands = vec![];
60    while let Some(block_id) = stack.pop() {
61        if visited[block_id.0] {
62            continue;
63        }
64        visited[block_id.0] = true;
65        let block = &lowered.blocks[block_id];
66        match &block.end {
67            FlatBlockEnd::Goto(block_id, _) => {
68                stack.push(*block_id);
69            }
70            FlatBlockEnd::Match { info } => {
71                let location = info.location().with_auto_generation_note(db, "withdraw_gas");
72                for arm in info.arms() {
73                    stack.push(arm.block_id);
74                    redeposit_commands.push((arm.block_id, location));
75                }
76            }
77            &FlatBlockEnd::Return(..) | FlatBlockEnd::Panic(_) => {}
78            FlatBlockEnd::NotSet => unreachable!("Block end not set"),
79        }
80    }
81    for (block_id, location) in redeposit_commands {
82        let block = &mut lowered.blocks[block_id];
83        block.statements.insert(
84            0,
85            Statement::Call(StatementCall {
86                function: redeposit_gas,
87                inputs: vec![],
88                with_coupon: false,
89                outputs: vec![],
90                location,
91            }),
92        );
93    }
94}
cairo_lang_lowering/optimizations/gas_redeposit.rs

cairo_lang_lowering/optimizations/
gas_redeposit.rs
