diff --git a/compiler/passes/src/function_inlining/function_inliner.rs b/compiler/passes/src/function_inlining/function_inliner.rs index 4d06af03a9..c28d2f18c3 100644 --- a/compiler/passes/src/function_inlining/function_inliner.rs +++ b/compiler/passes/src/function_inlining/function_inliner.rs @@ -14,7 +14,7 @@ // You should have received a copy of the GNU General Public License // along with the Leo library. If not, see . -use crate::{Assigner, CallGraph, StaticSingleAssigner, SymbolTable}; +use crate::{Assigner, AssignmentRenamer, CallGraph, SymbolTable}; use leo_ast::Function; use leo_span::Symbol; @@ -24,21 +24,18 @@ use indexmap::IndexMap; pub struct FunctionInliner<'a> { /// The call graph for the program. pub(crate) call_graph: &'a CallGraph, - /// A static single assigner used to create unique variable assignments. - pub(crate) static_single_assigner: StaticSingleAssigner<'a>, + /// A wrapper aroung an Assigner used to create unique variable assignments. + pub(crate) assignment_renamer: AssignmentRenamer, /// A map of reconstructed functions in the current program scope. pub(crate) reconstructed_functions: IndexMap, } impl<'a> FunctionInliner<'a> { /// Initializes a new `FunctionInliner`. - pub fn new(symbol_table: &'a SymbolTable, call_graph: &'a CallGraph, assigner: Assigner) -> Self { + pub fn new(_symbol_table: &'a SymbolTable, call_graph: &'a CallGraph, assigner: Assigner) -> Self { Self { call_graph, - // Note: Since we are using the `StaticSingleAssigner` to create unique variable assignments over `BlockStatement`s, we do not need to pass in - // an "accurate" symbol table. This assumption only holds if function inlining occurs after flattening. - // TODO: Refactor out the unique renamer from the static single assigner and use it instead. - static_single_assigner: StaticSingleAssigner::new(symbol_table, assigner), + assignment_renamer: AssignmentRenamer::new(assigner), reconstructed_functions: Default::default(), } } diff --git a/compiler/passes/src/function_inlining/inline_expression.rs b/compiler/passes/src/function_inlining/inline_expression.rs index 86e8a51f7e..21d1c08565 100644 --- a/compiler/passes/src/function_inlining/inline_expression.rs +++ b/compiler/passes/src/function_inlining/inline_expression.rs @@ -16,7 +16,10 @@ use crate::{FunctionInliner, Replacer}; -use leo_ast::{CallExpression, Expression, ExpressionReconstructor, Identifier, ReturnStatement, Statement, StatementConsumer, StatementReconstructor, UnitExpression, Variant}; +use leo_ast::{ + CallExpression, Expression, ExpressionReconstructor, Identifier, ReturnStatement, Statement, + StatementReconstructor, UnitExpression, Variant, +}; use indexmap::IndexMap; use itertools::Itertools; @@ -47,15 +50,27 @@ impl ExpressionReconstructor for FunctionInliner<'_> { .zip_eq(input.arguments.into_iter()) .collect::>(); - // Duplicate the body of the callee and replace each input variable with the appropriate parameter. - let replace = |identifier: Identifier| match parameter_to_argument.get(&identifier) { - Some(expression) => expression.clone(), - None => Expression::Identifier(identifier), - }; - let replaced_block = Replacer::new(replace).reconstruct_block(callee.block.clone()).0; + // Initializer `self.assignment_renamer` with the function parameters. + self.assignment_renamer.load( + callee + .input + .iter() + .map(|input| (input.identifier().name, input.identifier().name)), + ); - // Ensure that each assignment in the `replaced_block` is a unique assignment statement. - let mut inlined_statements = self.static_single_assigner.consume_block(replaced_block); + // Duplicate the body of the callee and create a unique assignment statement for each assignment in the body. + // This is necessary to ensure the inlined variables do not conflict with variables in the caller. + let unique_block = self.assignment_renamer.reconstruct_block(callee.block.clone()).0; + + // Reset `self.assignment_renamer`. + self.assignment_renamer.clear(); + + // Replace each input variable with the appropriate parameter. + let replace = |identifier: &Identifier| match parameter_to_argument.get(identifier) { + Some(expression) => expression.clone(), + None => Expression::Identifier(*identifier), + }; + let mut inlined_statements = Replacer::new(replace).reconstruct_block(unique_block).0.statements; // If the inlined block returns a value, then use the value in place of the call expression, otherwise, use the unit expression. let result = match inlined_statements.last() { diff --git a/compiler/passes/src/function_inlining/inline_statement.rs b/compiler/passes/src/function_inlining/inline_statement.rs index a8edbe1db7..cffd7560f4 100644 --- a/compiler/passes/src/function_inlining/inline_statement.rs +++ b/compiler/passes/src/function_inlining/inline_statement.rs @@ -17,11 +17,44 @@ use crate::FunctionInliner; use leo_ast::{ - Block, ConditionalStatement, ConsoleStatement, DefinitionStatement, IterationStatement, Statement, - StatementReconstructor, + AssignStatement, Block, ConditionalStatement, ConsoleStatement, DefinitionStatement, Expression, + ExpressionReconstructor, ExpressionStatement, IterationStatement, Statement, StatementReconstructor, }; impl StatementReconstructor for FunctionInliner<'_> { + /// Reconstruct an assignment statement by inlining any function calls. + /// This function also segments tuple assignment statements into multiple assignment statements. + fn reconstruct_assign(&mut self, input: AssignStatement) -> (Statement, Self::AdditionalOutput) { + let (value, mut statements) = self.reconstruct_expression(input.value.clone()); + match (input.place, value) { + // If the function call produces a tuple, we need to segment the tuple into multiple assignment statements. + (Expression::Tuple(left), Expression::Tuple(right)) if left.elements.len() == right.elements.len() => { + statements.extend( + left.elements + .into_iter() + .zip(right.elements.into_iter()) + .map(|(lhs, rhs)| { + Statement::Assign(Box::new(AssignStatement { + place: lhs, + value: rhs, + span: Default::default(), + })) + }), + ); + (Statement::dummy(Default::default()), statements) + } + + (place, value) => ( + Statement::Assign(Box::new(AssignStatement { + place, + value, + span: input.span, + })), + statements, + ), + } + } + /// Reconstructs the statements inside a basic block, accumulating any statements produced by function inlining. fn reconstruct_block(&mut self, block: Block) -> (Block, Self::AdditionalOutput) { let mut statements = Vec::with_capacity(block.statements.len()); @@ -56,6 +89,24 @@ impl StatementReconstructor for FunctionInliner<'_> { unreachable!("`DefinitionStatement`s should not exist in the AST at this phase of compilation.") } + /// Reconstructs expression statements by inlining any function calls. + fn reconstruct_expression_statement(&mut self, input: ExpressionStatement) -> (Statement, Self::AdditionalOutput) { + // Reconstruct the expression. + // Note that type checking guarantees that the expression is a function call. + let (expression, additional_statements) = self.reconstruct_expression(input.expression); + + // If the resulting expression is a unit expression, return a dummy statement. + let statement = match expression { + Expression::Unit(_) => Statement::dummy(Default::default()), + _ => Statement::Expression(ExpressionStatement { + expression, + span: input.span, + }), + }; + + (statement, additional_statements) + } + /// Loop unrolling unrolls and removes iteration statements from the program. fn reconstruct_iteration(&mut self, _: IterationStatement) -> (Statement, Self::AdditionalOutput) { unreachable!("`IterationStatement`s should not be in the AST at this phase of compilation."); diff --git a/compiler/passes/src/function_inlining/mod.rs b/compiler/passes/src/function_inlining/mod.rs index 88f064f8b1..d7ce5fd412 100644 --- a/compiler/passes/src/function_inlining/mod.rs +++ b/compiler/passes/src/function_inlining/mod.rs @@ -53,6 +53,9 @@ //! } //! ``` +pub mod assignment_renamer; +pub use assignment_renamer::*; + mod inline_expression; mod inline_statement; @@ -75,6 +78,6 @@ impl<'a> Pass for FunctionInliner<'a> { let mut reconstructor = FunctionInliner::new(st, call_graph, assigner); let program = reconstructor.reconstruct_program(ast.into_repr()); - Ok((Ast::new(program), reconstructor.static_single_assigner.assigner)) + Ok((Ast::new(program), reconstructor.assignment_renamer.assigner)) } }