cairo_lang_lowering/add_withdraw_gas/
mod.rs1use cairo_lang_diagnostics::Maybe;
2use cairo_lang_semantic::corelib::{
3 core_module, core_submodule, get_function_id, never_ty, option_none_variant,
4 option_some_variant, unit_ty,
5};
6use cairo_lang_semantic::items::constant::ConstValue;
7use cairo_lang_semantic::{ConcreteTypeId, GenericArgumentId, MatchArmSelector, TypeLongId};
8use cairo_lang_utils::{Intern, LookupIntern, extract_matches};
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, MatchEnumInfo, MatchExternInfo,
16 MatchInfo, Statement, StatementCall, VarUsage,
17};
18
19pub 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
33fn 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 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 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
88fn 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 never_ty = never_ty(db.upcast());
101 let never_var = variables.new_var(VarRequest { ty: never_ty, location });
102 lowered.variables = variables.variables;
103
104 let gas_panic_fn = get_function_id(
105 db.upcast(),
106 core_module(db.upcast()),
107 "panic_with_const_felt252".into(),
108 vec![GenericArgumentId::Constant(
109 ConstValue::Int(
110 BigInt::from_bytes_be(Sign::Plus, "Out of gas".as_bytes()),
111 db.core_info().felt252,
112 )
113 .intern(db),
114 )],
115 )
116 .lowered(db);
117
118 let never_enum_id = extract_matches!(
119 extract_matches!(never_ty.lookup_intern(db), TypeLongId::Concrete),
120 ConcreteTypeId::Enum
121 );
122
123 Ok(FlatBlock {
126 statements: vec![Statement::Call(StatementCall {
127 function: gas_panic_fn,
128 inputs: vec![],
129 with_coupon: false,
130 outputs: vec![never_var],
131 location,
132 })],
133 end: FlatBlockEnd::Match {
134 info: MatchInfo::Enum(MatchEnumInfo {
135 concrete_enum_id: never_enum_id,
136 input: VarUsage { var_id: never_var, location },
137 arms: vec![],
138 location,
139 }),
140 },
141 })
142}