Make Assembler single-use

CHANGELOG.md

@@ -20,6 +20,7 @@
 - Updated CI to support `CHANGELOG.md` modification checking and `no changelog` label (#1406)
 - Introduced `MastForestError` to enforce `MastForest` node count invariant (#1394)
 - Added functions to `MastForestBuilder` to allow ensuring of nodes with fewer LOC (#1404)
+- Make `Assembler` single-use (#1409)
 
 #### Changed
 

assembly/src/assembler/basic_block_builder.rs

@@ -1,6 +1,6 @@
 use super::{
-    mast_forest_builder::MastForestBuilder, AssemblyContext, BodyWrapper, Decorator, DecoratorList,
-    Instruction,
+    context::ProcedureContext, mast_forest_builder::MastForestBuilder, BodyWrapper, Decorator,
+    DecoratorList, Instruction,
 };
 use alloc::{borrow::Borrow, string::ToString, vec::Vec};
 use vm_core::{
@@ -85,8 +85,8 @@ impl BasicBlockBuilder {
     ///
     /// This indicates that the provided instruction should be tracked and the cycle count for
     /// this instruction will be computed when the call to set_instruction_cycle_count() is made.
-    pub fn track_instruction(&mut self, instruction: &Instruction, ctx: &AssemblyContext) {
-        let context_name = ctx.unwrap_current_procedure().name().to_string();
+    pub fn track_instruction(&mut self, instruction: &Instruction, proc_ctx: &ProcedureContext) {
+        let context_name = proc_ctx.name().to_string();
         let num_cycles = 0;
         let op = instruction.to_string();
         let should_break = instruction.should_break();

assembly/src/assembler/context.rs

@@ -1,193 +1,14 @@
 use alloc::{boxed::Box, sync::Arc};
 
-use super::{procedure::CallSet, ArtifactKind, GlobalProcedureIndex, Procedure};
+use super::{procedure::CallSet, GlobalProcedureIndex, Procedure};
 use crate::{
     ast::{FullyQualifiedProcedureName, Visibility},
     diagnostics::SourceFile,
-    AssemblyError, LibraryPath, RpoDigest, SourceSpan, Span, Spanned,
+    AssemblyError, LibraryPath, RpoDigest, SourceSpan, Spanned,
 };
 use vm_core::mast::{MastForest, MastNodeId};
 
-// ASSEMBLY CONTEXT
-// ================================================================================================
-
-/// An [AssemblyContext] is used to store configuration and state pertaining to the current
-/// compilation of a module/procedure by an [crate::Assembler].
-///
-/// The context specifies context-specific configuration, the type of artifact being compiled,
-/// the current module being compiled, and the current procedure being compiled.
-///
-/// To provide a custom context, you must compile by invoking the
-/// [crate::Assembler::assemble_in_context] API, which will use the provided context in place of
-/// the default one generated internally by the other `compile`-like APIs.
-#[derive(Default)]
-pub struct AssemblyContext {
-    /// What kind of artifact are we assembling
-    kind: ArtifactKind,
-    /// When true, promote warning diagnostics to errors
-    warnings_as_errors: bool,
-    /// When true, this permits calls to refer to procedures which are not locally available,
-    /// as long as they are referenced by MAST root, and not by name. As long as the MAST for those
-    /// roots is present when the code is executed, this works fine. However, if the VM tries to
-    /// execute a program with such calls, and the MAST is not available, the program will trap.
-    allow_phantom_calls: bool,
-    /// The current procedure being compiled
-    current_procedure: Option<ProcedureContext>,
-    /// The fully-qualified module path which should be compiled.
-    ///
-    /// If unset, it defaults to the module which represents the specified `kind`, i.e. if the kind
-    /// is executable, we compile the executable module, and so on.
-    ///
-    /// When set, the module graph is traversed from the given module only, so any code unreachable
-    /// from this module is not considered for compilation.
-    root: Option<LibraryPath>,
-}
-
-/// Builders
-impl AssemblyContext {
-    pub fn new(kind: ArtifactKind) -> Self {
-        Self {
-            kind,
-            ..Default::default()
-        }
-    }
-
-    /// Returns a new [AssemblyContext] for a non-executable kernel modules.
-    pub fn for_kernel(path: &LibraryPath) -> Self {
-        Self::new(ArtifactKind::Kernel).with_root(path.clone())
-    }
-
-    /// Returns a new [AssemblyContext] for library modules.
-    pub fn for_library(path: &LibraryPath) -> Self {
-        Self::new(ArtifactKind::Library).with_root(path.clone())
-    }
-
-    /// Returns a new [AssemblyContext] for an executable module.
-    pub fn for_program(path: &LibraryPath) -> Self {
-        Self::new(ArtifactKind::Executable).with_root(path.clone())
-    }
-
-    fn with_root(mut self, path: LibraryPath) -> Self {
-        self.root = Some(path);
-        self
-    }
-
-    /// When true, all warning diagnostics are promoted to errors
-    #[inline(always)]
-    pub fn set_warnings_as_errors(&mut self, yes: bool) {
-        self.warnings_as_errors = yes;
-    }
-
-    #[inline]
-    pub(super) fn set_current_procedure(&mut self, context: ProcedureContext) {
-        self.current_procedure = Some(context);
-    }
-
-    #[inline]
-    pub(super) fn take_current_procedure(&mut self) -> Option<ProcedureContext> {
-        self.current_procedure.take()
-    }
-
-    #[inline]
-    pub(super) fn unwrap_current_procedure(&self) -> &ProcedureContext {
-        self.current_procedure.as_ref().expect("missing current procedure context")
-    }
-
-    #[inline]
-    pub(super) fn unwrap_current_procedure_mut(&mut self) -> &mut ProcedureContext {
-        self.current_procedure.as_mut().expect("missing current procedure context")
-    }
-
-    /// Enables phantom calls when compiling with this context.
-    ///
-    /// # Panics
-    ///
-    /// This function will panic if you attempt to enable phantom calls for a kernel-mode context,
-    /// as non-local procedure calls are not allowed in kernel modules.
-    pub fn with_phantom_calls(mut self, allow_phantom_calls: bool) -> Self {
-        assert!(
-            !self.is_kernel() || !allow_phantom_calls,
-            "kernel modules cannot have phantom calls enabled"
-        );
-        self.allow_phantom_calls = allow_phantom_calls;
-        self
-    }
-
-    /// Returns true if this context is used for compiling a kernel.
-    pub fn is_kernel(&self) -> bool {
-        matches!(self.kind, ArtifactKind::Kernel)
-    }
-
-    /// Returns true if this context is used for compiling an executable.
-    pub fn is_executable(&self) -> bool {
-        matches!(self.kind, ArtifactKind::Executable)
-    }
-
-    /// Returns the type of artifact to produce with this context
-    pub fn kind(&self) -> ArtifactKind {
-        self.kind
-    }
-
-    /// Returns true if this context treats warning diagnostics as errors
-    #[inline(always)]
-    pub fn warnings_as_errors(&self) -> bool {
-        self.warnings_as_errors
-    }
-
-    /// Registers a "phantom" call to the procedure with the specified MAST root.
-    ///
-    /// A phantom call indicates that code for the procedure is not available. Executing a phantom
-    /// call will result in a runtime error. However, the VM may be able to execute a program with
-    /// phantom calls as long as the branches containing them are not taken.
-    ///
-    /// # Errors
-    /// Returns an error if phantom calls are not allowed in this assembly context.
-    pub fn register_phantom_call(
-        &mut self,
-        mast_root: Span<RpoDigest>,
-    ) -> Result<(), AssemblyError> {
-        if !self.allow_phantom_calls {
-            let source_file = self.unwrap_current_procedure().source_file().clone();
-            let (span, digest) = mast_root.into_parts();
-            Err(AssemblyError::PhantomCallsNotAllowed {
-                span,
-                source_file,
-                digest,
-            })
-        } else {
-            Ok(())
-        }
-    }
-
-    /// Registers a call to an externally-defined procedure which we have previously compiled.
-    ///
-    /// The call set of the callee is added to the call set of the procedure we are currently
-    /// compiling, to reflect that all of the code reachable from the callee is by extension
-    /// reachable by the caller.
-    pub fn register_external_call(
-        &mut self,
-        callee: &Procedure,
-        inlined: bool,
-        mast_forest: &MastForest,
-    ) -> Result<(), AssemblyError> {
-        let context = self.unwrap_current_procedure_mut();
-
-        // If we call the callee, it's callset is by extension part of our callset
-        context.extend_callset(callee.callset().iter().cloned());
-
-        // If the callee is not being inlined, add it to our callset
-        if !inlined {
-            context.insert_callee(callee.mast_root(mast_forest));
-        }
-
-        Ok(())
-    }
-}
-
-// PROCEDURE CONTEXT
-// ================================================================================================
-
-pub(super) struct ProcedureContext {
+pub struct ProcedureContext {
     span: SourceSpan,
     source_file: Option<Arc<SourceFile>>,
     gid: GlobalProcedureIndex,
@@ -198,7 +19,7 @@ pub(super) struct ProcedureContext {
 }
 
 impl ProcedureContext {
-    pub(super) fn new(
+    pub fn new(
         gid: GlobalProcedureIndex,
         name: FullyQualifiedProcedureName,
         visibility: Visibility,
@@ -214,26 +35,26 @@ impl ProcedureContext {
         }
     }
 
-    pub(super) fn with_span(mut self, span: SourceSpan) -> Self {
+    pub fn with_span(mut self, span: SourceSpan) -> Self {
         self.span = span;
         self
     }
 
-    pub(super) fn with_source_file(mut self, source_file: Option<Arc<SourceFile>>) -> Self {
+    pub fn with_source_file(mut self, source_file: Option<Arc<SourceFile>>) -> Self {
         self.source_file = source_file;
         self
     }
 
-    pub(super) fn with_num_locals(mut self, num_locals: u16) -> Self {
+    pub fn with_num_locals(mut self, num_locals: u16) -> Self {
         self.num_locals = num_locals;
         self
     }
 
-    pub(crate) fn insert_callee(&mut self, callee: RpoDigest) {
+    pub fn insert_callee(&mut self, callee: RpoDigest) {
         self.callset.insert(callee);
     }
 
-    pub(crate) fn extend_callset<I>(&mut self, callees: I)
+    pub fn extend_callset<I>(&mut self, callees: I)
     where
         I: IntoIterator<Item = RpoDigest>,
     {
@@ -265,6 +86,28 @@ impl ProcedureContext {
         self.visibility.is_syscall()
     }
 
+    /// Registers a call to an externally-defined procedure which we have previously compiled.
+    ///
+    /// The call set of the callee is added to the call set of the procedure we are currently
+    /// compiling, to reflect that all of the code reachable from the callee is by extension
+    /// reachable by the caller.
+    pub fn register_external_call(
+        &mut self,
+        callee: &Procedure,
+        inlined: bool,
+        mast_forest: &MastForest,
+    ) -> Result<(), AssemblyError> {
+        // If we call the callee, it's callset is by extension part of our callset
+        self.extend_callset(callee.callset().iter().cloned());
+
+        // If the callee is not being inlined, add it to our callset
+        if !inlined {
+            self.insert_callee(callee.mast_root(mast_forest));
+        }
+
+        Ok(())
+    }
+
     pub fn into_procedure(self, body_node_id: MastNodeId) -> Box<Procedure> {
         let procedure =
             Procedure::new(self.name, self.visibility, self.num_locals as u32, body_node_id)

assembly/src/assembler/instruction/env_ops.rs

@@ -1,5 +1,5 @@
-use super::{mem_ops::local_to_absolute_addr, push_felt, AssemblyContext, BasicBlockBuilder};
-use crate::{AssemblyError, Felt, Spanned};
+use super::{mem_ops::local_to_absolute_addr, push_felt, BasicBlockBuilder};
+use crate::{assembler::context::ProcedureContext, AssemblyError, Felt, Spanned};
 use vm_core::Operation::*;
 
 // CONSTANT INPUTS
@@ -41,9 +41,9 @@ where
 pub fn locaddr(
     span: &mut BasicBlockBuilder,
     index: u16,
-    context: &AssemblyContext,
+    proc_ctx: &ProcedureContext,
 ) -> Result<(), AssemblyError> {
-    local_to_absolute_addr(span, index, context.unwrap_current_procedure().num_locals())
+    local_to_absolute_addr(span, index, proc_ctx.num_locals())
 }
 
 /// Appends CALLER operation to the span which puts the hash of the function which initiated the
@@ -53,13 +53,12 @@ pub fn locaddr(
 /// Returns an error if the instruction is being executed outside of kernel context.
 pub fn caller(
     span: &mut BasicBlockBuilder,
-    context: &AssemblyContext,
+    proc_ctx: &ProcedureContext,
 ) -> Result<(), AssemblyError> {
-    let current_procedure = context.unwrap_current_procedure();
-    if !current_procedure.is_kernel() {
+    if !proc_ctx.is_kernel() {
         return Err(AssemblyError::CallerOutsideOfKernel {
-            span: current_procedure.span(),
-            source_file: current_procedure.source_file(),
+            span: proc_ctx.span(),
+            source_file: proc_ctx.source_file(),
         });
     }
     span.push_op(Caller);

assembly/src/assembler/instruction/field_ops.rs

@@ -1,5 +1,6 @@
-use super::{validate_param, AssemblyContext, BasicBlockBuilder};
+use super::{validate_param, BasicBlockBuilder};
 use crate::{
+    assembler::context::ProcedureContext,
     diagnostics::{RelatedError, Report},
     AssemblyError, Felt, Span, MAX_EXP_BITS, ONE, ZERO,
 };
@@ -88,11 +89,11 @@ pub fn mul_imm(span_builder: &mut BasicBlockBuilder, imm: Felt) {
 /// Returns an error if the immediate value is ZERO.
 pub fn div_imm(
     span_builder: &mut BasicBlockBuilder,
-    ctx: &mut AssemblyContext,
+    proc_ctx: &mut ProcedureContext,
     imm: Span<Felt>,
 ) -> Result<(), AssemblyError> {
     if imm == ZERO {
-        let source_file = ctx.unwrap_current_procedure().source_file();
+        let source_file = proc_ctx.source_file();
         let error = Report::new(crate::parser::ParsingError::DivisionByZero { span: imm.span() });
         return Err(if let Some(source_file) = source_file {
             AssemblyError::Other(RelatedError::new(error.with_source_code(source_file)))

assembly/src/assembler/instruction/mem_ops.rs

@@ -1,5 +1,5 @@
-use super::{push_felt, push_u32_value, validate_param, AssemblyContext, BasicBlockBuilder};
-use crate::{diagnostics::Report, AssemblyError};
+use super::{push_felt, push_u32_value, validate_param, BasicBlockBuilder};
+use crate::{assembler::context::ProcedureContext, diagnostics::Report, AssemblyError};
 use alloc::string::ToString;
 use vm_core::{Felt, Operation::*};
 
@@ -22,15 +22,15 @@ use vm_core::{Felt, Operation::*};
 /// the number of procedure locals.
 pub fn mem_read(
     span: &mut BasicBlockBuilder,
-    context: &AssemblyContext,
+    proc_ctx: &ProcedureContext,
     addr: Option<u32>,
     is_local: bool,
     is_single: bool,
 ) -> Result<(), AssemblyError> {
     // if the address was provided as an immediate value, put it onto the stack
     if let Some(addr) = addr {
         if is_local {
-            let num_locals = context.unwrap_current_procedure().num_locals();
+            let num_locals = proc_ctx.num_locals();
             local_to_absolute_addr(span, addr as u16, num_locals)?;
         } else {
             push_u32_value(span, addr);
@@ -73,14 +73,13 @@ pub fn mem_read(
 /// the number of procedure locals.
 pub fn mem_write_imm(
     span: &mut BasicBlockBuilder,
-    context: &AssemblyContext,
+    proc_ctx: &ProcedureContext,
     addr: u32,
     is_local: bool,
     is_single: bool,
 ) -> Result<(), AssemblyError> {
     if is_local {
-        let num_locals = context.unwrap_current_procedure().num_locals();
-        local_to_absolute_addr(span, addr as u16, num_locals)?;
+        local_to_absolute_addr(span, addr as u16, proc_ctx.num_locals())?;
     } else {
         push_u32_value(span, addr);
     }