leo_passes/function_inlining/
ast.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 crate::{Replacer, SsaFormingInput, static_single_assignment::visitor::SsaFormingVisitor};
19
20use leo_ast::*;
21
22use indexmap::IndexMap;
23use itertools::Itertools;
24
25impl AstReconstructor for FunctionInliningVisitor<'_> {
26    type AdditionalInput = ();
27    type AdditionalOutput = Vec<Statement>;
28
29    /* Expressions */
30    fn reconstruct_call(&mut self, input: CallExpression, _additional: &()) -> (Expression, Self::AdditionalOutput) {
31        // Type checking guarantees that only functions local to the program scope can be inlined.
32        if input.program.is_some_and(|prog| prog != self.program) {
33            return (input.into(), Default::default());
34        }
35
36        // Lookup the reconstructed callee function.
37        // Since this pass processes functions in post-order, the callee function is guaranteed to exist in `self.reconstructed_functions`
38        let (_, callee) = self
39            .reconstructed_functions
40            .iter()
41            .find(|(path, _)| *path == input.function.absolute_path())
42            .expect("guaranteed to exist due to post-order traversal of the call graph.");
43
44        // Inline the callee function, if required, otherwise, return the call expression.
45        match callee.variant {
46            Variant::Inline => {
47                // Construct a mapping from input variables of the callee function to arguments passed to the callee.
48                let parameter_to_argument = callee
49                    .input
50                    .iter()
51                    .map(|input| input.identifier().name)
52                    .zip_eq(input.arguments)
53                    .collect::<IndexMap<_, _>>();
54
55                // Function to replace path expressions with their corresponding const argument or keep them unchanged.
56                let replace_path = |expr: &Expression| match expr {
57                    Expression::Path(path) => parameter_to_argument
58                        .get(&path.identifier().name)
59                        .map_or(Expression::Path(path.clone()), |expr| expr.clone()),
60                    _ => expr.clone(),
61                };
62
63                // Replace path expressions with their corresponding const argument or keep them unchanged.
64                let reconstructed_block = Replacer::new(replace_path, false /* refresh IDs */, self.state)
65                    .reconstruct_block(callee.block.clone())
66                    .0;
67
68                // Run SSA formation on the inlined block and rename definitions. Renaming is necessary to avoid shadowing variables.
69                let mut inlined_statements =
70                    SsaFormingVisitor::new(self.state, SsaFormingInput { rename_defs: true }, self.program)
71                        .consume_block(reconstructed_block);
72
73                // If the inlined block returns a value, then use the value in place of the call expression; otherwise, use the unit expression.
74                let result = match inlined_statements.last() {
75                    Some(Statement::Return(_)) => {
76                        // Note that this unwrap is safe since we know that the last statement is a return statement.
77                        match inlined_statements.pop().unwrap() {
78                            Statement::Return(ReturnStatement { expression, .. }) => expression,
79                            _ => panic!("This branch checks that the last statement is a return statement."),
80                        }
81                    }
82                    _ => {
83                        let id = self.state.node_builder.next_id();
84                        self.state.type_table.insert(id, Type::Unit);
85                        UnitExpression { span: Default::default(), id }.into()
86                    }
87                };
88
89                (result, inlined_statements)
90            }
91            Variant::Function
92            | Variant::Script
93            | Variant::AsyncFunction
94            | Variant::Transition
95            | Variant::AsyncTransition => (input.into(), Default::default()),
96        }
97    }
98
99    /* Statements */
100    fn reconstruct_assign(&mut self, _input: AssignStatement) -> (Statement, Self::AdditionalOutput) {
101        panic!("`AssignStatement`s should not exist in the AST at this phase of compilation.")
102    }
103
104    /// Reconstructs the statements inside a basic block, accumulating any statements produced by function inlining.
105    fn reconstruct_block(&mut self, block: Block) -> (Block, Self::AdditionalOutput) {
106        let mut statements = Vec::with_capacity(block.statements.len());
107
108        for statement in block.statements {
109            let (reconstructed_statement, additional_statements) = self.reconstruct_statement(statement);
110            statements.extend(additional_statements);
111            statements.push(reconstructed_statement);
112        }
113
114        (Block { span: block.span, statements, id: block.id }, Default::default())
115    }
116
117    /// Flattening removes conditional statements from the program.
118    fn reconstruct_conditional(&mut self, input: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
119        if !self.is_async {
120            panic!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
121        } else {
122            (
123                ConditionalStatement {
124                    condition: self.reconstruct_expression(input.condition, &()).0,
125                    then: self.reconstruct_block(input.then).0,
126                    otherwise: input.otherwise.map(|n| Box::new(self.reconstruct_statement(*n).0)),
127                    span: input.span,
128                    id: input.id,
129                }
130                .into(),
131                Default::default(),
132            )
133        }
134    }
135
136    /// Reconstruct a definition statement by inlining any function calls.
137    /// This function also segments tuple assignment statements into multiple assignment statements.
138    fn reconstruct_definition(&mut self, mut input: DefinitionStatement) -> (Statement, Self::AdditionalOutput) {
139        let (value, mut statements) = self.reconstruct_expression(input.value, &());
140        match (input.place, value) {
141            // If we just inlined the production of a tuple literal, we need multiple definition statements.
142            (DefinitionPlace::Multiple(left), Expression::Tuple(right)) => {
143                assert_eq!(left.len(), right.elements.len());
144                for (identifier, rhs_value) in left.into_iter().zip(right.elements) {
145                    let stmt = DefinitionStatement {
146                        place: DefinitionPlace::Single(identifier),
147                        type_: None,
148                        value: rhs_value,
149                        span: Default::default(),
150                        id: self.state.node_builder.next_id(),
151                    }
152                    .into();
153
154                    statements.push(stmt);
155                }
156                (Statement::dummy(), statements)
157            }
158
159            (place, value) => {
160                input.value = value;
161                input.place = place;
162                (input.into(), statements)
163            }
164        }
165    }
166
167    /// Reconstructs expression statements by inlining any function calls.
168    fn reconstruct_expression_statement(&mut self, input: ExpressionStatement) -> (Statement, Self::AdditionalOutput) {
169        // Reconstruct the expression.
170        // Note that type checking guarantees that the expression is a function call.
171        let (expression, additional_statements) = self.reconstruct_expression(input.expression, &());
172
173        // If the resulting expression is a unit expression, return a dummy statement.
174        let statement = match expression {
175            Expression::Unit(_) => Statement::dummy(),
176            _ => ExpressionStatement { expression, ..input }.into(),
177        };
178
179        (statement, additional_statements)
180    }
181
182    /// Loop unrolling unrolls and removes iteration statements from the program.
183    fn reconstruct_iteration(&mut self, _: IterationStatement) -> (Statement, Self::AdditionalOutput) {
184        panic!("`IterationStatement`s should not be in the AST at this phase of compilation.");
185    }
186}
leo_passes/function_inlining/ast.rs

leo_passes/function_inlining/
ast.rs
