leo_passes/destructuring/
visitor.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 crate::CompilerState;
18
19use leo_ast::{
20    DefinitionPlace,
21    DefinitionStatement,
22    Expression,
23    ExpressionReconstructor as _,
24    Identifier,
25    Node as _,
26    Statement,
27    TupleExpression,
28    Type,
29    TypeReconstructor,
30};
31use leo_span::Symbol;
32
33use indexmap::IndexMap;
34
35pub struct DestructuringVisitor<'a> {
36    pub state: &'a mut CompilerState,
37    /// A mapping between variables and tuple elements.
38    pub tuples: IndexMap<Symbol, Vec<Identifier>>,
39    /// Whether or not we are currently traversing an async function block.
40    pub is_async: bool,
41}
42
43impl DestructuringVisitor<'_> {
44    /// Similar to `reconstruct_expression`, except that if `expression` is of tuple type, returns it as a tuple
45    /// literal `(mem1, mem2, mem3, ...)`.
46    pub fn reconstruct_expression_tuple(&mut self, expression: Expression) -> (Expression, Vec<Statement>) {
47        let Type::Tuple(tuple_type) =
48            self.state.type_table.get(&expression.id()).expect("Expressions should have types.")
49        else {
50            // It's not a tuple, so there's no more to do.
51            return self.reconstruct_expression(expression);
52        };
53
54        let (new_expression, mut statements) = self.reconstruct_expression(expression);
55
56        match new_expression {
57            Expression::Identifier(identifier) => {
58                // It's a variable name, so just get the member identifiers we've already made.
59                let identifiers = self.tuples.get(&identifier.name).expect("Tuples should have been found");
60                let elements: Vec<Expression> =
61                    identifiers.iter().map(|identifier| Expression::Identifier(*identifier)).collect();
62
63                let tuple: Expression =
64                    TupleExpression { elements, span: Default::default(), id: self.state.node_builder.next_id() }
65                        .into();
66
67                self.state.type_table.insert(tuple.id(), Type::Tuple(tuple_type.clone()));
68
69                (tuple, statements)
70            }
71
72            tuple @ Expression::Tuple(..) => {
73                // It's already a tuple literal.
74                (tuple, statements)
75            }
76
77            expr @ Expression::Call(..) => {
78                // It's a call, so we'll need to make a new definition for the variables.
79                let definition_stmt = self.assign_tuple(expr, Symbol::intern("destructure"));
80                let Statement::Definition(DefinitionStatement {
81                    place: DefinitionPlace::Multiple(identifiers), ..
82                }) = &definition_stmt
83                else {
84                    panic!("`assign_tuple` always creates a definition with `Multiple`");
85                };
86
87                let elements = identifiers.iter().map(|identifier| Expression::Identifier(*identifier)).collect();
88
89                let expr = Expression::Tuple(TupleExpression {
90                    elements,
91                    span: Default::default(),
92                    id: self.state.node_builder.next_id(),
93                });
94
95                self.state.type_table.insert(expr.id(), Type::Tuple(tuple_type.clone()));
96
97                statements.push(definition_stmt);
98
99                (expr, statements)
100            }
101
102            _ => panic!("Tuples may only be identifiers, tuple literals, or calls."),
103        }
104    }
105
106    // Given the `expression` of tuple type, make a definition assigning variable to its members.
107    //
108    // That is, `let (mem1, mem2, mem3...) = expression;`
109    pub fn assign_tuple(&mut self, expression: Expression, name: Symbol) -> Statement {
110        let Type::Tuple(tuple_type) =
111            self.state.type_table.get(&expression.id()).expect("Expressions should have types.")
112        else {
113            panic!("assign_tuple should only be called for tuple types.");
114        };
115
116        let new_identifiers: Vec<Identifier> = (0..tuple_type.length())
117            .map(|i| {
118                let new_symbol = self.state.assigner.unique_symbol(name, format_args!("#{i}#"));
119                Identifier::new(new_symbol, self.state.node_builder.next_id())
120            })
121            .collect();
122
123        Statement::Definition(DefinitionStatement {
124            place: DefinitionPlace::Multiple(new_identifiers),
125            type_: Some(Type::Tuple(tuple_type.clone())),
126            value: expression,
127            span: Default::default(),
128            id: self.state.node_builder.next_id(),
129        })
130    }
131}
132
133impl TypeReconstructor for DestructuringVisitor<'_> {}