leo_passes/destructuring/
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::DestructuringVisitor;
18
19use leo_ast::{
20    AssignStatement,
21    Block,
22    ConditionalStatement,
23    DefinitionPlace,
24    DefinitionStatement,
25    Expression,
26    ExpressionReconstructor,
27    Identifier,
28    IterationStatement,
29    Node,
30    ReturnStatement,
31    Statement,
32    StatementReconstructor,
33    Type,
34};
35use leo_span::Symbol;
36
37use itertools::{Itertools as _, izip};
38
39impl StatementReconstructor for DestructuringVisitor<'_> {
40    /// Modify assignments to tuples to become assignments to the corresponding variables.
41    ///
42    /// There are two cases we handle:
43    /// 1. An assignment to a tuple x, like `x = rhs;`.
44    ///    This we need to transform into individual assignments
45    ///    `x_i = rhs_i;`
46    ///    of the variables corresponding to members of `x` and `rhs`.
47    /// 2. An assignment to a tuple member, like `x.2[i].member = rhs;`.
48    ///    This we need to change into
49    ///    `x_2[i].member = rhs;`
50    ///    where `x_2` is the variable corresponding to `x.2`.
51    fn reconstruct_assign(&mut self, mut assign: AssignStatement) -> (Statement, Self::AdditionalOutput) {
52        let (value, mut statements) = self.reconstruct_expression(assign.value);
53
54        if let Expression::Identifier(identifier) = assign.place {
55            if let Type::Tuple(..) = self.state.type_table.get(&value.id()).expect("Expressions should have types.") {
56                // This is the first case, assigning to a variable of tuple type.
57                let identifiers = self.tuples.get(&identifier.name).expect("Tuple should have been encountered.");
58
59                let Expression::Identifier(rhs) = value else {
60                    panic!("SSA should have ensured this is an identifier.");
61                };
62
63                let rhs_identifiers = self.tuples.get(&rhs.name).expect("Tuple should have been encountered.");
64
65                // Again, make an assignment for each identifier.
66                for (&identifier, &rhs_identifier) in identifiers.iter().zip_eq(rhs_identifiers) {
67                    let stmt = AssignStatement {
68                        place: identifier.into(),
69                        value: rhs_identifier.into(),
70                        id: self.state.node_builder.next_id(),
71                        span: Default::default(),
72                    }
73                    .into();
74
75                    statements.push(stmt);
76                }
77
78                // We don't need the original assignment, just the ones we've created.
79                return (Statement::dummy(), statements);
80            }
81        }
82
83        // We need to check for case 2, so we loop and see if we find a tuple access.
84
85        assign.value = value;
86        let mut place = &mut assign.place;
87
88        loop {
89            // Loop through the places in the assignment to the top-level expression until an identifier or tuple access is reached.
90            match place {
91                Expression::TupleAccess(access) => {
92                    // We're assigning to a tuple member, case 2 mentioned above.
93                    let Expression::Identifier(identifier) = &access.tuple else {
94                        panic!("SSA should have ensured this is an identifier.");
95                    };
96
97                    let tuple_ids = self.tuples.get(&identifier.name).expect("Tuple should have been encountered.");
98
99                    // This is the corresponding variable name of the member we're assigning to.
100                    let identifier = tuple_ids[access.index.value()];
101
102                    *place = identifier.into();
103
104                    return (assign.into(), statements);
105                }
106
107                Expression::ArrayAccess(access) => {
108                    // We need to investigate the array, as maybe it's inside a tuple access, like `tupl.0[1u8]`.
109                    place = &mut access.array;
110                }
111
112                Expression::MemberAccess(access) => {
113                    // We need to investigate the struct, as maybe it's inside a tuple access, like `tupl.0.mem`.
114                    place = &mut access.inner;
115                }
116
117                Expression::Identifier(..) => {
118                    // There was no tuple access, so this is neither case 1 nor 2; there's nothing to do.
119                    return (assign.into(), statements);
120                }
121
122                _ => panic!("Type checking should have prevented this."),
123            }
124        }
125    }
126
127    fn reconstruct_block(&mut self, block: Block) -> (Block, Self::AdditionalOutput) {
128        let mut statements = Vec::with_capacity(block.statements.len());
129
130        // Reconstruct the statements in the block, accumulating any additional statements.
131        for statement in block.statements {
132            let (reconstructed_statement, additional_statements) = self.reconstruct_statement(statement);
133            statements.extend(additional_statements);
134            if !reconstructed_statement.is_empty() {
135                statements.push(reconstructed_statement);
136            }
137        }
138
139        (Block { statements, ..block }, Default::default())
140    }
141
142    fn reconstruct_conditional(&mut self, input: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
143        let (condition, mut statements) = self.reconstruct_expression(input.condition);
144        let (then, statements2) = self.reconstruct_block(input.then);
145        statements.extend(statements2);
146        let otherwise = input.otherwise.map(|oth| {
147            let (expr, statements3) = self.reconstruct_statement(*oth);
148            statements.extend(statements3);
149            Box::new(expr)
150        });
151        (ConditionalStatement { condition, then, otherwise, ..input }.into(), statements)
152    }
153
154    fn reconstruct_definition(&mut self, definition: DefinitionStatement) -> (Statement, Self::AdditionalOutput) {
155        use DefinitionPlace::*;
156
157        let make_identifiers = |slf: &mut Self, single: Symbol, count: usize| -> Vec<Identifier> {
158            (0..count)
159                .map(|i| {
160                    Identifier::new(
161                        slf.state.assigner.unique_symbol(format_args!("{single}#tuple{i}"), "$"),
162                        slf.state.node_builder.next_id(),
163                    )
164                })
165                .collect()
166        };
167
168        let (value, mut statements) = self.reconstruct_expression(definition.value);
169        let ty = self.state.type_table.get(&value.id()).expect("Expressions should have a type.");
170        match (definition.place, value, ty) {
171            (Single(identifier), Expression::Identifier(rhs), Type::Tuple(tuple_type)) => {
172                // We need to give the members new names, in case they are assigned to.
173                let identifiers = make_identifiers(self, identifier.name, tuple_type.length());
174
175                let rhs_identifiers = self.tuples.get(&rhs.name).unwrap();
176
177                for (identifier, rhs_identifier, ty) in izip!(&identifiers, rhs_identifiers, tuple_type.elements()) {
178                    // Make a definition for each.
179                    let stmt = DefinitionStatement {
180                        place: Single(*identifier),
181                        type_: Some(ty.clone()),
182                        value: Expression::Identifier(*rhs_identifier),
183                        span: Default::default(),
184                        id: self.state.node_builder.next_id(),
185                    }
186                    .into();
187                    statements.push(stmt);
188
189                    // Put each into the type table.
190                    self.state.type_table.insert(identifier.id(), ty.clone());
191                }
192
193                // Put the identifier in `self.tuples`. We don't need to keep our definition.
194                self.tuples.insert(identifier.name, identifiers);
195                (Statement::dummy(), statements)
196            }
197            (Single(identifier), Expression::Tuple(tuple), Type::Tuple(tuple_type)) => {
198                // Name each of the expressions on the right.
199                let identifiers = make_identifiers(self, identifier.name, tuple_type.length());
200
201                for (identifier, expr, ty) in izip!(&identifiers, tuple.elements, tuple_type.elements()) {
202                    // Make a definition for each.
203                    let stmt = DefinitionStatement {
204                        place: Single(*identifier),
205                        type_: Some(ty.clone()),
206                        value: expr,
207                        span: Default::default(),
208                        id: self.state.node_builder.next_id(),
209                    }
210                    .into();
211                    statements.push(stmt);
212
213                    // Put each into the type table.
214                    self.state.type_table.insert(identifier.id(), ty.clone());
215                }
216
217                // Put the identifier in `self.tuples`. We don't need to keep our definition.
218                self.tuples.insert(identifier.name, identifiers);
219                (Statement::dummy(), statements)
220            }
221            (Single(identifier), rhs @ Expression::Call(..), Type::Tuple(tuple_type)) => {
222                let definition_stmt = self.assign_tuple(rhs, identifier.name);
223
224                let Statement::Definition(DefinitionStatement {
225                    place: DefinitionPlace::Multiple(identifiers), ..
226                }) = &definition_stmt
227                else {
228                    panic!("assign_tuple creates `Multiple`.");
229                };
230
231                // Put it into `self.tuples`.
232                self.tuples.insert(identifier.name, identifiers.clone());
233
234                // Put each into the type table.
235                for (identifier, ty) in identifiers.iter().zip(tuple_type.elements()) {
236                    self.state.type_table.insert(identifier.id(), ty.clone());
237                }
238
239                (definition_stmt, statements)
240            }
241            (Multiple(identifiers), Expression::Tuple(tuple), Type::Tuple(..)) => {
242                // Just make a definition for each tuple element.
243                for (identifier, expr) in identifiers.into_iter().zip_eq(tuple.elements) {
244                    let stmt = DefinitionStatement {
245                        place: Single(identifier),
246                        type_: None,
247                        value: expr,
248                        span: Default::default(),
249                        id: self.state.node_builder.next_id(),
250                    }
251                    .into();
252                    statements.push(stmt);
253                }
254
255                // We don't need to keep the original definition.
256                (Statement::dummy(), statements)
257            }
258            (Multiple(identifiers), Expression::Identifier(rhs), Type::Tuple(..)) => {
259                // Again, make a definition for each tuple element.
260                let rhs_identifiers = self.tuples.get(&rhs.name).expect("We should have encountered this tuple by now");
261                for (identifier, rhs_identifier) in identifiers.into_iter().zip_eq(rhs_identifiers.iter()) {
262                    let stmt = DefinitionStatement {
263                        place: Single(identifier),
264                        type_: None,
265                        value: Expression::Identifier(*rhs_identifier),
266                        span: Default::default(),
267                        id: self.state.node_builder.next_id(),
268                    }
269                    .into();
270                    statements.push(stmt);
271                }
272
273                // We don't need to keep the original definition.
274                (Statement::dummy(), statements)
275            }
276            (m @ Multiple(..), value @ Expression::Call(..), Type::Tuple(..)) => {
277                // Just reconstruct the statement.
278                let stmt =
279                    DefinitionStatement { place: m, type_: None, value, span: definition.span, id: definition.id }
280                        .into();
281                (stmt, statements)
282            }
283            (_, _, Type::Tuple(..)) => {
284                panic!("Expressions of tuple type can only be tuple literals, identifiers, or calls.");
285            }
286            (s @ Single(..), rhs, _) => {
287                // This isn't a tuple. Just build the definition again.
288                let stmt = DefinitionStatement {
289                    place: s,
290                    type_: None,
291                    value: rhs,
292                    span: Default::default(),
293                    id: definition.id,
294                }
295                .into();
296                (stmt, statements)
297            }
298            (Multiple(_), _, _) => panic!("A definition with multiple identifiers must have tuple type"),
299        }
300    }
301
302    fn reconstruct_iteration(&mut self, _: IterationStatement) -> (Statement, Self::AdditionalOutput) {
303        panic!("`IterationStatement`s should not be in the AST at this phase of compilation.");
304    }
305
306    fn reconstruct_return(&mut self, input: ReturnStatement) -> (Statement, Self::AdditionalOutput) {
307        let (expression, statements) = self.reconstruct_expression_tuple(input.expression);
308        (ReturnStatement { expression, ..input }.into(), statements)
309    }
310}
leo_passes/destructuring/statement.rs

leo_passes/destructuring/
statement.rs
