cairo_lang_lowering/add_withdraw_gas/
mod.rs

1use cairo_lang_diagnostics::Maybe;
2use cairo_lang_semantic::corelib::{
3    core_array_felt252_ty, core_felt252_ty, core_module, core_submodule, get_function_id,
4    get_ty_by_name, option_none_variant, option_some_variant, unit_ty,
5};
6use cairo_lang_semantic::items::constant::ConstValue;
7use cairo_lang_semantic::{GenericArgumentId, MatchArmSelector, TypeLongId};
8use cairo_lang_utils::Intern;
9use num_bigint::{BigInt, Sign};
10
11use crate::db::LoweringGroup;
12use crate::ids::{ConcreteFunctionWithBodyId, LocationId, SemanticFunctionIdEx};
13use crate::lower::context::{VarRequest, VariableAllocator};
14use crate::{
15    BlockId, FlatBlock, FlatBlockEnd, FlatLowered, MatchArm, MatchExternInfo, MatchInfo, Statement,
16    StatementCall, StatementConst, StatementStructConstruct, VarUsage,
17};
18
19/// Main function for the add_withdraw_gas lowering phase. Adds a `withdraw_gas` statement to the
20/// given function, if needed.
21pub fn add_withdraw_gas(
22    db: &dyn LoweringGroup,
23    function: ConcreteFunctionWithBodyId,
24    lowered: &mut FlatLowered,
25) -> Maybe<()> {
26    if db.needs_withdraw_gas(function)? {
27        add_withdraw_gas_to_function(db, function, lowered)?;
28    }
29
30    Ok(())
31}
32
33/// Adds a `withdraw_gas` call statement to the given function.
34/// Creates a new root block that matches on `withdraw_gas`, moves the old root block to the success
35/// arm of it, and creates a new panic block for the failure arm.
36fn add_withdraw_gas_to_function(
37    db: &dyn LoweringGroup,
38    function: ConcreteFunctionWithBodyId,
39    lowered: &mut FlatLowered,
40) -> Maybe<()> {
41    let location = LocationId::from_stable_location(db, function.stable_location(db)?)
42        .with_auto_generation_note(db, "withdraw_gas");
43    let panic_block = create_panic_block(db, function, lowered, location)?;
44
45    let old_root_block = lowered.blocks.root_block()?.clone();
46    let old_root_new_id = lowered.blocks.push(old_root_block);
47    let panic_block_id = lowered.blocks.push(panic_block);
48    let new_root_block = FlatBlock {
49        statements: vec![],
50        end: FlatBlockEnd::Match {
51            info: MatchInfo::Extern(MatchExternInfo {
52                function: get_function_id(
53                    db.upcast(),
54                    core_submodule(db.upcast(), "gas"),
55                    "withdraw_gas".into(),
56                    vec![],
57                )
58                .lowered(db),
59                inputs: vec![],
60                arms: vec![
61                    MatchArm {
62                        arm_selector: MatchArmSelector::VariantId(option_some_variant(
63                            db.upcast(),
64                            GenericArgumentId::Type(unit_ty(db.upcast())),
65                        )),
66                        block_id: old_root_new_id,
67                        var_ids: vec![],
68                    },
69                    MatchArm {
70                        arm_selector: MatchArmSelector::VariantId(option_none_variant(
71                            db.upcast(),
72                            GenericArgumentId::Type(unit_ty(db.upcast())),
73                        )),
74                        block_id: panic_block_id,
75                        var_ids: vec![],
76                    },
77                ],
78                location,
79            }),
80        },
81    };
82
83    lowered.blocks.reset_block(BlockId(0), new_root_block);
84
85    Ok(())
86}
87
88/// Creates the panic block for the case `withdraw_gas` failure.
89fn create_panic_block(
90    db: &dyn LoweringGroup,
91    function: ConcreteFunctionWithBodyId,
92    lowered: &mut FlatLowered,
93    location: LocationId,
94) -> Maybe<FlatBlock> {
95    let mut variables = VariableAllocator::new(
96        db,
97        function.function_with_body_id(db).base_semantic_function(db),
98        lowered.variables.clone(),
99    )?;
100    let new_array_var =
101        variables.new_var(VarRequest { ty: core_array_felt252_ty(db.upcast()), location });
102    let out_of_gas_err_var =
103        variables.new_var(VarRequest { ty: core_felt252_ty(db.upcast()), location });
104    let panic_instance_var = variables.new_var(VarRequest {
105        ty: get_ty_by_name(db.upcast(), core_module(db.upcast()), "Panic".into(), vec![]),
106        location,
107    });
108    let panic_data_var =
109        variables.new_var(VarRequest { ty: core_array_felt252_ty(db.upcast()), location });
110    let err_data_var = variables.new_var(VarRequest {
111        ty: TypeLongId::Tuple(vec![variables[panic_instance_var].ty, variables[panic_data_var].ty])
112            .intern(db),
113        location,
114    });
115    lowered.variables = variables.variables;
116
117    let array_module = core_submodule(db.upcast(), "array");
118
119    let add_location = |var_id| VarUsage { var_id, location };
120
121    // The block consists of creating a new array, appending 'Out of gas' to it and panic with this
122    // array as panic data.
123    Ok(FlatBlock {
124        statements: vec![
125            Statement::Call(StatementCall {
126                function: get_function_id(db.upcast(), array_module, "array_new".into(), vec![
127                    GenericArgumentId::Type(core_felt252_ty(db.upcast())),
128                ])
129                .lowered(db),
130                inputs: vec![],
131                with_coupon: false,
132                outputs: vec![new_array_var],
133                location,
134            }),
135            Statement::Const(StatementConst {
136                value: ConstValue::Int(
137                    BigInt::from_bytes_be(Sign::Plus, "Out of gas".as_bytes()),
138                    core_felt252_ty(db.upcast()),
139                ),
140                output: out_of_gas_err_var,
141            }),
142            Statement::Call(StatementCall {
143                function: get_function_id(db.upcast(), array_module, "array_append".into(), vec![
144                    GenericArgumentId::Type(core_felt252_ty(db.upcast())),
145                ])
146                .lowered(db),
147                inputs: vec![new_array_var, out_of_gas_err_var]
148                    .into_iter()
149                    .map(add_location)
150                    .collect(),
151                with_coupon: false,
152                outputs: vec![panic_data_var],
153                location,
154            }),
155            Statement::StructConstruct(StatementStructConstruct {
156                inputs: vec![],
157                output: panic_instance_var,
158            }),
159            Statement::StructConstruct(StatementStructConstruct {
160                inputs: vec![panic_instance_var, panic_data_var]
161                    .into_iter()
162                    .map(add_location)
163                    .collect(),
164                output: err_data_var,
165            }),
166        ],
167        end: FlatBlockEnd::Panic(VarUsage { var_id: err_data_var, location }),
168    })
169}
cairo_lang_lowering/add_withdraw_gas/mod.rs

cairo_lang_lowering/add_withdraw_gas/
mod.rs
