cairo_lang_lowering/
utils.rs

1use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
2
3use crate::ids::LocationId;
4use crate::{
5    BlockId, FlatBlock, FlatBlockEnd, MatchArm, MatchEnumInfo, MatchEnumValue, MatchExternInfo,
6    MatchInfo, Statement, StatementCall, StatementConst, StatementDesnap, StatementEnumConstruct,
7    StatementSnapshot, StatementStructConstruct, StatementStructDestructure, VarRemapping,
8    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 {
29    fn map_var_id(&mut self, var: VariableId) -> VariableId;
30    fn map_var_usage(&mut self, var_usage: VarUsage) -> VarUsage {
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) -> LocationId {
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) {}
43    fn transform_remapping(&mut self, _remapping: &mut VarRemapping) {}
44    fn transform_end(&mut self, _end: &mut FlatBlockEnd) {}
45    fn transform_block(&mut self, _block: &mut FlatBlock) {}
46}
47
48pub trait RebuilderEx: Rebuilder {
49    /// Rebuilds the statement with renamed var and block ids.
50    fn rebuild_statement(&mut self, statement: &Statement) -> Statement {
51        let mut statement = match statement {
52            Statement::Const(stmt) => Statement::Const(StatementConst {
53                value: stmt.value.clone(),
54                output: self.map_var_id(stmt.output),
55            }),
56            Statement::Call(stmt) => Statement::Call(StatementCall {
57                function: stmt.function,
58                inputs: stmt.inputs.iter().map(|v| self.map_var_usage(*v)).collect(),
59                with_coupon: stmt.with_coupon,
60                outputs: stmt.outputs.iter().map(|v| self.map_var_id(*v)).collect(),
61                location: self.map_location(stmt.location),
62            }),
63            Statement::StructConstruct(stmt) => {
64                Statement::StructConstruct(StatementStructConstruct {
65                    inputs: stmt.inputs.iter().map(|v| self.map_var_usage(*v)).collect(),
66                    output: self.map_var_id(stmt.output),
67                })
68            }
69            Statement::StructDestructure(stmt) => {
70                Statement::StructDestructure(StatementStructDestructure {
71                    input: self.map_var_usage(stmt.input),
72                    outputs: stmt.outputs.iter().map(|v| self.map_var_id(*v)).collect(),
73                })
74            }
75            Statement::EnumConstruct(stmt) => Statement::EnumConstruct(StatementEnumConstruct {
76                variant: stmt.variant.clone(),
77                input: self.map_var_usage(stmt.input),
78                output: self.map_var_id(stmt.output),
79            }),
80            Statement::Snapshot(stmt) => Statement::Snapshot(StatementSnapshot::new(
81                self.map_var_usage(stmt.input),
82                self.map_var_id(stmt.original()),
83                self.map_var_id(stmt.snapshot()),
84            )),
85            Statement::Desnap(stmt) => Statement::Desnap(StatementDesnap {
86                input: self.map_var_usage(stmt.input),
87                output: self.map_var_id(stmt.output),
88            }),
89        };
90        self.transform_statement(&mut statement);
91        statement
92    }
93
94    /// Apply map_var_id to all the variable in the `remapping`.
95    fn rebuild_remapping(&mut self, remapping: &VarRemapping) -> VarRemapping {
96        let mut remapping = VarRemapping {
97            remapping: OrderedHashMap::from_iter(remapping.iter().map(|(dst, src_var_usage)| {
98                (self.map_var_id(*dst), self.map_var_usage(*src_var_usage))
99            })),
100        };
101        self.transform_remapping(&mut remapping);
102        remapping
103    }
104
105    /// Rebuilds the block end with renamed var and block ids.
106    fn rebuild_end(&mut self, end: &FlatBlockEnd) -> FlatBlockEnd {
107        let mut end = match end {
108            FlatBlockEnd::Return(returns, location) => FlatBlockEnd::Return(
109                returns.iter().map(|var_usage| self.map_var_usage(*var_usage)).collect(),
110                self.map_location(*location),
111            ),
112            FlatBlockEnd::Panic(data) => FlatBlockEnd::Panic(self.map_var_usage(*data)),
113            FlatBlockEnd::Goto(block_id, remapping) => {
114                FlatBlockEnd::Goto(self.map_block_id(*block_id), self.rebuild_remapping(remapping))
115            }
116            FlatBlockEnd::NotSet => unreachable!(),
117            FlatBlockEnd::Match { info } => FlatBlockEnd::Match {
118                info: match info {
119                    MatchInfo::Extern(stmt) => MatchInfo::Extern(MatchExternInfo {
120                        function: stmt.function,
121                        inputs: stmt.inputs.iter().map(|v| self.map_var_usage(*v)).collect(),
122                        arms: stmt
123                            .arms
124                            .iter()
125                            .map(|arm| MatchArm {
126                                arm_selector: arm.arm_selector.clone(),
127                                block_id: self.map_block_id(arm.block_id),
128                                var_ids: arm
129                                    .var_ids
130                                    .iter()
131                                    .map(|var_id| self.map_var_id(*var_id))
132                                    .collect(),
133                            })
134                            .collect(),
135                        location: self.map_location(stmt.location),
136                    }),
137                    MatchInfo::Enum(stmt) => MatchInfo::Enum(MatchEnumInfo {
138                        concrete_enum_id: stmt.concrete_enum_id,
139                        input: self.map_var_usage(stmt.input),
140                        arms: stmt
141                            .arms
142                            .iter()
143                            .map(|arm| MatchArm {
144                                arm_selector: arm.arm_selector.clone(),
145                                block_id: self.map_block_id(arm.block_id),
146                                var_ids: arm
147                                    .var_ids
148                                    .iter()
149                                    .map(|var_id| self.map_var_id(*var_id))
150                                    .collect(),
151                            })
152                            .collect(),
153                        location: self.map_location(stmt.location),
154                    }),
155                    MatchInfo::Value(stmt) => MatchInfo::Value(MatchEnumValue {
156                        num_of_arms: stmt.num_of_arms,
157                        input: self.map_var_usage(stmt.input),
158                        arms: stmt
159                            .arms
160                            .iter()
161                            .map(|arm| MatchArm {
162                                arm_selector: arm.arm_selector.clone(),
163                                block_id: self.map_block_id(arm.block_id),
164                                var_ids: arm
165                                    .var_ids
166                                    .iter()
167                                    .map(|var_id| self.map_var_id(*var_id))
168                                    .collect(),
169                            })
170                            .collect(),
171                        location: self.map_location(stmt.location),
172                    }),
173                },
174            },
175        };
176        self.transform_end(&mut end);
177        end
178    }
179
180    /// Rebuilds the block with renamed var and block ids.
181    fn rebuild_block(&mut self, block: &FlatBlock) -> FlatBlock {
182        let mut statements = vec![];
183        for stmt in &block.statements {
184            statements.push(self.rebuild_statement(stmt));
185        }
186        let end = self.rebuild_end(&block.end);
187        let mut block = FlatBlock { statements, end };
188        self.transform_block(&mut block);
189        block
190    }
191}
192
193impl<T: Rebuilder> RebuilderEx for T {}
cairo_lang_lowering/utils.rs

cairo_lang_lowering/
utils.rs
