diff --git a/assembly/src/assembler/basic_block_builder.rs b/assembly/src/assembler/basic_block_builder.rs
index 2545d4eea8..5263faeff2 100644
--- a/assembly/src/assembler/basic_block_builder.rs
+++ b/assembly/src/assembler/basic_block_builder.rs
@@ -4,7 +4,7 @@ use super::{
};
use alloc::{borrow::Borrow, string::ToString, vec::Vec};
use vm_core::{
- mast::{MastForestError, MastNode, MastNodeId},
+ mast::{MastForestError, MastNodeId},
AdviceInjector, AssemblyOp, Operation,
};
@@ -134,8 +134,7 @@ 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 basic_block_node_id = mast_forest_builder.ensure_block(ops, Some(decorators))?;
Ok(Some(basic_block_node_id))
} else if !self.decorators.is_empty() {
diff --git a/assembly/src/assembler/instruction/procedures.rs b/assembly/src/assembler/instruction/procedures.rs
index d6a96f0f1b..ee655af7fd 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 {
@@ -96,8 +96,7 @@ impl Assembler {
None => {
// 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)?
}
}
}
@@ -107,13 +106,11 @@ impl Assembler {
None => {
// 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 = match mast_forest_builder.find_procedure_root(mast_root) {
@@ -121,14 +118,11 @@ impl Assembler {
None => {
// 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)?
}
}
};
@@ -141,7 +135,7 @@ impl Assembler {
&self,
mast_forest_builder: &mut MastForestBuilder,
) -> Result