leo_passes/function_inlining/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::FunctionInliningVisitor;
18use leo_ast::{AstReconstructor, Constructor, Function, Program, ProgramReconstructor, ProgramScope};
19
20use snarkvm::prelude::Itertools;
21
22impl ProgramReconstructor for FunctionInliningVisitor<'_> {
23 fn reconstruct_program_scope(&mut self, input: ProgramScope) -> ProgramScope {
24 // Set the program name.
25 self.program = input.program_id.name.name;
26
27 // Get the post-order ordering of the call graph.
28 // Note that the post-order always contains all nodes in the call graph.
29 // Note that the unwrap is safe since type checking guarantees that the call graph is acyclic.
30 let order = self
31 .state
32 .call_graph
33 .post_order()
34 .unwrap()
35 .into_iter()
36 .filter_map(|location| (location.program == self.program).then_some(location.path))
37 .collect_vec();
38
39 // Reconstruct and accumulate each of the functions in post-order.
40 for function_name in order {
41 // None: If `function_name` is not in `input.functions`, then it must be an external function.
42 // TODO: Check that this is indeed an external function. Requires a redesign of the symbol table.
43 if let Some(function) = self.function_map.shift_remove(&function_name) {
44 // Reconstruct the function.
45 let reconstructed_function = self.reconstruct_function(function);
46 // Add the reconstructed function to the mapping.
47 self.reconstructed_functions.push((function_name.clone(), reconstructed_function));
48 }
49 }
50
51 // This is a sanity check to ensure that functions in the program scope have been processed.
52 assert!(self.function_map.is_empty(), "All functions in the program should have been processed.");
53
54 // Reconstruct the constructor.
55 // Note: This must be done after the functions have been reconstructed to ensure that every callee function has been inlined.
56 let constructor = input.constructor.map(|constructor| self.reconstruct_constructor(constructor));
57
58 // Note that this intentionally clears `self.reconstructed_functions` for the next program scope.
59 let functions = core::mem::take(&mut self.reconstructed_functions)
60 .iter()
61 .filter_map(|(path, f)| {
62 // Only consider functions defined at program scope. The rest are not relevant since they should all
63 // have been inlined by now.
64 path.split_last().filter(|(_, rest)| rest.is_empty()).map(|(last, _)| (*last, f.clone()))
65 })
66 .collect();
67
68 ProgramScope {
69 program_id: input.program_id,
70 structs: input.structs,
71 mappings: input.mappings,
72 constructor,
73 functions,
74 consts: input.consts,
75 span: input.span,
76 }
77 }
78
79 fn reconstruct_function(&mut self, input: Function) -> Function {
80 Function {
81 annotations: input.annotations,
82 variant: input.variant,
83 identifier: input.identifier,
84 const_parameters: input.const_parameters,
85 input: input.input,
86 output: input.output,
87 output_type: input.output_type,
88 block: {
89 // Set the `is_async` flag before reconstructing the block.
90 self.is_async = input.variant.is_async_function();
91 // Reconstruct the block.
92 let block = self.reconstruct_block(input.block).0;
93 // Reset the `is_async` flag.
94 self.is_async = false;
95 block
96 },
97 span: input.span,
98 id: input.id,
99 }
100 }
101
102 fn reconstruct_constructor(&mut self, input: Constructor) -> Constructor {
103 Constructor {
104 annotations: input.annotations,
105 block: {
106 // Set the `is_async` flag before reconstructing the block.
107 self.is_async = true;
108 // Reconstruct the block.
109 let block = self.reconstruct_block(input.block).0;
110 // Reset the `is_async` flag.
111 self.is_async = false;
112 block
113 },
114 span: input.span,
115 id: input.id,
116 }
117 }
118
119 fn reconstruct_program(&mut self, input: Program) -> Program {
120 // Populate `self.function_map` using the functions in the program scopes and the modules
121 input
122 .modules
123 .iter()
124 .flat_map(|(module_path, m)| {
125 m.functions.iter().map(move |(name, f)| {
126 (module_path.iter().cloned().chain(std::iter::once(*name)).collect(), f.clone())
127 })
128 })
129 .chain(
130 input
131 .program_scopes
132 .iter()
133 .flat_map(|(_, scope)| scope.functions.iter().map(|(name, f)| (vec![*name], f.clone()))),
134 )
135 .for_each(|(full_name, f)| {
136 self.function_map.insert(full_name, f);
137 });
138
139 // It's sufficient to reconstruct program scopes because `inline` functions defined in
140 // modules will be traversed using the call graph and reconstructed in the right order, so
141 // no need to reconstruct the modules explicitly.
142 Program {
143 program_scopes: input
144 .program_scopes
145 .into_iter()
146 .map(|(id, scope)| (id, self.reconstruct_program_scope(scope)))
147 .collect(),
148 ..input
149 }
150 }
151}