refactor: remove ProcedureCache from assembler

CHANGELOG.md

@@ -21,6 +21,7 @@
 - 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)
+- Remove `ProcedureCache` from the assembler (#1411).
 
 #### Changed
 

assembly/src/assembler/basic_block_builder.rs

@@ -1,12 +1,11 @@
+use crate::AssemblyError;
+
 use super::{
-    context::ProcedureContext, mast_forest_builder::MastForestBuilder, BodyWrapper, Decorator,
-    DecoratorList, Instruction,
+    mast_forest_builder::MastForestBuilder, BodyWrapper, Decorator, DecoratorList, Instruction,
+    ProcedureContext,
 };
 use alloc::{borrow::Borrow, string::ToString, vec::Vec};
-use vm_core::{
-    mast::{MastForestError, MastNodeId},
-    AdviceInjector, AssemblyOp, Operation,
-};
+use vm_core::{mast::MastNodeId, AdviceInjector, AssemblyOp, Operation};
 
 // BASIC BLOCK BUILDER
 // ================================================================================================
@@ -129,7 +128,7 @@ impl BasicBlockBuilder {
     pub fn make_basic_block(
         &mut self,
         mast_forest_builder: &mut MastForestBuilder,
-    ) -> Result<Option<MastNodeId>, MastForestError> {
+    ) -> Result<Option<MastNodeId>, AssemblyError> {
         if !self.ops.is_empty() {
             let ops = self.ops.drain(..).collect();
             let decorators = self.decorators.drain(..).collect();
@@ -157,7 +156,7 @@ impl BasicBlockBuilder {
     pub fn try_into_basic_block(
         mut self,
         mast_forest_builder: &mut MastForestBuilder,
-    ) -> Result<Option<MastNodeId>, MastForestError> {
+    ) -> Result<Option<MastNodeId>, AssemblyError> {
         self.ops.append(&mut self.epilogue);
         self.make_basic_block(mast_forest_builder)
     }

assembly/src/assembler/id.rs

@@ -16,10 +16,10 @@ use crate::ast::ProcedureIndex;
 /// </div>
 ///
 /// In addition to the [super::ModuleGraph], these indices are also used with an instance of a
-/// [super::ProcedureCache]. This is because the [super::ModuleGraph] and [super::ProcedureCache]
-/// instances are paired, i.e. the [super::ModuleGraph] stores the syntax trees and call graph
-/// analysis for a program, while the [super::ProcedureCache] caches the compiled
-/// [super::Procedure]s for the same program, as derived from the corresponding graph.
+/// [super::MastForestBuilder]. This is because the [super::ModuleGraph] and
+/// [super::MastForestBuilder] instances are paired, i.e. the [super::ModuleGraph] stores the syntax
+/// trees and call graph analysis for a program, while the [super::MastForestBuilder] caches the
+/// compiled [super::Procedure]s for the same program, as derived from the corresponding graph.
 ///
 /// This is intended for use when we are doing global inter-procedural analysis on a (possibly
 /// growable) set of modules. It is expected that the index of a module in the set, as well as the

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::context::ProcedureContext, AssemblyError, Felt, Spanned};
+use crate::{assembler::ProcedureContext, AssemblyError, Felt, Spanned};
 use vm_core::Operation::*;
 
 // CONSTANT INPUTS

assembly/src/assembler/instruction/field_ops.rs

@@ -1,6 +1,6 @@
 use super::{validate_param, BasicBlockBuilder};
 use crate::{
-    assembler::context::ProcedureContext,
+    assembler::ProcedureContext,
     diagnostics::{RelatedError, Report},
     AssemblyError, Felt, Span, MAX_EXP_BITS, ONE, ZERO,
 };

assembly/src/assembler/instruction/mem_ops.rs

@@ -1,5 +1,5 @@
 use super::{push_felt, push_u32_value, validate_param, BasicBlockBuilder};
-use crate::{assembler::context::ProcedureContext, diagnostics::Report, AssemblyError};
+use crate::{assembler::ProcedureContext, diagnostics::Report, AssemblyError};
 use alloc::string::ToString;
 use vm_core::{Felt, Operation::*};
 

assembly/src/assembler/instruction/mod.rs

@@ -1,6 +1,6 @@
 use super::{
-    ast::InvokeKind, context::ProcedureContext, mast_forest_builder::MastForestBuilder, Assembler,
-    BasicBlockBuilder, Felt, Instruction, Operation, ONE, ZERO,
+    ast::InvokeKind, mast_forest_builder::MastForestBuilder, Assembler, BasicBlockBuilder, Felt,
+    Instruction, Operation, ProcedureContext, ONE, ZERO,
 };
 use crate::{diagnostics::Report, utils::bound_into_included_u64, AssemblyError};
 use core::ops::RangeBounds;
@@ -391,7 +391,7 @@ impl Assembler {
             Instruction::DynExec => return self.dynexec(mast_forest_builder),
             Instruction::DynCall => return self.dyncall(mast_forest_builder),
             Instruction::ProcRef(ref callee) => {
-                self.procref(callee, proc_ctx, span_builder, mast_forest_builder.forest())?
+                self.procref(callee, proc_ctx, span_builder, mast_forest_builder)?
             }
 
             // ----- debug decorators -------------------------------------------------------------

assembly/src/assembler/instruction/procedures.rs

@@ -1,12 +1,12 @@
 use super::{Assembler, BasicBlockBuilder, Operation};
 use crate::{
-    assembler::{context::ProcedureContext, mast_forest_builder::MastForestBuilder},
+    assembler::{mast_forest_builder::MastForestBuilder, ProcedureContext},
     ast::{InvocationTarget, InvokeKind},
     AssemblyError, RpoDigest, SourceSpan, Spanned,
 };
 
 use smallvec::SmallVec;
-use vm_core::mast::{MastForest, MastNodeId};
+use vm_core::mast::MastNodeId;
 
 /// Procedure Invocation
 impl Assembler {
@@ -18,7 +18,7 @@ impl Assembler {
         mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         let span = callee.span();
-        let digest = self.resolve_target(kind, callee, proc_ctx, mast_forest_builder.forest())?;
+        let digest = self.resolve_target(kind, callee, proc_ctx, mast_forest_builder)?;
         self.invoke_mast_root(kind, span, digest, proc_ctx, mast_forest_builder)
     }
 
@@ -31,12 +31,11 @@ impl Assembler {
         mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Option<MastNodeId>, AssemblyError> {
         // Get the procedure from the assembler
-        let cache = &self.procedure_cache;
         let current_source_file = proc_ctx.source_file();
 
         // If the procedure is cached, register the call to ensure the callset
         // is updated correctly.
-        match cache.get_by_mast_root(&mast_root) {
+        match mast_forest_builder.find_procedure(&mast_root) {
             Some(proc) if matches!(kind, InvokeKind::SysCall) => {
                 // Verify if this is a syscall, that the callee is a kernel procedure
                 //
@@ -69,11 +68,9 @@ impl Assembler {
                             })
                         }
                     })?;
-                proc_ctx.register_external_call(&proc, false, mast_forest_builder.forest())?;
-            }
-            Some(proc) => {
-                proc_ctx.register_external_call(&proc, false, mast_forest_builder.forest())?
+                proc_ctx.register_external_call(&proc, false)?;
             }
+            Some(proc) => proc_ctx.register_external_call(&proc, false)?,
             None if matches!(kind, InvokeKind::SysCall) => {
                 return Err(AssemblyError::UnknownSysCallTarget {
                     span,
@@ -91,7 +88,7 @@ impl Assembler {
                     // procedures, such that when we assemble a procedure, all
                     // procedures that it calls will have been assembled, and
                     // hence be present in the `MastForest`.
-                    match mast_forest_builder.find_procedure_root(mast_root) {
+                    match mast_forest_builder.find_procedure_node_id(mast_root) {
                         Some(root) => root,
                         None => {
                             // If the MAST root called isn't known to us, make it an external
@@ -101,7 +98,7 @@ impl Assembler {
                     }
                 }
                 InvokeKind::Call => {
-                    let callee_id = match mast_forest_builder.find_procedure_root(mast_root) {
+                    let callee_id = match mast_forest_builder.find_procedure_node_id(mast_root) {
                         Some(callee_id) => callee_id,
                         None => {
                             // If the MAST root called isn't known to us, make it an external
@@ -113,7 +110,7 @@ impl Assembler {
                     mast_forest_builder.ensure_call(callee_id)?
                 }
                 InvokeKind::SysCall => {
-                    let callee_id = match mast_forest_builder.find_procedure_root(mast_root) {
+                    let callee_id = match mast_forest_builder.find_procedure_node_id(mast_root) {
                         Some(callee_id) => callee_id,
                         None => {
                             // If the MAST root called isn't known to us, make it an external
@@ -158,23 +155,25 @@ impl Assembler {
         callee: &InvocationTarget,
         proc_ctx: &mut ProcedureContext,
         span_builder: &mut BasicBlockBuilder,
-        mast_forest: &MastForest,
+        mast_forest_builder: &MastForestBuilder,
     ) -> Result<(), AssemblyError> {
-        let digest = self.resolve_target(InvokeKind::Exec, callee, proc_ctx, mast_forest)?;
-        self.procref_mast_root(digest, proc_ctx, span_builder, mast_forest)
+        let digest =
+            self.resolve_target(InvokeKind::Exec, callee, proc_ctx, mast_forest_builder)?;
+        self.procref_mast_root(digest, proc_ctx, span_builder, mast_forest_builder)
     }
 
     fn procref_mast_root(
         &self,
         mast_root: RpoDigest,
         proc_ctx: &mut ProcedureContext,
         span_builder: &mut BasicBlockBuilder,
-        mast_forest: &MastForest,
+        mast_forest_builder: &MastForestBuilder,
     ) -> Result<(), AssemblyError> {
         // Add the root to the callset to be able to use dynamic instructions
         // with the referenced procedure later
-        if let Some(proc) = self.procedure_cache.get_by_mast_root(&mast_root) {
-            proc_ctx.register_external_call(&proc, false, mast_forest)?;
+
+        if let Some(proc) = mast_forest_builder.find_procedure(&mast_root) {
+            proc_ctx.register_external_call(&proc, false)?;
         }
 
         // Create an array with `Push` operations containing root elements

assembly/src/assembler/instruction/u32_ops.rs

@@ -1,6 +1,6 @@
 use super::{field_ops::append_pow2_op, push_u32_value, validate_param, BasicBlockBuilder};
 use crate::{
-    assembler::context::ProcedureContext,
+    assembler::ProcedureContext,
     diagnostics::{RelatedError, Report},
     AssemblyError, Span, MAX_U32_ROTATE_VALUE, MAX_U32_SHIFT_VALUE,
 };