leo_passes/static_single_assignment/
statement.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;
18use crate::RenameTable;
19
20use leo_ast::{
21    AssertStatement,
22    AssertVariant,
23    AssignStatement,
24    Block,
25    ConditionalStatement,
26    ConstDeclaration,
27    DefinitionPlace,
28    DefinitionStatement,
29    Expression,
30    ExpressionConsumer,
31    ExpressionStatement,
32    Identifier,
33    IterationStatement,
34    Node,
35    ReturnStatement,
36    Statement,
37    StatementConsumer,
38    TernaryExpression,
39};
40use leo_span::Symbol;
41
42use indexmap::IndexSet;
43
44impl StatementConsumer for SsaFormingVisitor<'_> {
45    type Output = Vec<Statement>;
46
47    /// Consumes the expressions in an `AssertStatement`, returning the list of simplified statements.
48    fn consume_assert(&mut self, input: AssertStatement) -> Self::Output {
49        let (variant, mut statements) = match input.variant {
50            AssertVariant::Assert(expr) => {
51                let (expr, statements) = self.consume_expression_and_define(expr);
52                (AssertVariant::Assert(expr), statements)
53            }
54            AssertVariant::AssertEq(left, right) => {
55                // Reconstruct the lhs of the binary expression.
56                let (left, mut statements) = self.consume_expression_and_define(left);
57                // Reconstruct the rhs of the binary expression.
58                let (right, right_statements) = self.consume_expression_and_define(right);
59                // Accumulate any statements produced.
60                statements.extend(right_statements);
61
62                (AssertVariant::AssertEq(left, right), statements)
63            }
64            AssertVariant::AssertNeq(left, right) => {
65                // Reconstruct the lhs of the binary expression.
66                let (left, mut statements) = self.consume_expression_and_define(left);
67                // Reconstruct the rhs of the binary expression.
68                let (right, right_statements) = self.consume_expression_and_define(right);
69                // Accumulate any statements produced.
70                statements.extend(right_statements);
71
72                (AssertVariant::AssertNeq(left, right), statements)
73            }
74        };
75
76        // Add the assert statement to the list of produced statements.
77        statements.push(AssertStatement { variant, ..input }.into());
78
79        statements
80    }
81
82    /// Consume all `AssignStatement`s, renaming as necessary.
83    fn consume_assign(&mut self, mut assign: AssignStatement) -> Self::Output {
84        let (value, mut statements) = self.consume_expression(assign.value);
85        if let Expression::Path(path) = assign.place {
86            // Then assign a new unique name to the left-hand-side of the assignment.
87            // Note that this order is necessary to ensure that the right-hand-side uses the correct name when consuming a complex assignment.
88            // We really expect `path` to refer to a local variable so we only care about the result of `identifier().name`.
89            let new_place = self.rename_identifier(path.identifier());
90
91            statements.push(self.simple_definition(new_place, value));
92            statements
93        } else {
94            // It must be a sequence of accesses ending in an identifier.
95            // This loop will iterate until the identifier is reached.
96            // For example, `some_identifier[i].member` -> `some_identifier[i]` -> `some_identifier`.
97            // All we need to do is consume that identifier to possibly get a new name.
98            let mut place = &mut assign.place;
99            loop {
100                match place {
101                    Expression::ArrayAccess(array_access) => place = &mut array_access.array,
102                    Expression::MemberAccess(member_access) => place = &mut member_access.inner,
103                    Expression::TupleAccess(tuple_access) => place = &mut tuple_access.tuple,
104                    expr @ Expression::Path(..) => {
105                        let (new_expr, statements2) = self.consume_expression(std::mem::take(expr));
106                        *expr = new_expr;
107                        statements.extend(statements2);
108                        assign.value = value;
109                        statements.push(assign.into());
110                        return statements;
111                    }
112                    _ => panic!("Type checking should have ensured this is not possible."),
113                }
114            }
115        }
116    }
117
118    /// Consumes a `Block`, flattening its constituent `ConditionalStatement`s.
119    fn consume_block(&mut self, block: Block) -> Self::Output {
120        block.statements.into_iter().flat_map(|statement| self.consume_statement(statement)).collect()
121    }
122
123    /// Consumes a `ConditionalStatement`, producing phi functions (assign statements) for variables written in the then-block and otherwise-block.
124    /// For more information on phi functions, see https://en.wikipedia.org/wiki/Static_single_assignment_form.
125    /// Furthermore a new `AssignStatement` is introduced for non-trivial expressions in the condition of `ConditionalStatement`s.
126    /// For example,
127    ///   - `if x > 0 { x = x + 1 }` becomes `let $cond$0 = x > 0; if $cond$0 { x = x + 1; }`
128    ///   - `if true { x = x + 1 }` remains the same.
129    ///   - `if b { x = x + 1 }` remains the same.
130    fn consume_conditional(&mut self, conditional: ConditionalStatement) -> Self::Output {
131        // Simplify the condition and add it into the rename table.
132        let (condition, mut statements) = self.consume_expression_and_define(conditional.condition);
133
134        // Instantiate a `RenameTable` for the then-block.
135        self.push();
136
137        // Consume the then-block.
138        let then = Block {
139            span: conditional.then.span,
140            id: conditional.then.id,
141            statements: self.consume_block(conditional.then),
142        };
143
144        // Remove the `RenameTable` for the then-block.
145        let if_table = self.pop();
146
147        // Instantiate a `RenameTable` for the otherwise-block.
148        self.push();
149
150        // Consume the otherwise-block and flatten its constituent statements into the current block.
151        let otherwise = conditional.otherwise.map(|otherwise| Box::new(Statement::Block(match *otherwise {
152            Statement::Block(block) => Block {
153                span: block.span,
154                id: block.id,
155                statements: self.consume_block(block),
156            },
157            Statement::Conditional(conditional) => Block {
158                span: conditional.span,
159                id: conditional.id,
160                statements: self.consume_conditional(conditional),
161            },
162            _ => panic!("Type checking guarantees that the otherwise-block of a conditional statement is a block or another conditional statement."),
163        })));
164
165        // Remove the `RenameTable` for the otherwise-block.
166        let else_table = self.pop();
167
168        // Add reconstructed conditional statement to the list of produced statements.
169        statements.push(ConditionalStatement { condition: condition.clone(), then, otherwise, ..conditional }.into());
170
171        // Compute the write set for the variables written in the then-block or otherwise-block.
172        let if_write_set: IndexSet<&Symbol> = IndexSet::from_iter(if_table.local_names());
173        let else_write_set: IndexSet<&Symbol> = IndexSet::from_iter(else_table.local_names());
174        let write_set = if_write_set.union(&else_write_set);
175
176        // For each variable in the write set, instantiate and add a phi function to the list of produced statements.
177        for symbol in write_set {
178            // Note that phi functions only need to be instantiated if the variable exists before the `ConditionalStatement`.
179            if self.rename_table.lookup(**symbol).is_some() {
180                // Helper to lookup an and create an argument for the phi function.
181                let create_phi_argument = |table: &RenameTable, symbol: Symbol| -> Expression {
182                    let name =
183                        *table.lookup(symbol).unwrap_or_else(|| panic!("Symbol {symbol} should exist in the program."));
184                    let id = *table
185                        .lookup_id(&name)
186                        .unwrap_or_else(|| panic!("Symbol {name} should exist in the rename table."));
187                    Identifier { name, span: Default::default(), id }.into()
188                };
189
190                // Create a new name for the variable written to in the `ConditionalStatement`.
191                let new_name = self.state.assigner.unique_symbol(symbol, "$");
192
193                // Create the arguments for the phi function.
194                let if_true = create_phi_argument(&if_table, **symbol);
195                let if_false = create_phi_argument(&else_table, **symbol);
196
197                // Create a new node ID for the phi function.
198                let id = self.state.node_builder.next_id();
199                // Update the type of the node ID.
200                let Some(type_) = self.state.type_table.get(&if_true.id()) else {
201                    panic!("Type checking guarantees that all expressions have a type.");
202                };
203                self.state.type_table.insert(id, type_);
204
205                // Construct a ternary expression for the phi function.
206                let (value, stmts) = self.consume_ternary(TernaryExpression {
207                    condition: condition.clone(),
208                    if_true,
209                    if_false,
210                    span: Default::default(),
211                    id,
212                });
213
214                statements.extend(stmts);
215
216                // Get the ID for the new name of the variable.
217                let id = *self.rename_table.lookup_id(symbol).unwrap_or_else(|| {
218                    panic!("The ID for the symbol `{symbol}` should already exist in the rename table.")
219                });
220
221                // Update the `RenameTable` with the new name of the variable.
222                self.rename_table.update(**symbol, new_name, id);
223
224                // Create a new `DefinitionStatement` for the phi function.
225                let identifier = Identifier { name: new_name, span: Default::default(), id };
226                let assignment = self.simple_definition(identifier, value);
227
228                // Store the generated phi function.
229                statements.push(assignment);
230            }
231        }
232
233        statements
234    }
235
236    fn consume_const(&mut self, _: ConstDeclaration) -> Self::Output {
237        // Constants have been propagated everywhere by this point, so we no longer need const declarations.
238        Default::default()
239    }
240
241    /// Consumes the `DefinitionStatement` into an `AssignStatement`, renaming the left-hand-side as appropriate.
242    fn consume_definition(&mut self, definition: DefinitionStatement) -> Self::Output {
243        let mut statements = Vec::new();
244
245        match definition.place {
246            DefinitionPlace::Single(identifier) => {
247                // Consume the right-hand-side of the definition.
248                let (value, statements2) = self.consume_expression(definition.value);
249                statements = statements2;
250                let new_identifier = if self.rename_defs { self.rename_identifier(identifier) } else { identifier };
251                // Create a new assignment statement.
252                statements.push(self.simple_definition(new_identifier, value));
253            }
254            DefinitionPlace::Multiple(identifiers) => {
255                let new_identifiers: Vec<Identifier> = if self.rename_defs {
256                    identifiers
257                        .into_iter()
258                        .map(
259                            |identifier| if self.rename_defs { self.rename_identifier(identifier) } else { identifier },
260                        )
261                        .collect()
262                } else {
263                    identifiers
264                };
265
266                // We don't need to update the type table, as the new identifiers have
267                // the same IDs as the old ones.
268
269                // Construct the lhs of the assignment.
270                let place = DefinitionPlace::Multiple(new_identifiers);
271
272                let value = if let Expression::Call(mut call) = definition.value {
273                    for argument in call.arguments.iter_mut() {
274                        let (new_argument, new_statements) = self.consume_expression(std::mem::take(argument));
275                        *argument = new_argument;
276                        statements.extend(new_statements);
277                    }
278                    Expression::Call(call)
279                } else {
280                    let (value, new_statements) = self.consume_expression(definition.value);
281                    statements.extend(new_statements);
282                    value
283                };
284
285                // Create the definition.
286                let definition = DefinitionStatement { place, type_: None, value, ..definition }.into();
287
288                statements.push(definition);
289            }
290        }
291
292        statements
293    }
294
295    /// Consumes the expressions associated with `ExpressionStatement`, returning the simplified `ExpressionStatement`.
296    fn consume_expression_statement(&mut self, mut input: ExpressionStatement) -> Self::Output {
297        let (expr, mut statements) = self.consume_expression(input.expression);
298        input.expression = expr;
299        statements.push(input.into());
300        statements
301    }
302
303    fn consume_iteration(&mut self, _input: IterationStatement) -> Self::Output {
304        panic!("`IterationStatement`s should not be in the AST at this phase of compilation.");
305    }
306
307    /// Reconstructs the expression associated with the return statement, returning a simplified `ReturnStatement`.
308    /// Note that type checking guarantees that there is at most one `ReturnStatement` in a block.
309    fn consume_return(&mut self, mut input: ReturnStatement) -> Self::Output {
310        if let Expression::Tuple(tuple_expr) = &mut input.expression {
311            // Leave tuple expressions alone.
312            let mut statements = Vec::new();
313            for element in tuple_expr.elements.iter_mut() {
314                let (new_element, new_statements) = self.consume_expression_and_define(std::mem::take(element));
315                *element = new_element;
316                statements.extend(new_statements);
317            }
318            statements.push(input.into());
319            statements
320        } else {
321            let (expression, mut statements) = self.consume_expression_and_define(input.expression);
322            input.expression = expression;
323            statements.push(input.into());
324            statements
325        }
326    }
327}
leo_passes/static_single_assignment/statement.rs

leo_passes/static_single_assignment/
statement.rs
