feat: add Call block to program MAST

assembly/src/context/mod.rs

@@ -14,22 +14,29 @@ use super::{BTreeMap, CodeBlock, ProcMap, Procedure, String, ToString, MODULE_PA
 pub struct AssemblyContext<'a> {
     local_procs: ProcMap,
     imported_procs: BTreeMap<String, &'a Procedure>,
+    in_debug_mode: bool,
 }
 
 impl<'a> AssemblyContext<'a> {
     // CONSTRUCTOR
     // --------------------------------------------------------------------------------------------
     /// Returns a new empty [AssemblyContext].
-    pub fn new() -> Self {
+    pub fn new(in_debug_mode: bool) -> Self {
         Self {
             local_procs: BTreeMap::new(),
             imported_procs: BTreeMap::new(),
+            in_debug_mode,
         }
     }
 
     // STATE ACCESSORS
     // --------------------------------------------------------------------------------------------
 
+    /// Returns true if the assembly context was instantiated in debug mode.
+    pub fn in_debug_mode(&self) -> bool {
+        self.in_debug_mode
+    }
+
     /// Returns true if a procedure with the specified label exists in this context.
     pub fn contains_proc(&self, label: &str) -> bool {
         self.local_procs.contains_key(label) || self.imported_procs.contains_key(label)

assembly/src/lib.rs

@@ -10,7 +10,7 @@ use vm_core::{
         collections::{BTreeMap, Vec},
         string::{String, ToString},
     },
-    Library, Program,
+    CodeBlockTable, Library, Program,
 };
 use vm_stdlib::StdLibrary;
 
@@ -73,19 +73,20 @@ impl Assembler {
     /// on Miden VM.
     pub fn compile(&self, source: &str) -> Result<Program, AssemblyError> {
         let mut tokens = TokenStream::new(source)?;
-        let mut context = AssemblyContext::new();
+        let mut context = AssemblyContext::new(self.in_debug_mode);
+        let mut cb_table = CodeBlockTable::default();
 
         // parse imported modules (if any), and add exported procedures from these modules to the
         // current context; since we are in the root context here, we initialize dependency chain
         // with an empty vector.
-        self.parse_imports(&mut tokens, &mut context, &mut Vec::new())?;
+        self.parse_imports(&mut tokens, &mut context, &mut Vec::new(), &mut cb_table)?;
 
         // parse locally defined procedures (if any), and add these procedures to the current
         // context
         while let Some(token) = tokens.read() {
             let proc = match token.parts()[0] {
                 Token::PROC | Token::EXPORT => {
-                    Procedure::parse(&mut tokens, &context, false, self.in_debug_mode)?
+                    Procedure::parse(&mut tokens, &context, &mut cb_table, false)?
                 }
                 _ => break,
             };
@@ -101,8 +102,8 @@ impl Assembler {
         }
 
         // parse program body and return the resulting program
-        let program_root = parse_program(&mut tokens, &context, self.in_debug_mode)?;
-        Ok(Program::new(program_root))
+        let program_root = parse_program(&mut tokens, &context, &mut cb_table)?;
+        Ok(Program::with_table(program_root, cb_table))
     }
 
     // IMPORT PARSERS
@@ -126,6 +127,7 @@ impl Assembler {
         tokens: &mut TokenStream,
         context: &mut AssemblyContext<'a>,
         dep_chain: &mut Vec<String>,
+        cb_table: &mut CodeBlockTable,
     ) -> Result<(), AssemblyError> {
         // read tokens from the token stream until all `use` tokens are consumed
         while let Some(token) = tokens.read() {
@@ -152,7 +154,7 @@ impl Assembler {
                             self.stdlib.get_module_source(module_path).map_err(|_| {
                                 AssemblyError::missing_import_source(token, module_path)
                             })?;
-                        self.parse_module(module_source, module_path, dep_chain)?;
+                        self.parse_module(module_source, module_path, dep_chain, cb_table)?;
                     }
 
                     // get procedures from the module at the specified path; we are guaranteed to
@@ -190,20 +192,21 @@ impl Assembler {
         source: &str,
         path: &str,
         dep_chain: &mut Vec<String>,
+        cb_table: &mut CodeBlockTable,
     ) -> Result<(), AssemblyError> {
         let mut tokens = TokenStream::new(source)?;
-        let mut context = AssemblyContext::new();
+        let mut context = AssemblyContext::new(self.in_debug_mode);
 
         // parse imported modules (if any), and add exported procedures from these modules to
         // the current context
-        self.parse_imports(&mut tokens, &mut context, dep_chain)?;
+        self.parse_imports(&mut tokens, &mut context, dep_chain, cb_table)?;
 
         // parse procedures defined in the module, and add these procedures to the current
         // context
         while let Some(token) = tokens.read() {
             let proc = match token.parts()[0] {
                 Token::PROC | Token::EXPORT => {
-                    Procedure::parse(&mut tokens, &context, true, self.in_debug_mode)?
+                    Procedure::parse(&mut tokens, &context, cb_table, true)?
                 }
                 _ => break,
             };
@@ -246,7 +249,7 @@ impl Default for Assembler {
 fn parse_program(
     tokens: &mut TokenStream,
     context: &AssemblyContext,
-    in_debug_mode: bool,
+    cb_table: &mut CodeBlockTable,
 ) -> Result<CodeBlock, AssemblyError> {
     let program_start = tokens.pos();
     // consume the 'begin' token
@@ -255,7 +258,7 @@ fn parse_program(
     tokens.advance();
 
     // parse the program body
-    let root = parse_code_blocks(tokens, context, 0, in_debug_mode)?;
+    let root = parse_code_blocks(tokens, context, cb_table, 0)?;
 
     // consume the 'end' token
     match tokens.read() {

assembly/src/parsers/blocks.rs

@@ -2,7 +2,7 @@ use super::{
     parse_op_token, AssemblyContext, AssemblyError, CodeBlock, Operation, String, Token,
     TokenStream, Vec,
 };
-use vm_core::{utils::group_vector_elements, DecoratorList};
+use vm_core::{utils::group_vector_elements, CodeBlockTable, DecoratorList};
 
 // BLOCK PARSER
 // ================================================================================================
@@ -11,8 +11,8 @@ use vm_core::{utils::group_vector_elements, DecoratorList};
 pub fn parse_code_blocks(
     tokens: &mut TokenStream,
     context: &AssemblyContext,
+    cb_table: &mut CodeBlockTable,
     num_proc_locals: u32,
-    in_debug_mode: bool,
 ) -> Result<CodeBlock, AssemblyError> {
     // make sure there is something to be read
     let start_pos = tokens.pos();
@@ -23,7 +23,7 @@ pub fn parse_code_blocks(
     // parse the sequence of blocks and add each block to the list
     let mut blocks = Vec::new();
     while let Some(parser) = BlockParser::next(tokens)? {
-        let block = parser.parse(tokens, context, num_proc_locals, in_debug_mode)?;
+        let block = parser.parse(tokens, context, cb_table, num_proc_locals)?;
         blocks.push(block);
     }
 
@@ -48,6 +48,7 @@ enum BlockParser {
     While,
     Repeat(u32),
     Exec(String),
+    Call(String),
 }
 
 impl BlockParser {
@@ -56,12 +57,12 @@ impl BlockParser {
         &self,
         tokens: &mut TokenStream,
         context: &AssemblyContext,
+        cb_table: &mut CodeBlockTable,
         num_proc_locals: u32,
-        in_debug_mode: bool,
     ) -> Result<CodeBlock, AssemblyError> {
         match self {
+            // ------------------------------------------------------------------------------------
             Self::Span => {
-                // --------------------------------------------------------------------------------
                 let mut span_ops = Vec::new();
                 let mut decorators = DecoratorList::new();
                 while let Some(op) = tokens.read() {
@@ -73,20 +74,20 @@ impl BlockParser {
                         &mut span_ops,
                         num_proc_locals,
                         &mut decorators,
-                        in_debug_mode,
+                        context.in_debug_mode(),
                     )?;
                     tokens.advance();
                 }
                 Ok(CodeBlock::new_span_with_decorators(span_ops, decorators))
             }
+            // ------------------------------------------------------------------------------------
             Self::IfElse => {
-                // --------------------------------------------------------------------------------
                 // record start of the if-else block and consume the 'if' token
                 let if_start = tokens.pos();
                 tokens.advance();
 
                 // read the `if` clause
-                let t_branch = parse_code_blocks(tokens, context, num_proc_locals, in_debug_mode)?;
+                let t_branch = parse_code_blocks(tokens, context, cb_table, num_proc_locals)?;
 
                 // build the `else` clause; if the else clause is specified, then read it;
                 // otherwise, set to a Span with a single noop
@@ -100,7 +101,7 @@ impl BlockParser {
 
                             // parse the `false` branch
                             let f_branch =
-                                parse_code_blocks(tokens, context, num_proc_locals, in_debug_mode)?;
+                                parse_code_blocks(tokens, context, cb_table, num_proc_locals)?;
 
                             // consume the `end` token
                             match tokens.read() {
@@ -143,14 +144,14 @@ impl BlockParser {
 
                 Ok(CodeBlock::new_split(t_branch, f_branch))
             }
+            // ------------------------------------------------------------------------------------
             Self::While => {
-                // --------------------------------------------------------------------------------
                 // record start of the while block and consume the 'while' token
                 let while_start = tokens.pos();
                 tokens.advance();
 
                 // read the loop body
-                let loop_body = parse_code_blocks(tokens, context, num_proc_locals, in_debug_mode)?;
+                let loop_body = parse_code_blocks(tokens, context, cb_table, num_proc_locals)?;
 
                 // consume the `end` token
                 match tokens.read() {
@@ -169,14 +170,14 @@ impl BlockParser {
 
                 Ok(CodeBlock::new_loop(loop_body))
             }
+            // ------------------------------------------------------------------------------------
             Self::Repeat(iter_count) => {
-                // --------------------------------------------------------------------------------
                 // record start of the repeat block and consume the 'repeat' token
                 let repeat_start = tokens.pos();
                 tokens.advance();
 
                 // read the loop body
-                let loop_body = parse_code_blocks(tokens, context, num_proc_locals, in_debug_mode)?;
+                let loop_body = parse_code_blocks(tokens, context, cb_table, num_proc_locals)?;
 
                 // consume the `end` token
                 match tokens.read() {
@@ -205,17 +206,29 @@ impl BlockParser {
                     Ok(combine_blocks(blocks))
                 }
             }
+            // ------------------------------------------------------------------------------------
             Self::Exec(label) => {
-                // --------------------------------------------------------------------------------
                 // retrieve the procedure block from the proc map and consume the 'exec' token
-                let proc_root = context
-                    .get_proc_code(label)
-                    .ok_or_else(|| {
-                        AssemblyError::undefined_proc(tokens.read().expect("no exec token"), label)
-                    })?
-                    .clone();
+                let proc_block = context.get_proc_code(label).ok_or_else(|| {
+                    AssemblyError::undefined_proc(tokens.read().expect("no exec token"), label)
+                })?;
                 tokens.advance();
-                Ok(proc_root)
+                Ok(proc_block.clone())
+            }
+            // ------------------------------------------------------------------------------------
+            Self::Call(label) => {
+                // retrieve the procedure block from the proc map and consume the 'call' token
+                let proc_block = context.get_proc_code(label).ok_or_else(|| {
+                    AssemblyError::undefined_proc(tokens.read().expect("no call token"), label)
+                })?;
+                tokens.advance();
+
+                // if the procedure hasn't been inserted into code block table yet, insert it
+                if !cb_table.has(proc_block.hash()) {
+                    cb_table.insert(proc_block.clone());
+                }
+
+                Ok(CodeBlock::new_call(proc_block.hash()))
             }
         }
     }
@@ -245,6 +258,10 @@ impl BlockParser {
                     let label = token.parse_exec()?;
                     Some(Self::Exec(label))
                 }
+                Token::CALL => {
+                    let label = token.parse_call()?;
+                    Some(Self::Call(label))
+                }
                 Token::END => {
                     token.validate_end()?;
                     None

assembly/src/procedures/mod.rs

@@ -1,6 +1,6 @@
 use super::{
-    combine_blocks, parse_code_blocks, AssemblyContext, AssemblyError, CodeBlock, String, Token,
-    TokenStream, Vec,
+    combine_blocks, parse_code_blocks, AssemblyContext, AssemblyError, CodeBlock, CodeBlockTable,
+    String, Token, TokenStream, Vec,
 };
 use vm_core::{Felt, Operation};
 
@@ -51,8 +51,8 @@ impl Procedure {
     pub fn parse(
         tokens: &mut TokenStream,
         context: &AssemblyContext,
+        cb_table: &mut CodeBlockTable,
         allow_export: bool,
-        in_debug_mode: bool,
     ) -> Result<Self, AssemblyError> {
         let proc_start = tokens.pos();
 
@@ -68,7 +68,7 @@ impl Procedure {
         tokens.advance();
 
         // parse procedure body, and handle memory allocation/deallocation of locals if any are declared
-        let code_root = parse_proc_blocks(tokens, context, num_locals, in_debug_mode)?;
+        let code_root = parse_proc_blocks(tokens, context, cb_table, num_locals)?;
 
         // consume the 'end' token
         match tokens.read() {
@@ -100,11 +100,11 @@ impl Procedure {
 pub fn parse_proc_blocks(
     tokens: &mut TokenStream,
     context: &AssemblyContext,
+    cb_table: &mut CodeBlockTable,
     num_proc_locals: u32,
-    in_debug_mode: bool,
 ) -> Result<CodeBlock, AssemblyError> {
     // parse the procedure body
-    let body = parse_code_blocks(tokens, context, num_proc_locals, in_debug_mode)?;
+    let body = parse_code_blocks(tokens, context, cb_table, num_proc_locals)?;
 
     if num_proc_locals == 0 {
         // if no allocation of locals is required, return the procedure body

assembly/src/tokens/mod.rs

@@ -27,6 +27,7 @@ impl<'a> Token<'a> {
     pub const WHILE: &'static str = "while";
     pub const REPEAT: &'static str = "repeat";
     pub const EXEC: &'static str = "exec";
+    pub const CALL: &'static str = "call";
     pub const END: &'static str = "end";
 
     // CONSTRUCTOR
@@ -74,6 +75,7 @@ impl<'a> Token<'a> {
                 | Self::WHILE
                 | Self::REPEAT
                 | Self::EXEC
+                | Self::CALL
                 | Self::END
         )
     }
@@ -192,6 +194,15 @@ impl<'a> Token<'a> {
         }
     }
 
+    pub fn parse_call(&self) -> Result<String, AssemblyError> {
+        assert_eq!(Self::CALL, self.parts[0], "not a call");
+        match self.num_parts() {
+            1 => Err(AssemblyError::missing_param(self)),
+            2 => validate_proc_invocation_label(self.parts[1], self),
+            _ => Err(AssemblyError::extra_param(self)),
+        }
+    }
+
     pub fn validate_end(&self) -> Result<(), AssemblyError> {
         assert_eq!(Self::END, self.parts[0], "not an end");
         if self.num_parts() > 1 {

core/src/lib.rs

@@ -14,7 +14,7 @@ pub mod range;
 pub use math::{fields::f64::BaseElement as Felt, ExtensionOf, FieldElement, StarkField};
 
 mod program;
-pub use program::{blocks as code_blocks, Library, Program};
+pub use program::{blocks as code_blocks, CodeBlockTable, Library, Program};
 
 mod operations;
 pub use operations::{