Skip to content

Commit

Permalink
feat: make Assembler single-use (#1409)
Browse files Browse the repository at this point in the history
  • Loading branch information
@plafer
plafer authored Jul 23, 2024
1 parent b10567b commit f187552
Show file tree
Hide file tree
Showing 25 changed files with 249 additions and 942 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
8 changes: 4 additions & 4 deletions assembly/src/assembler/basic_block_builder.rs
View file
Original file line number Diff line number Diff line change
@@ -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::{
Expand Down Expand Up @@ -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();
Expand Down
259 changes: 58 additions & 201 deletions assembly/src/assembler/context.rs
View file
Original file line number Diff line number Diff line change
@@ -1,193 +1,18 @@
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 {
/// Information about a procedure currently being compiled.
pub struct ProcedureContext {
span: SourceSpan,
source_file: Option<Arc<SourceFile>>,
gid: GlobalProcedureIndex,
Expand All @@ -197,8 +22,10 @@ pub(super) struct ProcedureContext {
callset: CallSet,
}

// ------------------------------------------------------------------------------------------------
/// Constructors
impl ProcedureContext {
pub(super) fn new(
pub fn new(
gid: GlobalProcedureIndex,
name: FullyQualifiedProcedureName,
visibility: Visibility,
Expand All @@ -214,36 +41,25 @@ impl ProcedureContext {
}
}

pub(super) fn with_span(mut self, span: SourceSpan) -> Self {
self.span = span;
pub fn with_num_locals(mut self, num_locals: u16) -> Self {
self.num_locals = num_locals;
self
}

pub(super) fn with_source_file(mut self, source_file: Option<Arc<SourceFile>>) -> Self {
self.source_file = source_file;
pub fn with_span(mut self, span: SourceSpan) -> Self {
self.span = span;
self
}

pub(super) fn with_num_locals(mut self, num_locals: u16) -> Self {
self.num_locals = num_locals;
pub fn with_source_file(mut self, source_file: Option<Arc<SourceFile>>) -> Self {
self.source_file = source_file;
self
}
}

pub(crate) fn insert_callee(&mut self, callee: RpoDigest) {
self.callset.insert(callee);
}

pub(crate) fn extend_callset<I>(&mut self, callees: I)
where
I: IntoIterator<Item = RpoDigest>,
{
self.callset.extend(callees);
}

pub fn num_locals(&self) -> u16 {
self.num_locals
}

// ------------------------------------------------------------------------------------------------
/// Public accessors
impl ProcedureContext {
pub fn id(&self) -> GlobalProcedureIndex {
self.gid
}
Expand All @@ -252,6 +68,10 @@ impl ProcedureContext {
&self.name
}

pub fn num_locals(&self) -> u16 {
self.num_locals
}

#[allow(unused)]
pub fn module(&self) -> &LibraryPath {
&self.name.module
Expand All @@ -264,6 +84,43 @@ impl ProcedureContext {
pub fn is_kernel(&self) -> bool {
self.visibility.is_syscall()
}
}

// ------------------------------------------------------------------------------------------------
/// State mutators
impl ProcedureContext {
pub fn insert_callee(&mut self, callee: RpoDigest) {
self.callset.insert(callee);
}

pub fn extend_callset<I>(&mut self, callees: I)
where
I: IntoIterator<Item = RpoDigest>,
{
self.callset.extend(callees);
}

/// 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 =
Expand Down
17 changes: 8 additions & 9 deletions assembly/src/assembler/instruction/env_ops.rs
View file
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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);
Expand Down
7 changes: 4 additions & 3 deletions assembly/src/assembler/instruction/field_ops.rs
View file
Original file line number Diff line number Diff line change
@@ -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,
};
Expand Down Expand Up @@ -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)))
Expand Down
Loading

0 comments on commit f187552

Please sign in to comment.