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

leo_passes/code_generation/
statement.rs
