feat: add execution context to block stack table

miden/tests/integration/flow_control/mod.rs

@@ -1,4 +1,4 @@
-use crate::build_test;
+use crate::{build_test, helpers::TestError};
 
 // SIMPLE FLOW CONTROL TESTS
 // ================================================================================================
@@ -123,19 +123,42 @@ fn if_in_loop_in_if() {
 
 #[test]
 fn local_fn_call() {
+    // returning from a function with non-empty overflow table should result in an error
     let source = "
         proc.foo
-            add
+            push.1
         end
 
         begin
             call.foo
         end";
 
-    let test = build_test!(source, &[1, 2]);
-    test.expect_stack(&[3]);
+    let expected_err = TestError::ExecutionError("InvalidStackDepthOnReturn(17)");
+    build_test!(source, &[1, 2]).expect_error(expected_err);
+
+    // dropping values from the stack in the current execution context should not affect values
+    // in the overflow table from the parent execution context
+    let source = "
+        proc.foo
+            repeat.20
+                drop
+            end
+        end
+
+        begin
+            push.18
+            call.foo
+            repeat.16
+                drop
+            end
+        end";
+
+    let inputs = (1_u64..18).collect::<Vec<_>>();
+
+    let test = build_test!(source, &inputs);
+    test.expect_stack(&[2, 1]);
 
-    test.prove_and_verify(vec![1, 2], false);
+    test.prove_and_verify(inputs, false);
 }
 
 #[test]

processor/src/decoder/aux_hints.rs

@@ -276,21 +276,30 @@ pub enum OpGroupTableUpdate {
 // ================================================================================================
 
 /// Describes a single entry in the block stack table. An entry in the block stack table is a tuple
-/// (block_id, parent_id, is_loop).
+/// (block_id, parent_id, is_loop, parent_ctx, fmp, stack_depth, next_overflow_addr).
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct BlockStackTableRow {
     block_id: Felt,
     parent_id: Felt,
     is_loop: bool,
+    parent_ctx: u32,
+    fmp: Felt,
+    stack_depth: u32,
+    next_overflow_addr: Felt,
 }
 
 impl BlockStackTableRow {
     /// Returns a new [BlockStackTableRow] instantiated from the specified block info.
     pub fn new(block_info: &BlockInfo) -> Self {
+        let ctx_info = block_info.ctx_info.unwrap_or_default();
         Self {
             block_id: block_info.addr,
             parent_id: block_info.parent_addr,
             is_loop: block_info.is_entered_loop() == ONE,
+            parent_ctx: ctx_info.parent_ctx,
+            fmp: ctx_info.fmp,
+            stack_depth: ctx_info.stack_depth,
+            next_overflow_addr: ctx_info.next_overflow_addr,
         }
     }
 
@@ -302,19 +311,47 @@ impl BlockStackTableRow {
             block_id,
             parent_id,
             is_loop,
+            parent_ctx: 0,
+            fmp: ZERO,
+            stack_depth: 0,
+            next_overflow_addr: ZERO,
+        }
+    }
+
+    #[cfg(test)]
+    /// Returns a new [BlockStackTableRow] corresponding to a CALL code block. This is used for
+    /// test purpose only.
+    pub fn new_test_with_ctx(
+        block_id: Felt,
+        parent_id: Felt,
+        is_loop: bool,
+        ctx_info: super::ExecutionContextInfo,
+    ) -> Self {
+        Self {
+            block_id,
+            parent_id,
+            is_loop,
+            parent_ctx: ctx_info.parent_ctx,
+            fmp: ctx_info.fmp,
+            stack_depth: ctx_info.stack_depth,
+            next_overflow_addr: ctx_info.next_overflow_addr,
         }
     }
 }
 
 impl LookupTableRow for BlockStackTableRow {
     /// Reduces this row to a single field element in the field specified by E. This requires
-    /// at least 4 alpha values.
+    /// at least 8 alpha values.
     fn to_value<E: FieldElement<BaseField = Felt>>(&self, alphas: &[E]) -> E {
         let is_loop = if self.is_loop { ONE } else { ZERO };
         alphas[0]
             + alphas[1].mul_base(self.block_id)
             + alphas[2].mul_base(self.parent_id)
             + alphas[3].mul_base(is_loop)
+            + alphas[4].mul_base(Felt::from(self.parent_ctx))
+            + alphas[5].mul_base(self.fmp)
+            + alphas[6].mul_base(Felt::from(self.stack_depth))
+            + alphas[7].mul_base(self.next_overflow_addr)
     }
 }
 

processor/src/decoder/block_stack.rs

@@ -14,44 +14,55 @@ impl BlockStack {
     // --------------------------------------------------------------------------------------------
 
     /// Pushes a new code block onto the block stack and returns the address of the block's parent.
-    /// block_type can be anything except for CALL block.
     ///
     /// The block is identified by its address, and we also need to know what type of a block this
-    /// is. Other information (i.e., the block's parent, whether the block is a body of
-    /// a loop or a first child of a JOIN block) is determined from the information already on the
-    /// stack.
-    ///
-    /// For non-CALL blocks, we set parent context and free memory pointer to zero values.
-    ///
-    /// # Panics
-    /// Panics if the block type is a CALL block.
-    pub fn push(&mut self, addr: Felt, block_type: BlockType) -> Felt {
-        assert_ne!(block_type, BlockType::Call, "cannot be a call block");
-        let (parent_addr, is_loop_body, is_first_child) = self.get_new_block_context();
-        self.blocks.push(BlockInfo {
-            addr,
-            block_type,
-            parent_addr,
-            parent_ctx: 0,
-            parent_fmp: ZERO,
-            is_loop_body,
-            is_first_child,
-        });
-        parent_addr
-    }
+    /// is. Additionally, for CALL blocks, execution context info must be provided. Other
+    /// information (i.e., the block's parent, whether the block is a body of a loop or a first
+    /// child of a JOIN block) is determined from the information already on the stack.
+    pub fn push(
+        &mut self,
+        addr: Felt,
+        block_type: BlockType,
+        ctx_info: Option<ExecutionContextInfo>,
+    ) -> Felt {
+        // make sure execution context was provided when expected
+        if block_type == BlockType::Call {
+            debug_assert!(
+                ctx_info.is_some(),
+                "no execution context provided for a CALL block"
+            );
+        } else {
+            debug_assert!(
+                ctx_info.is_none(),
+                "execution context provided for a non-CALL block"
+            );
+        }
+
+        // determine additional info about the new block based on its parent
+        let (parent_addr, is_loop_body, is_first_child) = match self.blocks.last() {
+            Some(parent) => match parent.block_type {
+                // if the parent block is a LOOP block, the new block must be a loop body
+                BlockType::Loop(loop_entered) => {
+                    debug_assert!(loop_entered, "parent is un-entered loop");
+                    (parent.addr, true, false)
+                }
+                // if the parent block is a JOIN block, figure out if the new block is the first
+                // or the second child
+                BlockType::Join(first_child_executed) => {
+                    (parent.addr, false, !first_child_executed)
+                }
+                _ => (parent.addr, false, false),
+            },
+            // if the block stack is empty, a new block is neither a body of a loop nor the first
+            // child of a JOIN block; also, we set the parent address to ZERO.
+            None => (ZERO, false, false),
+        };
 
-    /// Pushes a new CALL block onto the block stack and returns the address of the block's parent.
-    ///
-    /// This is similar to pushing any other block onto the stack but unlike with all other blocks,
-    /// we initialize parent context and free memory pointer with the specified values.
-    pub fn push_call(&mut self, addr: Felt, parent_ctx: u32, parent_fmp: Felt) -> Felt {
-        let (parent_addr, is_loop_body, is_first_child) = self.get_new_block_context();
         self.blocks.push(BlockInfo {
             addr,
-            block_type: BlockType::Call,
+            block_type,
             parent_addr,
-            parent_ctx,
-            parent_fmp,
+            ctx_info,
             is_loop_body,
             is_first_child,
         });
@@ -82,37 +93,6 @@ impl BlockStack {
     pub fn peek_mut(&mut self) -> &mut BlockInfo {
         self.blocks.last_mut().expect("block stack is empty")
     }
-
-    // HELPER METHODS
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns context information for a block which is about to be pushed onto the block stack.
-    ///
-    /// Context information consists of a tuple (parent_addr, is_loop_body, is_first_child), where:
-    /// - parent_addr is the address of the block currently at the top of the stack.
-    /// - is_loop_body is set to true if the newly added block will be a body of a loop.
-    /// - is_first_child is set to true if the newly added block will be a first child in the
-    ///   JOIN block.
-    fn get_new_block_context(&self) -> (Felt, bool, bool) {
-        match self.blocks.last() {
-            Some(parent) => match parent.block_type {
-                // if the current block is a LOOP block, the new block must be a loop body
-                BlockType::Loop(loop_entered) => {
-                    debug_assert!(loop_entered, "parent is un-entered loop");
-                    (parent.addr, true, false)
-                }
-                // if the current block is a JOIN block, figure out if the new block is the first
-                // or the second child
-                BlockType::Join(first_child_executed) => {
-                    (parent.addr, false, !first_child_executed)
-                }
-                _ => (parent.addr, false, false),
-            },
-            // if the block stack is empty, a new block is neither a body of a loop nor the first
-            // child of a JOIN block; also, we set the parent address to ZERO.
-            None => (ZERO, false, false),
-        }
-    }
 }
 
 // BLOCK INFO
@@ -124,8 +104,7 @@ pub struct BlockInfo {
     pub addr: Felt,
     block_type: BlockType,
     pub parent_addr: Felt,
-    pub parent_ctx: u32,
-    pub parent_fmp: Felt,
+    pub ctx_info: Option<ExecutionContextInfo>,
     pub is_loop_body: bool,
     pub is_first_child: bool,
 }
@@ -177,6 +156,35 @@ impl BlockInfo {
     }
 }
 
+// BLOCK CONTEXT INFO
+// ================================================================================================
+
+/// Contains information about an execution context. Execution context are relevant only for CALL
+/// blocks.
+#[derive(Debug, Default, Clone, Copy)]
+pub struct ExecutionContextInfo {
+    /// Context ID of the block's parent.
+    pub parent_ctx: u32,
+    /// Value of free memory pointer right before a CALL instruction is executed.
+    pub fmp: Felt,
+    /// Depth of the operand stack right before a CALL operation is executed.
+    pub stack_depth: u32,
+    /// Address of the top row in the overflow table right before a CALL operations is executed.
+    pub next_overflow_addr: Felt,
+}
+
+impl ExecutionContextInfo {
+    /// Returns an new [ExecutionContextInfo] instantiated with the specified parameters.
+    pub fn new(parent_ctx: u32, fmp: Felt, stack_depth: u32, next_overflow_addr: Felt) -> Self {
+        Self {
+            parent_ctx,
+            fmp,
+            stack_depth,
+            next_overflow_addr,
+        }
+    }
+}
+
 // BLOCK TYPE
 // ================================================================================================