leo_passes/option_lowering/
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::OptionLoweringVisitor;
18use leo_ast::{
19    AstReconstructor,
20    ConstParameter,
21    Function,
22    Input,
23    Module,
24    Output,
25    Program,
26    ProgramReconstructor,
27    ProgramScope,
28    Statement,
29};
30use leo_span::Symbol;
31
32impl ProgramReconstructor for OptionLoweringVisitor<'_> {
33    fn reconstruct_program(&mut self, input: Program) -> Program {
34        // Reconstruct all structs first and keep track of them in `self.reconstructed_structs`.
35        for (_, scope) in &input.program_scopes {
36            for (_, c) in &scope.structs {
37                let new_struct = self.reconstruct_struct(c.clone());
38                self.reconstructed_structs.insert(vec![new_struct.name()], new_struct);
39            }
40        }
41        for (module_path, module) in &input.modules {
42            for (_, c) in &module.structs {
43                let full_name = module_path.iter().cloned().chain(std::iter::once(c.name())).collect::<Vec<Symbol>>();
44                let new_struct = self.reconstruct_struct(c.clone());
45                self.reconstructed_structs.insert(full_name, new_struct.clone());
46            }
47        }
48
49        // Now we're ready to reconstruct everything else.
50        Program {
51            imports: input
52                .imports
53                .into_iter()
54                .map(|(id, import)| (id, (self.reconstruct_import(import.0), import.1)))
55                .collect(),
56            stubs: input.stubs.into_iter().map(|(id, stub)| (id, self.reconstruct_stub(stub))).collect(),
57            program_scopes: input
58                .program_scopes
59                .into_iter()
60                .map(|(id, scope)| (id, self.reconstruct_program_scope(scope)))
61                .collect(),
62            modules: input.modules.into_iter().map(|(id, module)| (id, self.reconstruct_module(module))).collect(),
63        }
64    }
65
66    fn reconstruct_program_scope(&mut self, input: ProgramScope) -> ProgramScope {
67        self.program = input.program_id.name.name;
68
69        let mut program = ProgramScope {
70            consts: input
71                .consts
72                .into_iter()
73                .map(|(i, c)| match self.reconstruct_const(c) {
74                    (Statement::Const(decl), _) => (i, decl),
75                    _ => panic!("`reconstruct_const` can only return `Statement::Const`"),
76                })
77                .collect(),
78            structs: self
79                .reconstructed_structs
80                .iter()
81                .filter_map(|(path, s)| {
82                    path.split_last().filter(|(_, rest)| rest.is_empty()).map(|(last, _)| (*last, s.clone()))
83                })
84                .collect(),
85            mappings: input.mappings.into_iter().map(|(id, m)| (id, self.reconstruct_mapping(m))).collect(),
86            functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
87            constructor: input.constructor.map(|c| self.reconstruct_constructor(c)),
88            ..input
89        };
90
91        program.structs.extend(self.new_structs.drain(..));
92        program
93    }
94
95    fn reconstruct_module(&mut self, input: Module) -> Module {
96        self.program = input.program_name;
97        self.in_module_scope(&input.path.clone(), |slf| Module {
98            consts: input
99                .consts
100                .into_iter()
101                .map(|(i, c)| match slf.reconstruct_const(c) {
102                    (Statement::Const(declaration), _) => (i, declaration),
103                    _ => panic!("`reconstruct_const` can only return `Statement::Const`"),
104                })
105                .collect(),
106            structs: slf
107                .reconstructed_structs
108                .iter()
109                .filter_map(|(path, c)| path.split_last().map(|(last, rest)| (last, rest, c)))
110                .filter(|&(_, rest, _)| input.path == rest)
111                .map(|(last, _, c)| (*last, c.clone()))
112                .collect(),
113            functions: input.functions.into_iter().map(|(i, f)| (i, slf.reconstruct_function(f))).collect(),
114            ..input
115        })
116    }
117
118    fn reconstruct_function(&mut self, input: Function) -> Function {
119        self.function = Some(input.identifier.name);
120        Function {
121            const_parameters: input
122                .const_parameters
123                .iter()
124                .map(|param| ConstParameter { type_: self.reconstruct_type(param.type_.clone()).0, ..param.clone() })
125                .collect(),
126            input: input
127                .input
128                .iter()
129                .map(|input| Input { type_: self.reconstruct_type(input.type_.clone()).0, ..input.clone() })
130                .collect(),
131            output: input
132                .output
133                .iter()
134                .map(|output| Output { type_: self.reconstruct_type(output.type_.clone()).0, ..output.clone() })
135                .collect(),
136            output_type: self.reconstruct_type(input.output_type).0,
137            block: self.reconstruct_block(input.block).0,
138            ..input
139        }
140    }
141}
leo_passes/option_lowering/program.rs

leo_passes/option_lowering/
program.rs
