From a8b317574968ac91f2f96f803c04a1c37d6b7bfe Mon Sep 17 00:00:00 2001
From: Bobbin Threadbare <bobbinth@protonmail.com>
Date: Mon, 11 Sep 2023 11:03:30 -0700
Subject: [PATCH] refactor: make ProgramAst parameter for AssemblyContext
 optional

---
 assembly/src/assembler/context.rs | 13 ++++++++-----
 assembly/src/assembler/mod.rs     |  2 +-
 assembly/src/tests.rs             |  8 ++++----
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/assembly/src/assembler/context.rs b/assembly/src/assembler/context.rs
index a5e035d3f5..4c119d4104 100644
--- a/assembly/src/assembler/context.rs
+++ b/assembly/src/assembler/context.rs
@@ -25,7 +25,7 @@ impl AssemblyContext {
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
 
-    /// Returns a new [AssemblyContext] for non-execurable kernel and non-kernel modules.
+    /// Returns a new [AssemblyContext] for non-executable kernel and non-kernel modules.
     ///
     /// The `is_kernel_module` specifies whether provided module is a kernel module.
     pub fn for_module(is_kernel_module: bool) -> Self {
@@ -37,12 +37,15 @@ impl AssemblyContext {
         }
     }
 
-    /// Returns a new [AssemblyContext] for execurable module.
+    /// Returns a new [AssemblyContext] for executable module.
     ///
-    /// The `program` is required to provide data about IDs and names of imported procedures.
-    pub fn for_program(program: &ProgramAst) -> Self {
+    /// If [ProgramAst] is provided, the context will contain info about the procedures imported
+    /// by the program, and thus, will be able to determine names of imported procedures for error
+    /// reporting purposes.
+    pub fn for_program(program: Option<&ProgramAst>) -> Self {
+        let program_imports = program.map(|p| p.get_imported_procedures_map()).unwrap_or_default();
         Self {
-            module_stack: vec![ModuleContext::for_program(program.get_imported_procedures_map())],
+            module_stack: vec![ModuleContext::for_program(program_imports)],
             is_kernel: false,
             kernel: None,
             allow_phantom_calls: false,
diff --git a/assembly/src/assembler/mod.rs b/assembly/src/assembler/mod.rs
index e4cb4776c5..68ca9ffae0 100644
--- a/assembly/src/assembler/mod.rs
+++ b/assembly/src/assembler/mod.rs
@@ -131,7 +131,7 @@ impl Assembler {
         let program = ProgramAst::parse(source)?;
 
         // compile the program
-        let mut context = AssemblyContext::for_program(&program);
+        let mut context = AssemblyContext::for_program(Some(&program));
         let program_root = self.compile_in_context(&program, &mut context)?;
 
         // convert the context into a call block table for the program
diff --git a/assembly/src/tests.rs b/assembly/src/tests.rs
index 69c12b065f..b3edddd042 100644
--- a/assembly/src/tests.rs
+++ b/assembly/src/tests.rs
@@ -140,14 +140,14 @@ fn simple_main_call() {
     let note_1 =
         ProgramAst::parse("use.context::account begin call.account::account_method_1 end").unwrap();
     let _note_1_root = assembler
-        .compile_in_context(&note_1, &mut super::AssemblyContext::for_program(&note_1))
+        .compile_in_context(&note_1, &mut super::AssemblyContext::for_program(Some(&note_1)))
         .unwrap();
 
     // compile note 2 program
     let note_2 =
         ProgramAst::parse("use.context::account begin call.account::account_method_2 end").unwrap();
     let _note_2_root = assembler
-        .compile_in_context(&note_2, &mut super::AssemblyContext::for_program(&note_2))
+        .compile_in_context(&note_2, &mut super::AssemblyContext::for_program(Some(&note_2)))
         .unwrap();
 }
 
@@ -687,14 +687,14 @@ fn program_with_phantom_mast_call() {
     let ast = ProgramAst::parse(source).unwrap();
 
     // phantom calls not allowed
-    let mut context = AssemblyContext::for_program(&ast).with_phantom_calls(false);
+    let mut context = AssemblyContext::for_program(Some(&ast)).with_phantom_calls(false);
     let result = assembler.compile_in_context(&ast, &mut context);
     let err = result.err().unwrap();
     let expected_error = "cannot call phantom procedure with MAST root 0xc2545da99d3a1f3f38d957c7893c44d78998d8ea8b11aba7e22c8c2b2a213dae: phantom calls not allowed";
     assert_eq!(expected_error, err.to_string());
 
     // phantom calls allowed
-    let mut context = AssemblyContext::for_program(&ast).with_phantom_calls(true);
+    let mut context = AssemblyContext::for_program(Some(&ast)).with_phantom_calls(true);
     let result = assembler.compile_in_context(&ast, &mut context);
     assert!(result.is_ok());
 }