From b2eaaf825d7354793da21a9592363c087820e30e Mon Sep 17 00:00:00 2001
From: evan-schott <53463459+evan-schott@users.noreply.github.com>
Date: Thu, 12 Oct 2023 15:46:09 -0700
Subject: [PATCH] Make compatible w/ LU and CG passes

---
 .../src/code_generation/visit_expressions.rs      | 12 +++++++++++-
 .../passes/src/code_generation/visit_program.rs   |  9 ++++++++-
 .../passes/src/loop_unrolling/unroll_program.rs   | 15 +++++++++++++++
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/compiler/passes/src/code_generation/visit_expressions.rs b/compiler/passes/src/code_generation/visit_expressions.rs
index bccd21f1d3..60811f91f6 100644
--- a/compiler/passes/src/code_generation/visit_expressions.rs
+++ b/compiler/passes/src/code_generation/visit_expressions.rs
@@ -30,6 +30,7 @@ use leo_ast::{
     Identifier,
     Literal,
     MemberAccess,
+    ProgramScope,
     StructExpression,
     TernaryExpression,
     TupleExpression,
@@ -528,6 +529,7 @@ impl<'a> CodeGenerator<'a> {
                     Expression::Identifier(identifier) => identifier.name,
                     _ => unreachable!("Parsing guarantees that a program name is always an identifier."),
                 };
+                let stub_scope: ProgramScope;
                 // Lookup the imported program scope.
                 let imported_program_scope = match self
                     .program
@@ -536,7 +538,15 @@ impl<'a> CodeGenerator<'a> {
                     .and_then(|(program, _)| program.program_scopes.get(&program_name))
                 {
                     Some(program) => program,
-                    None => unreachable!("Type checking guarantees that imported programs are well defined."),
+                    None => {
+                        let stub_program = self.program.stubs.get(&program_name);
+                        if stub_program.is_some() {
+                            stub_scope = ProgramScope::from(stub_program.unwrap().clone());
+                            &stub_scope
+                        } else {
+                            unreachable!("Type checking guarantees that imported and stub programs are well defined.")
+                        }
+                    }
                 };
                 // Check if the external function has a finalize block.
                 let function_name = match *input.function {
diff --git a/compiler/passes/src/code_generation/visit_program.rs b/compiler/passes/src/code_generation/visit_program.rs
index 26fa6ee491..bdc92c1e4e 100644
--- a/compiler/passes/src/code_generation/visit_program.rs
+++ b/compiler/passes/src/code_generation/visit_program.rs
@@ -16,7 +16,7 @@
 
 use crate::CodeGenerator;
 
-use leo_ast::{functions, Function, Mapping, Mode, Program, ProgramScope, Struct, Type, Variant};
+use leo_ast::{functions, Function, Mapping, Mode, Program, ProgramScope, Struct, Stub, Type, Variant};
 
 use indexmap::IndexMap;
 use itertools::Itertools;
@@ -42,6 +42,9 @@ impl<'a> CodeGenerator<'a> {
             program_string.push('\n');
         }
 
+        // Import stub programs
+        program_string.push_str(&input.stubs.values().map(|stub| self.visit_stub(stub)).join("\n"));
+
         // Retrieve the program scope.
         // Note that type checking guarantees that there is exactly one program scope.
         let program_scope: &ProgramScope = input.program_scopes.values().next().unwrap();
@@ -109,6 +112,10 @@ impl<'a> CodeGenerator<'a> {
         program_string
     }
 
+    fn visit_stub(&mut self, input: &'a Stub) -> String {
+        format!("import {}.aleo;", input.stub_id.name)
+    }
+
     fn visit_import(&mut self, import_name: &'a Symbol, import_program: &'a Program) -> String {
         // Load symbols into composite mapping.
         let _import_program_string = self.visit_program(import_program);
diff --git a/compiler/passes/src/loop_unrolling/unroll_program.rs b/compiler/passes/src/loop_unrolling/unroll_program.rs
index de75835318..ead8847ee7 100644
--- a/compiler/passes/src/loop_unrolling/unroll_program.rs
+++ b/compiler/passes/src/loop_unrolling/unroll_program.rs
@@ -34,6 +34,21 @@ impl ProgramReconstructor for Unroller<'_> {
         }
     }
 
+    // Don't need to reconstruct anything, just need to add child scopes for constant propagation table
+    fn reconstruct_function_stub(&mut self, input: FunctionStub) -> FunctionStub {
+        // Lookup function metadata in the symbol table.
+        // Note that this unwrap is safe since function metadata is stored in a prior pass.
+        let function_index = self.symbol_table.borrow().lookup_fn_symbol(input.identifier.name).unwrap().id;
+
+        // Enter the function's scope.
+        let previous_function_index = self.enter_scope(function_index);
+
+        // Exit the function's scope.
+        self.exit_scope(previous_function_index);
+
+        input
+    }
+
     fn reconstruct_function(&mut self, function: Function) -> Function {
         // Lookup function metadata in the symbol table.
         // Note that this unwrap is safe since function metadata is stored in a prior pass.