feat: use new debug info/source code management in assembler

assembly/src/assembler/basic_block_builder.rs

@@ -88,15 +88,8 @@ impl BasicBlockBuilder {
         instruction: &Span<Instruction>,
         proc_ctx: &ProcedureContext,
     ) {
-        let location = proc_ctx.source_file().map(|source_file| {
-            let source_file = vm_core::SourceFile::new(source_file.name());
-            let span = instruction.span();
-            vm_core::SourceLocation {
-                source_file,
-                start: span.start().try_into().expect("invalid byte offset: value is > 2^32 bytes"),
-                end: span.end().try_into().expect("invalid byte offset: value is > 2^32 bytes"),
-            }
-        });
+        let span = instruction.span();
+        let location = proc_ctx.source_manager().location(span).ok();
         let context_name = proc_ctx.name().to_string();
         let num_cycles = 0;
         let op = instruction.to_string();

assembly/src/assembler/instruction/env_ops.rs

@@ -1,5 +1,5 @@
 use super::{mem_ops::local_to_absolute_addr, push_felt, BasicBlockBuilder};
-use crate::{assembler::ProcedureContext, AssemblyError, Felt, Spanned};
+use crate::{assembler::ProcedureContext, AssemblyError, Felt, SourceSpan};
 use vm_core::Operation::*;
 
 // CONSTANT INPUTS
@@ -54,11 +54,12 @@ pub fn locaddr(
 pub fn caller(
     span: &mut BasicBlockBuilder,
     proc_ctx: &ProcedureContext,
+    source_span: SourceSpan,
 ) -> Result<(), AssemblyError> {
     if !proc_ctx.is_kernel() {
         return Err(AssemblyError::CallerOutsideOfKernel {
-            span: proc_ctx.span(),
-            source_file: proc_ctx.source_file(),
+            span: source_span,
+            source_file: proc_ctx.source_manager().get(source_span.source_id()).ok(),
         });
     }
     span.push_op(Caller);

assembly/src/assembler/instruction/field_ops.rs

@@ -93,8 +93,9 @@ pub fn div_imm(
     imm: Span<Felt>,
 ) -> Result<(), AssemblyError> {
     if imm == ZERO {
-        let source_file = proc_ctx.source_file();
-        let error = Report::new(crate::parser::ParsingError::DivisionByZero { span: imm.span() });
+        let source_span = imm.span();
+        let source_file = proc_ctx.source_manager().get(source_span.source_id()).ok();
+        let error = Report::new(crate::parser::ParsingError::DivisionByZero { span: source_span });
         return Err(if let Some(source_file) = source_file {
             AssemblyError::Other(RelatedError::new(error.with_source_code(source_file)))
         } else {

assembly/src/assembler/instruction/mod.rs

@@ -313,7 +313,7 @@ impl Assembler {
             Instruction::PushU32List(imms) => env_ops::push_many(imms, span_builder),
             Instruction::PushFeltList(imms) => env_ops::push_many(imms, span_builder),
             Instruction::Sdepth => span_builder.push_op(SDepth),
-            Instruction::Caller => env_ops::caller(span_builder, proc_ctx)?,
+            Instruction::Caller => env_ops::caller(span_builder, proc_ctx, instruction.span())?,
             Instruction::Clk => span_builder.push_op(Clk),
             Instruction::AdvPipe => span_builder.push_op(Pipe),
             Instruction::AdvPush(n) => adv_ops::adv_push(span_builder, n.expect_value())?,

assembly/src/assembler/instruction/procedures.rs

@@ -31,7 +31,7 @@ impl Assembler {
         mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         // Get the procedure from the assembler
-        let current_source_file = proc_ctx.source_file();
+        let current_source_file = self.source_manager.get(span.source_id()).ok();
 
         // If the procedure is cached, register the call to ensure the callset
         // is updated correctly.

assembly/src/assembler/instruction/u32_ops.rs

@@ -372,9 +372,9 @@ fn handle_division(
 ) -> Result<(), AssemblyError> {
     if let Some(imm) = imm {
         if imm == 0 {
-            let source_file = proc_ctx.source_file();
-            let error =
-                Report::new(crate::parser::ParsingError::DivisionByZero { span: imm.span() });
+            let imm_span = imm.span();
+            let source_file = proc_ctx.source_manager().get(imm_span.source_id()).ok();
+            let error = Report::new(crate::parser::ParsingError::DivisionByZero { span: imm_span });
             return if let Some(source_file) = source_file {
                 Err(AssemblyError::Other(RelatedError::new(error.with_source_code(source_file))))
             } else {

assembly/src/assembler/mod.rs

@@ -4,7 +4,7 @@ use crate::{
     library::{CompiledLibrary, KernelLibrary},
     sema::SemanticAnalysisError,
     AssemblyError, Compile, CompileOptions, Library, LibraryNamespace, LibraryPath, RpoDigest,
-    Spanned,
+    SourceManager, Spanned,
 };
 use alloc::{sync::Arc, vec::Vec};
 use mast_forest_builder::MastForestBuilder;
@@ -47,8 +47,10 @@ use self::module_graph::{CallerInfo, ModuleGraph, ResolvedTarget};
 ///   procedures, build the assembler with them first, using the various builder methods on
 ///   [Assembler], e.g. [Assembler::with_module], [Assembler::with_library], etc. Then, call
 ///   [Assembler::assemble_program] to get your compiled program.
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub struct Assembler {
+    /// The source manager to use for compilation and source location information
+    source_manager: Arc<dyn SourceManager>,
     /// The global [ModuleGraph] for this assembler.
     module_graph: ModuleGraph,
     /// Whether to treat warning diagnostics as errors
@@ -57,14 +59,42 @@ pub struct Assembler {
     in_debug_mode: bool,
 }
 
+impl Default for Assembler {
+    #[allow(clippy::arc_with_non_send_sync)]
+    fn default() -> Self {
+        use vm_core::debuginfo::SingleThreadedSourceManager;
+        let source_manager = Arc::new(SingleThreadedSourceManager::default());
+        let module_graph = ModuleGraph::new(source_manager.clone());
+        Self {
+            source_manager,
+            module_graph,
+            warnings_as_errors: false,
+            in_debug_mode: false,
+        }
+    }
+}
+
 // ------------------------------------------------------------------------------------------------
 /// Constructors
 impl Assembler {
+    /// Start building an [Assembler]
+    pub fn new(source_manager: Arc<dyn SourceManager>) -> Self {
+        let module_graph = ModuleGraph::new(source_manager.clone());
+        Self {
+            source_manager,
+            module_graph,
+            warnings_as_errors: false,
+            in_debug_mode: false,
+        }
+    }
+
     /// Start building an [`Assembler`] with a kernel defined by the provided [KernelLibrary].
-    pub fn with_kernel(kernel_lib: KernelLibrary) -> Self {
+    pub fn with_kernel(source_manager: Arc<dyn SourceManager>, kernel_lib: KernelLibrary) -> Self {
         let (kernel, kernel_module, _) = kernel_lib.into_parts();
+        let module_graph = ModuleGraph::with_kernel(source_manager.clone(), kernel, kernel_module);
         Self {
-            module_graph: ModuleGraph::with_kernel(kernel, kernel_module),
+            source_manager,
+            module_graph,
             ..Default::default()
         }
     }
@@ -130,7 +160,7 @@ impl Assembler {
             ));
         }
 
-        let module = module.compile_with_options(options)?;
+        let module = module.compile_with_options(&self.source_manager, options)?;
         assert_eq!(module.kind(), kind, "expected module kind to match compilation options");
 
         self.module_graph.add_ast_module(module)?;
@@ -240,7 +270,8 @@ impl Assembler {
     ) -> Result<CompiledLibrary, Report> {
         let ast_module_indices: Vec<ModuleIndex> = modules
             .map(|module| {
-                let module = module.compile_with_options(CompileOptions::for_library())?;
+                let module = module
+                    .compile_with_options(&self.source_manager, CompileOptions::for_library())?;
 
                 Ok(self.module_graph.add_ast_module(module)?)
             })
@@ -281,7 +312,7 @@ impl Assembler {
             path: Some(LibraryPath::from(LibraryNamespace::Kernel)),
         };
 
-        let module = module.compile_with_options(options)?;
+        let module = module.compile_with_options(&self.source_manager, options)?;
         let module_idx = self.module_graph.add_ast_module(module)?;
 
         self.module_graph.recompute()?;
@@ -318,7 +349,7 @@ impl Assembler {
             path: Some(LibraryPath::from(LibraryNamespace::Exec)),
         };
 
-        let program = source.compile_with_options(options)?;
+        let program = source.compile_with_options(&self.source_manager, options)?;
         assert!(program.is_executable());
 
         // Recompute graph with executable module, and start compiling
@@ -418,10 +449,14 @@ impl Assembler {
                 module: module.path().clone(),
                 name: ast.name().clone(),
             };
-            let pctx = ProcedureContext::new(procedure_gid, name, ast.visibility())
-                .with_num_locals(num_locals as u16)
-                .with_span(ast.span())
-                .with_source_file(ast.source_file());
+            let pctx = ProcedureContext::new(
+                procedure_gid,
+                name,
+                ast.visibility(),
+                self.source_manager.clone(),
+            )
+            .with_num_locals(num_locals as u16)
+            .with_span(ast.span());
 
             // Compile this procedure
             let procedure = self.compile_procedure(pctx, mast_forest_builder)?;
@@ -578,7 +613,6 @@ impl Assembler {
     ) -> Result<RpoDigest, AssemblyError> {
         let caller = CallerInfo {
             span: target.span(),
-            source_file: proc_ctx.source_file(),
             module: proc_ctx.id().module,
             kind,
         };

assembly/src/assembler/module_graph/analysis/rewrite_check.rs

@@ -1,4 +1,3 @@
-use alloc::sync::Arc;
 use core::ops::ControlFlow;
 
 use crate::{
@@ -7,7 +6,6 @@ use crate::{
         ModuleIndex, ResolvedTarget,
     },
     ast::{visit::Visit, InvocationTarget, InvokeKind, Module},
-    diagnostics::SourceFile,
     AssemblyError, Spanned,
 };
 
@@ -38,7 +36,6 @@ impl<'a, 'b: 'a> MaybeRewriteCheck<'a, 'b> {
         let mut visitor = RewriteCheckVisitor {
             resolver: self.resolver,
             module_id,
-            source_file: module.source_file(),
         };
         match visitor.visit_module(module) {
             ControlFlow::Break(result) => result,
@@ -53,7 +50,6 @@ impl<'a, 'b: 'a> MaybeRewriteCheck<'a, 'b> {
 struct RewriteCheckVisitor<'a, 'b: 'a> {
     resolver: &'a NameResolver<'b>,
     module_id: ModuleIndex,
-    source_file: Option<Arc<SourceFile>>,
 }
 
 impl<'a, 'b: 'a> RewriteCheckVisitor<'a, 'b> {
@@ -64,7 +60,6 @@ impl<'a, 'b: 'a> RewriteCheckVisitor<'a, 'b> {
     ) -> ControlFlow<Result<bool, AssemblyError>> {
         let caller = CallerInfo {
             span: target.span(),
-            source_file: self.source_file.clone(),
             module: self.module_id,
             kind,
         };