cairo_lang_lowering/
utils.rs

1use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
2
3use crate::ids::LocationId;
4use crate::{
5    Block, BlockEnd, BlockId, MatchArm, MatchEnumInfo, MatchEnumValue, MatchExternInfo, MatchInfo,
6    Statement, StatementCall, StatementConst, StatementDesnap, StatementEnumConstruct,
7    StatementIntoBox, StatementSnapshot, StatementStructConstruct, StatementStructDestructure,
8    StatementUnbox, VarRemapping, VarUsage, VariableId,
9};
10
11/// Options for the `inlining-strategy` arguments.
12#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
13pub enum InliningStrategy {
14    /// Do not override inlining strategy.
15    ///
16    /// Note: equivalent to `InlineSmallFunctions(DEFAULT_INLINE_SMALL_FUNCTIONS_THRESHOLD)`.
17    #[default]
18    Default,
19    /// Should inline small functions up to the given weight.
20    ///
21    /// Note: the weight exact definition is subject to change.
22    InlineSmallFunctions(usize),
23    /// Inline only in the case of an `inline(always)` annotation.
24    Avoid,
25}
26
27/// A rebuilder trait for rebuilding lowered representation.
28pub trait Rebuilder<'db> {
29    fn map_var_id(&mut self, var: VariableId) -> VariableId;
30    fn map_var_usage(&mut self, var_usage: VarUsage<'db>) -> VarUsage<'db> {
31        VarUsage {
32            var_id: self.map_var_id(var_usage.var_id),
33            location: self.map_location(var_usage.location),
34        }
35    }
36    fn map_location(&mut self, location: LocationId<'db>) -> LocationId<'db> {
37        location
38    }
39    fn map_block_id(&mut self, block: BlockId) -> BlockId {
40        block
41    }
42    fn transform_statement(&mut self, _statement: &mut Statement<'db>) {}
43    fn transform_remapping(&mut self, _remapping: &mut VarRemapping<'db>) {}
44    fn transform_end(&mut self, _end: &mut BlockEnd<'db>) {}
45    fn transform_block(&mut self, _block: &mut Block<'db>) {}
46}
47
48pub trait RebuilderEx<'db>: Rebuilder<'db> {
49    /// Rebuilds the statement with renamed var and block ids.
50    fn rebuild_statement(&mut self, statement: &Statement<'db>) -> Statement<'db> {
51        let mut statement = match statement {
52            Statement::Const(stmt) => Statement::Const(StatementConst::new(
53                stmt.value,
54                self.map_var_id(stmt.output),
55                stmt.boxed,
56            )),
57            Statement::Call(stmt) => Statement::Call(StatementCall {
58                function: stmt.function,
59                inputs: stmt.inputs.iter().map(|v| self.map_var_usage(*v)).collect(),
60                with_coupon: stmt.with_coupon,
61                outputs: stmt.outputs.iter().map(|v| self.map_var_id(*v)).collect(),
62                location: self.map_location(stmt.location),
63                is_specialization_base_call: stmt.is_specialization_base_call,
64            }),
65            Statement::StructConstruct(stmt) => {
66                Statement::StructConstruct(StatementStructConstruct {
67                    inputs: stmt.inputs.iter().map(|v| self.map_var_usage(*v)).collect(),
68                    output: self.map_var_id(stmt.output),
69                })
70            }
71            Statement::StructDestructure(stmt) => {
72                Statement::StructDestructure(StatementStructDestructure {
73                    input: self.map_var_usage(stmt.input),
74                    outputs: stmt.outputs.iter().map(|v| self.map_var_id(*v)).collect(),
75                })
76            }
77            Statement::EnumConstruct(stmt) => Statement::EnumConstruct(StatementEnumConstruct {
78                variant: stmt.variant,
79                input: self.map_var_usage(stmt.input),
80                output: self.map_var_id(stmt.output),
81            }),
82            Statement::Snapshot(stmt) => Statement::Snapshot(StatementSnapshot::new(
83                self.map_var_usage(stmt.input),
84                self.map_var_id(stmt.original()),
85                self.map_var_id(stmt.snapshot()),
86            )),
87            Statement::Desnap(stmt) => Statement::Desnap(StatementDesnap {
88                input: self.map_var_usage(stmt.input),
89                output: self.map_var_id(stmt.output),
90            }),
91            Statement::IntoBox(stmt) => Statement::IntoBox(StatementIntoBox {
92                input: self.map_var_usage(stmt.input),
93                output: self.map_var_id(stmt.output),
94            }),
95            Statement::Unbox(stmt) => Statement::Unbox(StatementUnbox {
96                input: self.map_var_usage(stmt.input),
97                output: self.map_var_id(stmt.output),
98            }),
99        };
100        self.transform_statement(&mut statement);
101        statement
102    }
103
104    /// Apply map_var_id to all the variables in the `remapping`.
105    fn rebuild_remapping(&mut self, remapping: &VarRemapping<'db>) -> VarRemapping<'db> {
106        let mut remapping = VarRemapping {
107            remapping: OrderedHashMap::from_iter(remapping.iter().map(|(dst, src_var_usage)| {
108                (self.map_var_id(*dst), self.map_var_usage(*src_var_usage))
109            })),
110        };
111        self.transform_remapping(&mut remapping);
112        remapping
113    }
114
115    /// Rebuilds the block end with renamed var and block ids.
116    fn rebuild_end(&mut self, end: &BlockEnd<'db>) -> BlockEnd<'db> {
117        let mut end = match end {
118            BlockEnd::Return(returns, location) => BlockEnd::Return(
119                returns.iter().map(|var_usage| self.map_var_usage(*var_usage)).collect(),
120                self.map_location(*location),
121            ),
122            BlockEnd::Panic(data) => BlockEnd::Panic(self.map_var_usage(*data)),
123            BlockEnd::Goto(block_id, remapping) => {
124                BlockEnd::Goto(self.map_block_id(*block_id), self.rebuild_remapping(remapping))
125            }
126            BlockEnd::NotSet => unreachable!(),
127            BlockEnd::Match { info } => BlockEnd::Match {
128                info: match info {
129                    MatchInfo::Extern(stmt) => MatchInfo::Extern(MatchExternInfo {
130                        function: stmt.function,
131                        inputs: stmt.inputs.iter().map(|v| self.map_var_usage(*v)).collect(),
132                        arms: stmt
133                            .arms
134                            .iter()
135                            .map(|arm| MatchArm {
136                                arm_selector: arm.arm_selector.clone(),
137                                block_id: self.map_block_id(arm.block_id),
138                                var_ids: arm
139                                    .var_ids
140                                    .iter()
141                                    .map(|var_id| self.map_var_id(*var_id))
142                                    .collect(),
143                            })
144                            .collect(),
145                        location: self.map_location(stmt.location),
146                    }),
147                    MatchInfo::Enum(stmt) => MatchInfo::Enum(MatchEnumInfo {
148                        concrete_enum_id: stmt.concrete_enum_id,
149                        input: self.map_var_usage(stmt.input),
150                        arms: stmt
151                            .arms
152                            .iter()
153                            .map(|arm| MatchArm {
154                                arm_selector: arm.arm_selector.clone(),
155                                block_id: self.map_block_id(arm.block_id),
156                                var_ids: arm
157                                    .var_ids
158                                    .iter()
159                                    .map(|var_id| self.map_var_id(*var_id))
160                                    .collect(),
161                            })
162                            .collect(),
163                        location: self.map_location(stmt.location),
164                    }),
165                    MatchInfo::Value(stmt) => MatchInfo::Value(MatchEnumValue {
166                        num_of_arms: stmt.num_of_arms,
167                        input: self.map_var_usage(stmt.input),
168                        arms: stmt
169                            .arms
170                            .iter()
171                            .map(|arm| MatchArm {
172                                arm_selector: arm.arm_selector.clone(),
173                                block_id: self.map_block_id(arm.block_id),
174                                var_ids: arm
175                                    .var_ids
176                                    .iter()
177                                    .map(|var_id| self.map_var_id(*var_id))
178                                    .collect(),
179                            })
180                            .collect(),
181                        location: self.map_location(stmt.location),
182                    }),
183                },
184            },
185        };
186        self.transform_end(&mut end);
187        end
188    }
189
190    /// Rebuilds the block with renamed var and block ids.
191    fn rebuild_block(&mut self, block: &Block<'db>) -> Block<'db> {
192        let statements = block.statements.iter().map(|stmt| self.rebuild_statement(stmt)).collect();
193        let end = self.rebuild_end(&block.end);
194        let mut block = Block { statements, end };
195        self.transform_block(&mut block);
196        block
197    }
198}
199
200impl<'db, T: Rebuilder<'db>> RebuilderEx<'db> for T {}
cairo_lang_lowering/utils.rs

cairo_lang_lowering/
utils.rs
