Ensure type consistency during function inlining

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, AssignmentRenamer, CallGraph};
+use crate::{Assigner, AssignmentRenamer, CallGraph, TypeTable};
 
 use leo_ast::{Function, NodeBuilder};
 use leo_span::Symbol;
@@ -26,18 +26,21 @@ pub struct FunctionInliner<'a> {
     pub(crate) call_graph: &'a CallGraph,
     /// A wrapper around an Assigner used to create unique variable assignments.
     pub(crate) assignment_renamer: AssignmentRenamer<'a>,
+    /// A mapping between node IDs and their types.
+    pub(crate) type_table: &'a mut TypeTable,
     /// A map of reconstructed functions in the current program scope.
     pub(crate) reconstructed_functions: Vec<(Symbol, Function)>,
 }
 
 impl<'a> FunctionInliner<'a> {
     /// Initializes a new `FunctionInliner`.
-    pub fn new(node_builder: &'a NodeBuilder, call_graph: &'a CallGraph, assigner: &'a Assigner) -> Self {
+    pub fn new(node_builder: &'a NodeBuilder, call_graph: &'a CallGraph, assigner: &'a Assigner, type_table: &'a mut TypeTable) -> Self {
         Self {
             node_builder,
             call_graph,
             assignment_renamer: AssignmentRenamer::new(assigner),
             reconstructed_functions: Default::default(),
+            type_table
         }
     }
 }

compiler/passes/src/function_inlining/inline_expression.rs

@@ -16,17 +16,7 @@
 
 use crate::{FunctionInliner, Replacer};
 
-use leo_ast::{
-    CallExpression,
-    Expression,
-    ExpressionReconstructor,
-    Identifier,
-    ReturnStatement,
-    Statement,
-    StatementReconstructor,
-    UnitExpression,
-    Variant,
-};
+use leo_ast::{CallExpression, Expression, ExpressionReconstructor, Identifier, ReturnStatement, Statement, StatementReconstructor, Type, UnitExpression, Variant};
 
 use indexmap::IndexMap;
 use itertools::Itertools;
@@ -93,7 +83,11 @@ impl ExpressionReconstructor for FunctionInliner<'_> {
                             _ => unreachable!("This branch checks that the last statement is a return statement."),
                         }
                     }
-                    _ => Expression::Unit(UnitExpression { span: Default::default(), id: self.node_builder.next_id() }),
+                    _ => {
+                        let id = self.node_builder.next_id();
+                        self.type_table.insert(id, Type::Unit);
+                        Expression::Unit(UnitExpression { span: Default::default(), id })
+                    },
                 };
 
                 (result, inlined_statements)

compiler/passes/src/function_inlining/mod.rs

@@ -70,11 +70,11 @@ use leo_ast::{Ast, NodeBuilder, ProgramReconstructor};
 use leo_errors::Result;
 
 impl<'a> Pass for FunctionInliner<'a> {
-    type Input = (Ast, &'a NodeBuilder, &'a CallGraph, &'a Assigner);
+    type Input = (Ast, &'a NodeBuilder, &'a CallGraph, &'a Assigner, &'a mut TypeTable);
     type Output = Result<Ast>;
 
-    fn do_pass((ast, node_builder, call_graph, assigner): Self::Input) -> Self::Output {
-        let mut reconstructor = FunctionInliner::new(node_builder, call_graph, assigner);
+    fn do_pass((ast, node_builder, call_graph, assigner, tt): Self::Input) -> Self::Output {
+        let mut reconstructor = FunctionInliner::new(node_builder, call_graph, assigner, tt);
         let program = reconstructor.reconstruct_program(ast.into_repr());
 
         Ok(Ast::new(program))