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 storage_variables: input.storage_variables,
73 constructor,
74 functions,
75 consts: input.consts,
76 span: input.span,
77 }
78 }
79
80 fn reconstruct_function(&mut self, input: Function) -> Function {
81 Function {
82 annotations: input.annotations,
83 variant: input.variant,
84 identifier: input.identifier,
85 const_parameters: input.const_parameters,
86 input: input.input,
87 output: input.output,
88 output_type: input.output_type,
89 block: {
90 // Set the `is_async` flag before reconstructing the block.
91 self.is_async = input.variant.is_async_function();
92 // Reconstruct the block.
93 let block = self.reconstruct_block(input.block).0;
94 // Reset the `is_async` flag.
95 self.is_async = false;
96 block
97 },
98 span: input.span,
99 id: input.id,
100 }
101 }
102
103 fn reconstruct_constructor(&mut self, input: Constructor) -> Constructor {
104 Constructor {
105 annotations: input.annotations,
106 block: {
107 // Set the `is_async` flag before reconstructing the block.
108 self.is_async = true;
109 // Reconstruct the block.
110 let block = self.reconstruct_block(input.block).0;
111 // Reset the `is_async` flag.
112 self.is_async = false;
113 block
114 },
115 span: input.span,
116 id: input.id,
117 }
118 }
119
120 fn reconstruct_program(&mut self, input: Program) -> Program {
121 // Populate `self.function_map` using the functions in the program scopes and the modules
122 input
123 .modules
124 .iter()
125 .flat_map(|(module_path, m)| {
126 m.functions.iter().map(move |(name, f)| {
127 (module_path.iter().cloned().chain(std::iter::once(*name)).collect(), f.clone())
128 })
129 })
130 .chain(
131 input
132 .program_scopes
133 .iter()
134 .flat_map(|(_, scope)| scope.functions.iter().map(|(name, f)| (vec![*name], f.clone()))),
135 )
136 .for_each(|(full_name, f)| {
137 self.function_map.insert(full_name, f);
138 });
139
140 // It's sufficient to reconstruct program scopes because `inline` functions defined in
141 // modules will be traversed using the call graph and reconstructed in the right order, so
142 // no need to reconstruct the modules explicitly.
143 Program {
144 program_scopes: input
145 .program_scopes
146 .into_iter()
147 .map(|(id, scope)| (id, self.reconstruct_program_scope(scope)))
148 .collect(),
149 ..input
150 }
151 }
152}