diff --git a/assembly/src/assembler/basic_block_builder.rs b/assembly/src/assembler/basic_block_builder.rs
index dfbdbd84c2..a48cceea08 100644
--- a/assembly/src/assembler/basic_block_builder.rs
+++ b/assembly/src/assembler/basic_block_builder.rs
@@ -3,10 +3,7 @@ use super::{
Instruction,
};
use alloc::{borrow::Borrow, string::ToString, vec::Vec};
-use vm_core::{
- mast::{MastNode, MastNodeId},
- AdviceInjector, AssemblyOp, Operation,
-};
+use vm_core::{mast::MastNodeId, AdviceInjector, AssemblyOp, Operation};
// BASIC BLOCK BUILDER
// ================================================================================================
@@ -134,10 +131,9 @@ impl BasicBlockBuilder {
let ops = self.ops.drain(..).collect();
let decorators = self.decorators.drain(..).collect();
- let basic_block_node = MastNode::new_basic_block_with_decorators(ops, decorators);
- let basic_block_node_id = mast_forest_builder.ensure_node(basic_block_node);
+ let block_node_id = mast_forest_builder.ensure_block(ops, Some(decorators));
- Some(basic_block_node_id)
+ Some(block_node_id)
} else if !self.decorators.is_empty() {
// this is a bug in the assembler. we shouldn't have decorators added without their
// associated operations
diff --git a/assembly/src/assembler/instruction/procedures.rs b/assembly/src/assembler/instruction/procedures.rs
index 9894a82092..aa14943997 100644
--- a/assembly/src/assembler/instruction/procedures.rs
+++ b/assembly/src/assembler/instruction/procedures.rs
@@ -6,7 +6,7 @@ use crate::{
};
use smallvec::SmallVec;
-use vm_core::mast::{MastForest, MastNode, MastNodeId};
+use vm_core::mast::{MastForest, MastNodeId};
/// Procedure Invocation
impl Assembler {
@@ -94,8 +94,7 @@ impl Assembler {
mast_forest_builder.find_procedure_root(mast_root).unwrap_or_else(|| {
// If the MAST root called isn't known to us, make it an external
// reference.
- let external_node = MastNode::new_external(mast_root);
- mast_forest_builder.ensure_node(external_node)
+ mast_forest_builder.ensure_external(mast_root)
})
}
InvokeKind::Call => {
@@ -103,25 +102,20 @@ impl Assembler {
mast_forest_builder.find_procedure_root(mast_root).unwrap_or_else(|| {
// If the MAST root called isn't known to us, make it an external
// reference.
- let external_node = MastNode::new_external(mast_root);
- mast_forest_builder.ensure_node(external_node)
+ mast_forest_builder.ensure_external(mast_root)
});
- let call_node = MastNode::new_call(callee_id, mast_forest_builder.forest());
- mast_forest_builder.ensure_node(call_node)
+ mast_forest_builder.ensure_call(callee_id)
}
InvokeKind::SysCall => {
let callee_id =
mast_forest_builder.find_procedure_root(mast_root).unwrap_or_else(|| {
// If the MAST root called isn't known to us, make it an external
// reference.
- let external_node = MastNode::new_external(mast_root);
- mast_forest_builder.ensure_node(external_node)
+ mast_forest_builder.ensure_external(mast_root)
});
- let syscall_node =
- MastNode::new_syscall(callee_id, mast_forest_builder.forest());
- mast_forest_builder.ensure_node(syscall_node)
+ mast_forest_builder.ensure_syscall(callee_id)
}
}
};
@@ -134,7 +128,7 @@ impl Assembler {
&self,
mast_forest_builder: &mut MastForestBuilder,
) -> Result