leo_passes/code_generation/
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::*;
18
19use leo_ast::{
20    AssertStatement,
21    AssertVariant,
22    AssignStatement,
23    Block,
24    ConditionalStatement,
25    DefinitionPlace,
26    DefinitionStatement,
27    Expression,
28    ExpressionStatement,
29    IterationStatement,
30    Mode,
31    Output,
32    ReturnStatement,
33    Statement,
34    Type,
35};
36
37use indexmap::IndexMap;
38
39impl CodeGeneratingVisitor<'_> {
40    fn visit_statement(&mut self, input: &Statement) -> Vec<AleoStmt> {
41        match input {
42            Statement::Assert(stmt) => self.visit_assert(stmt),
43            Statement::Assign(stmt) => vec![self.visit_assign(stmt)],
44            Statement::Block(stmt) => self.visit_block(stmt),
45            Statement::Conditional(stmt) => self.visit_conditional(stmt),
46            Statement::Const(_) => {
47                panic!("`ConstStatement`s should not be in the AST at this phase of compilation.")
48            }
49            Statement::Definition(stmt) => self.visit_definition(stmt),
50            Statement::Expression(stmt) => self.visit_expression_statement(stmt),
51            Statement::Iteration(stmt) => vec![self.visit_iteration(stmt)],
52            Statement::Return(stmt) => self.visit_return(stmt),
53        }
54    }
55
56    fn visit_assert(&mut self, input: &AssertStatement) -> Vec<AleoStmt> {
57        match &input.variant {
58            AssertVariant::Assert(expr) => {
59                let (operand, mut instructions) = self.visit_expression(expr);
60                let operand = operand.expect("Trying to assert an empty expression.");
61                instructions.push(AleoStmt::AssertEq(operand, AleoExpr::Bool(true)));
62                instructions
63            }
64            AssertVariant::AssertEq(left, right) => {
65                let (left, left_stmts) = self.visit_expression(left);
66                let (right, right_stmts) = self.visit_expression(right);
67                let left = left.expect("Trying to assert an empty expression.");
68                let right = right.expect("Trying to assert an empty expression.");
69                let assert_instruction = AleoStmt::AssertEq(left, right);
70
71                // Concatenate the instructions.
72                let mut instructions = left_stmts;
73                instructions.extend(right_stmts);
74                instructions.push(assert_instruction);
75                instructions
76            }
77            AssertVariant::AssertNeq(left, right) => {
78                let (left, left_stmts) = self.visit_expression(left);
79                let (right, right_stmts) = self.visit_expression(right);
80                let left = left.expect("Trying to assert an empty expression.");
81                let right = right.expect("Trying to assert an empty expression.");
82                let assert_instruction = AleoStmt::AssertNeq(left, right);
83
84                // Concatenate the instructions.
85                let mut instructions = left_stmts;
86                instructions.extend(right_stmts);
87                instructions.push(assert_instruction);
88                instructions
89            }
90        }
91    }
92
93    fn visit_return(&mut self, input: &ReturnStatement) -> Vec<AleoStmt> {
94        let mut instructions = vec![];
95        let mut operands: IndexMap<AleoExpr, &Output> =
96            IndexMap::with_capacity(self.current_function.unwrap().output.len());
97
98        if let Expression::Tuple(tuple) = &input.expression {
99            // Now tuples only appear in return position, so let's handle this
100            // ourselves.
101            let outputs = &self.current_function.unwrap().output;
102            assert_eq!(tuple.elements.len(), outputs.len());
103
104            for (expr, output) in tuple.elements.iter().zip_eq(outputs) {
105                let (operand, op_instructions) = self.visit_expression(expr);
106                instructions.extend(op_instructions);
107                if let Some(operand) = operand {
108                    if self.internal_record_inputs.contains(&operand) || operands.contains_key(&operand) {
109                        // We can't output an internal record we received as input.
110                        // We also can't output the same value twice.
111                        // Either way, clone it.
112                        let (new_operand, new_instr) = self.clone_register(&operand, &output.type_);
113                        instructions.extend(new_instr);
114                        operands.insert(new_operand, output);
115                    } else {
116                        operands.insert(operand, output);
117                    }
118                }
119            }
120        } else {
121            // Not a tuple - only one output.
122            let (operand, op_instructions) = self.visit_expression(&input.expression);
123            if let Some(operand) = operand {
124                let output = &self.current_function.unwrap().output[0];
125
126                if self.internal_record_inputs.contains(&operand) {
127                    // We can't output an internal record we received as input.
128                    let (new_operand, new_instr) =
129                        self.clone_register(&operand, &self.current_function.unwrap().output_type);
130                    instructions.extend(new_instr);
131                    operands.insert(new_operand, output);
132                } else {
133                    instructions = op_instructions;
134                    operands.insert(operand, output);
135                }
136            }
137        }
138
139        for (operand, output) in operands.iter() {
140            // Transitions outputs with no mode are private.
141            let visibility = match (self.variant.unwrap().is_transition(), output.mode) {
142                (true, Mode::None) => Some(AleoVisibility::Private),
143                (_, mode) => AleoVisibility::maybe_from(mode),
144            };
145            if let Type::Future(_) = output.type_ {
146                instructions.push(AleoStmt::Output(
147                    operand.clone(),
148                    AleoType::Future {
149                        name: self.current_function.unwrap().identifier.to_string(),
150                        program: self.program_id.unwrap().name.to_string(),
151                    },
152                    None,
153                ));
154            } else if output.type_.is_empty() {
155                // do nothing
156            } else {
157                let (output_type, output_viz) = self.visit_type_with_visibility(&output.type_, visibility);
158                instructions.push(AleoStmt::Output(operand.clone(), output_type, output_viz));
159            }
160        }
161
162        instructions
163    }
164
165    fn visit_definition(&mut self, input: &DefinitionStatement) -> Vec<AleoStmt> {
166        match (&input.place, &input.value) {
167            (DefinitionPlace::Single(identifier), _) => {
168                let (operand, expression_instructions) = self.visit_expression(&input.value);
169                if let Some(operand) = operand {
170                    self.variable_mapping.insert(identifier.name, operand);
171                }
172                expression_instructions
173            }
174            (DefinitionPlace::Multiple(identifiers), Expression::Call(_)) => {
175                let (operand, expression_instructions) = self.visit_expression(&input.value);
176                let Some(AleoExpr::Tuple(elems)) = operand else {
177                    panic!("Definition with multiple identifiers should yield a tuple")
178                };
179                // Add the destinations to the variable mapping.
180                for (identifier, operand) in identifiers.iter().zip_eq(elems.iter()) {
181                    self.variable_mapping.insert(identifier.name, operand.clone());
182                }
183                expression_instructions
184            }
185            _ => panic!("Previous passes should have ensured that a definition with multiple identifiers is a `Call`."),
186        }
187    }
188
189    fn visit_expression_statement(&mut self, input: &ExpressionStatement) -> Vec<AleoStmt> {
190        self.visit_expression(&input.expression).1
191    }
192
193    fn visit_assign(&mut self, _input: &AssignStatement) -> AleoStmt {
194        panic!("AssignStatement's should not exist in SSA form.")
195    }
196
197    fn visit_conditional(&mut self, _input: &ConditionalStatement) -> Vec<AleoStmt> {
198        // Note that this unwrap is safe because we set the variant before traversing the function.
199        if !self.variant.unwrap().is_async_function() {
200            panic!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
201        } else {
202            // Construct a label for the end of the `then` block.
203            let end_then_label = format!("end_then_{}_{}", self.conditional_depth, self.next_label);
204            self.next_label += 1;
205            // Construct a label for the end of the `otherwise` block if it exists.
206            let (has_otherwise, end_otherwise_label) = {
207                match _input.otherwise.is_some() {
208                    true => {
209                        // Construct a label for the end of the `otherwise` block.
210                        let end_otherwise_label =
211                            { format!("end_otherwise_{}_{}", self.conditional_depth, self.next_label) };
212                        self.next_label += 1;
213                        (true, end_otherwise_label)
214                    }
215                    false => (false, String::new()),
216                }
217            };
218
219            // Increment the conditional depth.
220            self.conditional_depth += 1;
221
222            // Create a `branch` instruction.
223            let (condition, mut instructions) = self.visit_expression(&_input.condition);
224            let condition = condition.expect("Trying to branch on an empty expression");
225            instructions.push(AleoStmt::BranchEq(condition, AleoExpr::Bool(false), end_then_label.clone()));
226
227            // Visit the `then` block.
228            instructions.extend(self.visit_block(&_input.then));
229            // If the `otherwise` block is present, add a branch instruction to jump to the end of the `otherwise` block.
230            if has_otherwise {
231                instructions.push(AleoStmt::BranchEq(
232                    AleoExpr::Bool(true),
233                    AleoExpr::Bool(true),
234                    end_otherwise_label.clone(),
235                ));
236            }
237
238            // Add a label for the end of the `then` block.
239            instructions.push(AleoStmt::Position(end_then_label));
240
241            // Visit the `otherwise` block.
242            if let Some(else_block) = &_input.otherwise {
243                // Visit the `otherwise` block.
244                instructions.extend(self.visit_statement(else_block));
245                // Add a label for the end of the `otherwise` block.
246                instructions.push(AleoStmt::Position(end_otherwise_label));
247            }
248
249            // Decrement the conditional depth.
250            self.conditional_depth -= 1;
251
252            instructions
253        }
254    }
255
256    fn visit_iteration(&mut self, _input: &IterationStatement) -> AleoStmt {
257        panic!("`IterationStatement`s should not be in the AST at this phase of compilation.");
258    }
259
260    pub(crate) fn visit_block(&mut self, input: &Block) -> Vec<AleoStmt> {
261        // For each statement in the block, visit it and add its instructions to the list.
262        input.statements.iter().flat_map(|stmt| self.visit_statement(stmt)).collect()
263    }
264}