cairo_lang_lowering/optimizations/
gas_redeposit.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
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,
            }),
        );
    }
}