leo_passes/destructuring/
expression.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    ArrayAccess,
21    Expression,
22    ExpressionReconstructor,
23    Identifier,
24    IntegerType,
25    Literal,
26    Node as _,
27    Statement,
28    TernaryExpression,
29    TupleAccess,
30    TupleExpression,
31    Type,
32};
33
34use itertools::izip;
35
36impl ExpressionReconstructor for DestructuringVisitor<'_> {
37    type AdditionalOutput = Vec<Statement>;
38
39    /// Replaces a tuple access expression with the appropriate expression.
40    fn reconstruct_tuple_access(&mut self, input: TupleAccess) -> (Expression, Self::AdditionalOutput) {
41        let Expression::Identifier(identifier) = &input.tuple else {
42            panic!("SSA guarantees that subexpressions are identifiers or literals.");
43        };
44
45        // Look up the expression in the tuple map.
46        match self.tuples.get(&identifier.name).and_then(|tuple_names| tuple_names.get(input.index.value())) {
47            Some(id) => ((*id).into(), Default::default()),
48            None => {
49                if !matches!(self.state.type_table.get(&identifier.id), Some(Type::Future(_))) {
50                    panic!("Type checking guarantees that all tuple accesses are declared and indices are valid.");
51                }
52
53                let index = Literal::integer(
54                    IntegerType::U32,
55                    input.index.to_string(),
56                    input.span,
57                    self.state.node_builder.next_id(),
58                );
59                self.state.type_table.insert(index.id(), Type::Integer(IntegerType::U32));
60
61                let expr =
62                    ArrayAccess { array: (*identifier).into(), index: index.into(), span: input.span, id: input.id }
63                        .into();
64
65                (expr, Default::default())
66            }
67        }
68    }
69
70    /// If this is a ternary expression on tuples of length `n`, we'll need to change it into
71    /// `n` ternary expressions on the members.
72    fn reconstruct_ternary(&mut self, mut input: TernaryExpression) -> (Expression, Self::AdditionalOutput) {
73        let (if_true, mut statements) = self.reconstruct_expression_tuple(std::mem::take(&mut input.if_true));
74        let (if_false, statements2) = self.reconstruct_expression_tuple(std::mem::take(&mut input.if_false));
75        statements.extend(statements2);
76
77        match (if_true, if_false) {
78            (Expression::Tuple(tuple_true), Expression::Tuple(tuple_false)) => {
79                // Aleo's `ternary` opcode doesn't know about tuples, so we have to handle this.
80                let Some(Type::Tuple(tuple_type)) = self.state.type_table.get(&tuple_true.id()) else {
81                    panic!("Should have tuple type");
82                };
83
84                // We'll be reusing `input.condition`, so assign it to a variable.
85                let cond = if let Expression::Identifier(..) = input.condition {
86                    input.condition
87                } else {
88                    let place = Identifier::new(
89                        self.state.assigner.unique_symbol("cond", "$$"),
90                        self.state.node_builder.next_id(),
91                    );
92
93                    let definition = self.state.assigner.simple_definition(
94                        place,
95                        input.condition,
96                        self.state.node_builder.next_id(),
97                    );
98
99                    statements.push(definition);
100
101                    self.state.type_table.insert(place.id(), Type::Boolean);
102
103                    Expression::Identifier(place)
104                };
105
106                // These will be the `elements` of our resulting tuple.
107                let mut elements = Vec::with_capacity(tuple_true.elements.len());
108
109                // Create an individual `ternary` for each tuple member and assign the
110                // result to a new variable.
111                for (i, (lhs, rhs, ty)) in
112                    izip!(tuple_true.elements, tuple_false.elements, tuple_type.elements()).enumerate()
113                {
114                    let identifier = Identifier::new(
115                        self.state.assigner.unique_symbol(format_args!("ternary_{i}"), "$$"),
116                        self.state.node_builder.next_id(),
117                    );
118
119                    let expression: Expression = TernaryExpression {
120                        condition: cond.clone(),
121                        if_true: lhs,
122                        if_false: rhs,
123                        span: Default::default(),
124                        id: self.state.node_builder.next_id(),
125                    }
126                    .into();
127
128                    self.state.type_table.insert(identifier.id(), ty.clone());
129                    self.state.type_table.insert(expression.id(), ty.clone());
130
131                    let definition = self.state.assigner.simple_definition(
132                        identifier,
133                        expression,
134                        self.state.node_builder.next_id(),
135                    );
136
137                    statements.push(definition);
138                    elements.push(identifier.into());
139                }
140
141                let expr: Expression =
142                    TupleExpression { elements, span: Default::default(), id: self.state.node_builder.next_id() }
143                        .into();
144
145                self.state.type_table.insert(expr.id(), Type::Tuple(tuple_type.clone()));
146
147                (expr, statements)
148            }
149            (if_true, if_false) => {
150                // This isn't a tuple. Just rebuild it and otherwise leave it alone.
151                (TernaryExpression { if_true, if_false, ..input }.into(), statements)
152            }
153        }
154    }
155}