leo_passes/flattening/
program.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::{FlatteningVisitor, ReturnGuard};
18
19use leo_ast::{AstReconstructor, Expression, Function, ProgramReconstructor, ProgramScope, ReturnStatement, Statement};
20
21impl ProgramReconstructor for FlatteningVisitor<'_> {
22    /// Flattens a program scope.
23    fn reconstruct_program_scope(&mut self, input: ProgramScope) -> ProgramScope {
24        self.program = input.program_id.name.name;
25        ProgramScope {
26            program_id: input.program_id,
27            structs: input.structs.into_iter().map(|(i, c)| (i, self.reconstruct_struct(c))).collect(),
28            mappings: input.mappings.into_iter().map(|(id, mapping)| (id, self.reconstruct_mapping(mapping))).collect(),
29            functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
30            consts: input
31                .consts
32                .into_iter()
33                .map(|(i, c)| match self.reconstruct_const(c) {
34                    (Statement::Const(declaration), _) => (i, declaration),
35                    _ => panic!("`reconstruct_const` can only return `Statement::Const`"),
36                })
37                .collect(),
38            span: input.span,
39        }
40    }
41
42    /// Flattens a function's body
43    fn reconstruct_function(&mut self, function: Function) -> Function {
44        // Set when the function is an async function.
45        self.is_async = function.variant.is_async_function();
46
47        // Flatten the function body.
48        let mut block = self.reconstruct_block(function.block).0;
49
50        // Fold the return statements into the block.
51        let returns = std::mem::take(&mut self.returns);
52        let expression_returns: Vec<(Option<Expression>, ReturnStatement)> = returns
53            .into_iter()
54            .map(|(guard, statement)| match guard {
55                ReturnGuard::None => (None, statement),
56                ReturnGuard::Unconstructed(plain) | ReturnGuard::Constructed { plain, .. } => {
57                    (Some(plain.into()), statement)
58                }
59            })
60            .collect();
61
62        self.fold_returns(&mut block, expression_returns);
63
64        Function {
65            annotations: function.annotations,
66            variant: function.variant,
67            identifier: function.identifier,
68            const_parameters: function.const_parameters,
69            input: function.input,
70            output: function.output,
71            output_type: function.output_type,
72            block,
73            span: function.span,
74            id: function.id,
75        }
76    }
77}
leo_passes/flattening/program.rs

leo_passes/flattening/
program.rs
