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    AstReconstructor,
21    DefinitionPlace,
22    DefinitionStatement,
23    Expression,
24    Identifier,
25    Node as _,
26    Path,
27    Statement,
28    TupleExpression,
29    Type,
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::Path(path) => {
58                // It's a variable name, so just get the member identifiers we've already made.
59                let identifiers = self.tuples.get(&path.identifier().name).expect("Tuples should have been found");
60                let elements: Vec<Expression> =
61                    identifiers.iter().map(|identifier| Path::from(*identifier).into_absolute().into()).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 =
88                    identifiers.iter().map(|identifier| Path::from(*identifier).into_absolute().into()).collect();
89
90                let expr = Expression::Tuple(TupleExpression {
91                    elements,
92                    span: Default::default(),
93                    id: self.state.node_builder.next_id(),
94                });
95
96                self.state.type_table.insert(expr.id(), Type::Tuple(tuple_type.clone()));
97
98                statements.push(definition_stmt);
99
100                (expr, statements)
101            }
102
103            _ => panic!("Tuples may only be identifiers, tuple literals, or calls."),
104        }
105    }
106
107    // Given the `expression` of tuple type, make a definition assigning variable to its members.
108    //
109    // That is, `let (mem1, mem2, mem3...) = expression;`
110    pub fn assign_tuple(&mut self, expression: Expression, name: Symbol) -> Statement {
111        let Type::Tuple(tuple_type) =
112            self.state.type_table.get(&expression.id()).expect("Expressions should have types.")
113        else {
114            panic!("assign_tuple should only be called for tuple types.");
115        };
116
117        let new_identifiers: Vec<Identifier> = (0..tuple_type.length())
118            .map(|i| {
119                let new_symbol = self.state.assigner.unique_symbol(name, format_args!("#{i}#"));
120                Identifier::new(new_symbol, self.state.node_builder.next_id())
121            })
122            .collect();
123
124        Statement::Definition(DefinitionStatement {
125            place: DefinitionPlace::Multiple(new_identifiers),
126            type_: Some(Type::Tuple(tuple_type.clone())),
127            value: expression,
128            span: Default::default(),
129            id: self.state.node_builder.next_id(),
130        })
131    }
132}