leo_passes/static_single_assignment/
expression.rs

1// Copyright (C) 2019-2025 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17use super::SsaFormingVisitor;
18
19use leo_ast::{
20    ArrayAccess,
21    ArrayExpression,
22    AssociatedFunctionExpression,
23    BinaryExpression,
24    CallExpression,
25    CastExpression,
26    Composite,
27    Expression,
28    ExpressionConsumer,
29    Identifier,
30    Literal,
31    Location,
32    LocatorExpression,
33    MemberAccess,
34    RepeatExpression,
35    Statement,
36    StructExpression,
37    StructVariableInitializer,
38    TernaryExpression,
39    TupleAccess,
40    TupleExpression,
41    UnaryExpression,
42    UnitExpression,
43};
44use leo_span::{Symbol, sym};
45
46use indexmap::IndexMap;
47
48impl SsaFormingVisitor<'_> {
49    /// Consume this expression and assign it to a variable (unless it's already an Identifier).
50    pub fn consume_expression_and_define(&mut self, input: Expression) -> (Expression, Vec<Statement>) {
51        let (expr, mut statements) = self.consume_expression(input);
52        if matches!(expr, Expression::Identifier(..) | Expression::Unit(..) | Expression::Err(..)) {
53            (expr, statements)
54        } else {
55            let (place, statement) = self.unique_simple_definition(expr);
56            statements.push(statement);
57            (place.into(), statements)
58        }
59    }
60}
61
62impl ExpressionConsumer for SsaFormingVisitor<'_> {
63    type Output = (Expression, Vec<Statement>);
64
65    fn consume_array_access(&mut self, input: ArrayAccess) -> Self::Output {
66        let (array, statements) = self.consume_expression_and_define(input.array);
67        (ArrayAccess { array, ..input }.into(), statements)
68    }
69
70    fn consume_member_access(&mut self, input: MemberAccess) -> Self::Output {
71        // If the access expression is of the form `self.<name>`, then don't rename it.
72        if let Expression::Identifier(Identifier { name, .. }) = input.inner {
73            if name == sym::SelfLower {
74                return (input.into(), Vec::new());
75            }
76        }
77
78        let (inner, statements) = self.consume_expression_and_define(input.inner);
79        (MemberAccess { inner, ..input }.into(), statements)
80    }
81
82    fn consume_tuple_access(&mut self, input: TupleAccess) -> Self::Output {
83        let (tuple, statements) = self.consume_expression_and_define(input.tuple);
84        (TupleAccess { tuple, ..input }.into(), statements)
85    }
86
87    /// Consumes an array expression, accumulating any statements that are generated.
88    fn consume_array(&mut self, input: ArrayExpression) -> Self::Output {
89        let mut statements = Vec::new();
90
91        // Process the elements, accumulating any statements produced.
92        let elements = input
93            .elements
94            .into_iter()
95            .map(|element| {
96                let (element, mut stmts) = self.consume_expression_and_define(element);
97                statements.append(&mut stmts);
98                element
99            })
100            .collect();
101
102        (ArrayExpression { elements, ..input }.into(), statements)
103    }
104
105    /// Consumes a binary expression, accumulating any statements that are generated.
106    fn consume_binary(&mut self, input: BinaryExpression) -> Self::Output {
107        // Reconstruct the lhs of the binary expression.
108        let (left, mut statements) = self.consume_expression_and_define(input.left);
109        // Reconstruct the rhs of the binary expression.
110        let (right, mut right_statements) = self.consume_expression_and_define(input.right);
111        // Accumulate any statements produced.
112        statements.append(&mut right_statements);
113
114        (BinaryExpression { left, right, ..input }.into(), statements)
115    }
116
117    /// Consumes a call expression without visiting the function name, accumulating any statements that are generated.
118    fn consume_call(&mut self, input: CallExpression) -> Self::Output {
119        let mut statements = Vec::new();
120
121        // Process the arguments, accumulating any statements produced.
122        let arguments = input
123            .arguments
124            .into_iter()
125            .map(|argument| {
126                let (argument, mut stmts) = self.consume_expression_and_define(argument);
127                statements.append(&mut stmts);
128                argument
129            })
130            .collect();
131
132        (
133            CallExpression {
134                // Note that we do not rename the function name.
135                arguments,
136                ..input
137            }
138            .into(),
139            statements,
140        )
141    }
142
143    /// Consumes a cast expression, accumulating any statements that are generated.
144    fn consume_cast(&mut self, input: CastExpression) -> Self::Output {
145        // Reconstruct the expression being casted.
146        let (expression, statements) = self.consume_expression_and_define(input.expression);
147        (CastExpression { expression, ..input }.into(), statements)
148    }
149
150    /// Consumes a struct initialization expression with renamed variables, accumulating any statements that are generated.
151    fn consume_struct_init(&mut self, input: StructExpression) -> Self::Output {
152        let mut statements = Vec::new();
153
154        // Process the members, accumulating any statements produced.
155        let members: Vec<StructVariableInitializer> = input
156            .members
157            .into_iter()
158            .map(|arg| {
159                let (expression, mut stmts) = if let Some(expr) = arg.expression {
160                    self.consume_expression_and_define(expr)
161                } else {
162                    self.consume_identifier(arg.identifier)
163                };
164                // Accumulate any statements produced.
165                statements.append(&mut stmts);
166
167                // Return the new member.
168                StructVariableInitializer { expression: Some(expression), ..arg }
169            })
170            .collect();
171
172        // Reorder the members to match that of the struct definition.
173
174        // Lookup the struct definition.
175        let struct_definition: &Composite = self
176            .state
177            .symbol_table
178            .lookup_record(Location::new(self.program, input.name.name))
179            .or_else(|| self.state.symbol_table.lookup_struct(input.name.name))
180            .expect("Type checking guarantees this definition exists.");
181
182        // Initialize the list of reordered members.
183        let mut reordered_members = Vec::with_capacity(members.len());
184
185        // Collect the members of the init expression into a map.
186        let mut member_map: IndexMap<Symbol, StructVariableInitializer> =
187            members.into_iter().map(|member| (member.identifier.name, member)).collect();
188
189        // If we are initializing a record, add the `owner` first.
190        // Note that type checking guarantees that the above fields exist.
191        if struct_definition.is_record {
192            // Add the `owner` field.
193            // Note that the `unwrap` is safe, since type checking guarantees that the member exists.
194            reordered_members.push(member_map.shift_remove(&sym::owner).unwrap());
195        }
196
197        // For each member of the struct definition, push the corresponding member of the init expression.
198        for member in &struct_definition.members {
199            // If the member is part of a record and it is `owner` then we have already added it.
200            if !(struct_definition.is_record && matches!(member.identifier.name, sym::owner)) {
201                // Lookup and push the member of the init expression.
202                // Note that the `unwrap` is safe, since type checking guarantees that the member exists.
203                reordered_members.push(member_map.shift_remove(&member.identifier.name).unwrap());
204            }
205        }
206
207        (StructExpression { members: reordered_members, ..input }.into(), statements)
208    }
209
210    /// Retrieve the new name for this `Identifier`.
211    ///
212    /// Note that this shouldn't be used for `Identifier`s on the lhs of definitions or
213    /// assignments.
214    fn consume_identifier(&mut self, identifier: Identifier) -> Self::Output {
215        // If lookup fails, either it's the name of a mapping or we didn't rename it.
216        let name = *self.rename_table.lookup(identifier.name).unwrap_or(&identifier.name);
217        (Identifier { name, ..identifier }.into(), Default::default())
218    }
219
220    /// Consumes and returns the literal without making any modifications.
221    fn consume_literal(&mut self, input: Literal) -> Self::Output {
222        (input.into(), Default::default())
223    }
224
225    /// Consumes and returns the locator expression without making any modifications
226    fn consume_locator(&mut self, input: LocatorExpression) -> Self::Output {
227        (input.into(), Vec::new())
228    }
229
230    fn consume_repeat(&mut self, input: RepeatExpression) -> Self::Output {
231        let (expr, statements) = self.consume_expression_and_define(input.expr);
232
233        // By now, the repeat count should be a literal. So we just ignore it. There is no need to SSA it.
234        (RepeatExpression { expr, ..input }.into(), statements)
235    }
236
237    /// Consumes a ternary expression, accumulating any statements that are generated.
238    fn consume_ternary(&mut self, input: TernaryExpression) -> Self::Output {
239        // Reconstruct the condition of the ternary expression.
240        let (cond_expr, mut statements) = self.consume_expression_and_define(input.condition);
241        // Reconstruct the if-true case of the ternary expression.
242        let (if_true_expr, if_true_statements) = self.consume_expression_and_define(input.if_true);
243        // Reconstruct the if-false case of the ternary expression.
244        let (if_false_expr, if_false_statements) = self.consume_expression_and_define(input.if_false);
245
246        // Accumulate any statements produced.
247        statements.extend(if_true_statements);
248        statements.extend(if_false_statements);
249
250        (
251            TernaryExpression { condition: cond_expr, if_true: if_true_expr, if_false: if_false_expr, ..input }.into(),
252            statements,
253        )
254    }
255
256    /// Consumes a tuple expression, accumulating any statements that are generated
257    fn consume_tuple(&mut self, input: TupleExpression) -> Self::Output {
258        let mut statements = Vec::new();
259
260        // Process the elements, accumulating any statements produced.
261        let elements = input
262            .elements
263            .into_iter()
264            .map(|element| {
265                let (element, mut stmts) = self.consume_expression_and_define(element);
266                statements.append(&mut stmts);
267                element
268            })
269            .collect();
270
271        (TupleExpression { elements, ..input }.into(), statements)
272    }
273
274    /// Consumes a unary expression, accumulating any statements that are generated.
275    fn consume_unary(&mut self, input: UnaryExpression) -> Self::Output {
276        // Reconstruct the operand of the unary expression.
277        let (receiver, statements) = self.consume_expression_and_define(input.receiver);
278        (UnaryExpression { receiver, ..input }.into(), statements)
279    }
280
281    fn consume_unit(&mut self, input: UnitExpression) -> Self::Output {
282        (input.into(), Default::default())
283    }
284
285    fn consume_associated_constant(&mut self, input: leo_ast::AssociatedConstantExpression) -> Self::Output {
286        (input.into(), Default::default())
287    }
288
289    fn consume_associated_function(&mut self, input: leo_ast::AssociatedFunctionExpression) -> Self::Output {
290        let mut statements = Vec::new();
291        let expr = AssociatedFunctionExpression {
292            arguments: input
293                .arguments
294                .into_iter()
295                .map(|arg| {
296                    let (arg, mut stmts) = self.consume_expression_and_define(arg);
297                    statements.append(&mut stmts);
298                    arg
299                })
300                .collect(),
301            ..input
302        }
303        .into();
304
305        (expr, statements)
306    }
307}