Inliner uses AssignmentRenamer instead of SSA; address edge cases

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 <https://www.gnu.org/licenses/>.
 
-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<Symbol, Function>,
 }
 
 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(),
         }
     }

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::<IndexMap<_, _>>();
 
-                // 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() {

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.");

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))
     }
 }