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

leo_passes/static_single_assignment/
statement.rs
