miden_processor/
continuation_stack.rs

1use alloc::{sync::Arc, vec::Vec};
2
3use miden_core::{mast::MastNodeId, program::Program};
4use miden_mast_package::debug_info::{DebugSourceNodeId, PackageDebugInfo};
5
6/// A hint for the initial size of the continuation stack.
7const CONTINUATION_STACK_SIZE_HINT: usize = 64;
8
9// CONTINUATION
10// ================================================================================================
11
12/// Represents a unit of work in the continuation stack.
13///
14/// This enum defines the different types of continuations that can be performed on MAST nodes
15/// during program execution.
16///
17/// The type parameter `F` is the representation of a MAST forest carried by the
18/// [`Continuation::EnterForest`] variant. For live execution this is `Arc<MastForest>`; for the
19/// snapshotted continuation stack inside a trace fragment it is a `usize` index into the
20/// `mast_forest_store` of the trace generation context.
21#[derive(Debug, Clone)]
22pub enum Continuation<F> {
23    /// Start processing a node in the MAST forest.
24    StartNode(MastNodeId),
25    /// Process the finish phase of a Join node.
26    FinishJoin(MastNodeId),
27    /// Process the finish phase of a Split node.
28    FinishSplit(MastNodeId),
29    /// Process the finish phase of a Loop node.
30    ///
31    /// Reached after the loop body has finished executing. Inspects the condition the body left on
32    /// top of the stack and either fires REPEAT (re-enter the body) or END (exit the loop). Loop
33    /// bodies are entered unconditionally — a `while.true` source construct is desugared into a
34    /// SPLIT that wraps the LOOP, so the LOOP itself sees a do-while body.
35    FinishLoop(MastNodeId),
36    /// Process the finish phase of a Call node.
37    FinishCall(MastNodeId),
38    /// Process the finish phase of a Dyn node.
39    FinishDyn(MastNodeId),
40    /// Resume execution at the specified operation of the specified batch in the given basic block
41    /// node.
42    ResumeBasicBlock {
43        node_id: MastNodeId,
44        batch_index: usize,
45        op_idx_in_batch: usize,
46    },
47    /// Resume execution at the RESPAN operation before the specific batch within a basic block
48    /// node.
49    Respan { node_id: MastNodeId, batch_index: usize },
50    /// Process the finish phase of a basic block node.
51    ///
52    /// This corresponds to incrementing the clock to account for the inserted END operation.
53    FinishBasicBlock(MastNodeId),
54    /// Enter a new MAST forest, where all subsequent `MastNodeId`s will be relative to this forest.
55    ///
56    /// When we encounter an `ExternalNode`, we enter the corresponding MAST forest directly, and
57    /// push an `EnterForest` continuation to restore the previous forest when done.
58    EnterForest {
59        forest: F,
60        package_debug_info: Option<Arc<PackageDebugInfo>>,
61    },
62}
63
64impl<F> Continuation<F> {
65    /// Returns true if executing this continuation increments the processor clock, and false
66    /// otherwise.
67    pub fn increments_clk(&self) -> bool {
68        use Continuation::*;
69
70        // Note: we prefer naming all the variants over using a wildcard arm to ensure that if new
71        // variants are added in the future, we consciously decide whether they should increment the
72        // clock or not.
73        match self {
74            StartNode(_)
75            | FinishJoin(_)
76            | FinishSplit(_)
77            | FinishLoop(_)
78            | FinishCall(_)
79            | FinishDyn(_)
80            | ResumeBasicBlock {
81                node_id: _,
82                batch_index: _,
83                op_idx_in_batch: _,
84            }
85            | Respan { node_id: _, batch_index: _ }
86            | FinishBasicBlock(_) => true,
87
88            EnterForest { .. } => false,
89        }
90    }
91
92    pub fn exec_node(&self) -> Option<MastNodeId> {
93        match self {
94            Self::StartNode(node_id)
95            | Self::FinishJoin(node_id)
96            | Self::FinishSplit(node_id)
97            | Self::FinishLoop(node_id)
98            | Self::FinishCall(node_id)
99            | Self::FinishDyn(node_id)
100            | Self::ResumeBasicBlock { node_id, .. }
101            | Self::Respan { node_id, .. }
102            | Self::FinishBasicBlock(node_id) => Some(*node_id),
103            Self::EnterForest { .. } => None,
104        }
105    }
106}
107
108// CONTINUATION STACK
109// ================================================================================================
110
111/// [ContinuationStack] reifies the call stack used by the processor when executing a program made
112/// up of possibly multiple MAST forests.
113///
114/// This allows the processor to execute a program iteratively in a loop rather than recursively
115/// traversing the nodes. It also allows the processor to pass the state of execution to another
116/// processor for further processing, which is useful for parallel execution of MAST forests.
117#[derive(Debug, Clone)]
118pub struct ContinuationStack<F> {
119    stack: Vec<Continuation<F>>,
120    source_node_ids: Option<Vec<Option<DebugSourceNodeId>>>,
121}
122
123impl<F> Default for ContinuationStack<F> {
124    fn default() -> Self {
125        Self { stack: Vec::new(), source_node_ids: None }
126    }
127}
128
129impl<F> ContinuationStack<F> {
130    /// Creates a new continuation stack for a program.
131    ///
132    /// # Arguments
133    /// * `program` - The program whose execution will be managed by this continuation stack
134    pub fn new(program: &Program) -> Self {
135        let mut stack = Vec::with_capacity(CONTINUATION_STACK_SIZE_HINT);
136        stack.push(Continuation::StartNode(program.entrypoint()));
137
138        Self { stack, source_node_ids: None }
139    }
140
141    pub(crate) fn new_with_source_node_id(
142        program: &Program,
143        source_node_id: DebugSourceNodeId,
144    ) -> Self {
145        Self::new_with_optional_source_node_id(program, Some(source_node_id))
146    }
147
148    pub(crate) fn new_with_optional_source_node_id(
149        program: &Program,
150        source_node_id: Option<DebugSourceNodeId>,
151    ) -> Self {
152        let mut stack = Vec::with_capacity(CONTINUATION_STACK_SIZE_HINT);
153        stack.push(Continuation::StartNode(program.entrypoint()));
154
155        let mut source_node_ids = Vec::with_capacity(CONTINUATION_STACK_SIZE_HINT);
156        source_node_ids.push(source_node_id);
157
158        Self {
159            stack,
160            source_node_ids: Some(source_node_ids),
161        }
162    }
163
164    // STATE MUTATORS
165    // --------------------------------------------------------------------------------------------
166
167    /// Pushes a continuation onto the continuation stack.
168    pub fn push_continuation(&mut self, continuation: Continuation<F>) {
169        self.stack.push(continuation);
170        self.push_source_node_id(None);
171    }
172
173    pub(crate) fn push_with_source_node_id(
174        &mut self,
175        continuation: Continuation<F>,
176        source_node_id: Option<DebugSourceNodeId>,
177    ) {
178        self.stack.push(continuation);
179        self.push_source_node_id(source_node_id);
180    }
181
182    /// Pushes a continuation to enter the given MAST forest on the continuation stack.
183    ///
184    /// # Arguments
185    /// * `forest` - The MAST forest to enter
186    pub fn push_enter_forest(&mut self, forest: F) {
187        self.push_enter_forest_with_package_debug_info(forest, None);
188    }
189
190    pub(crate) fn push_enter_forest_with_package_debug_info(
191        &mut self,
192        forest: F,
193        package_debug_info: Option<Arc<PackageDebugInfo>>,
194    ) {
195        self.stack.push(Continuation::EnterForest { forest, package_debug_info });
196        self.push_source_node_id(None);
197    }
198
199    /// Pushes a join finish continuation onto the stack.
200    pub fn push_finish_join(&mut self, node_id: MastNodeId) {
201        self.stack.push(Continuation::FinishJoin(node_id));
202        self.push_source_node_id(None);
203    }
204
205    /// Pushes a split finish continuation onto the stack.
206    pub fn push_finish_split(&mut self, node_id: MastNodeId) {
207        self.stack.push(Continuation::FinishSplit(node_id));
208        self.push_source_node_id(None);
209    }
210
211    /// Pushes a loop finish continuation onto the stack.
212    pub fn push_finish_loop(&mut self, node_id: MastNodeId) {
213        self.stack.push(Continuation::FinishLoop(node_id));
214        self.push_source_node_id(None);
215    }
216
217    /// Pushes a call finish continuation onto the stack.
218    pub fn push_finish_call(&mut self, node_id: MastNodeId) {
219        self.stack.push(Continuation::FinishCall(node_id));
220        self.push_source_node_id(None);
221    }
222
223    /// Pushes a dyn finish continuation onto the stack.
224    pub fn push_finish_dyn(&mut self, node_id: MastNodeId) {
225        self.stack.push(Continuation::FinishDyn(node_id));
226        self.push_source_node_id(None);
227    }
228
229    /// Pushes a continuation to start processing the given node.
230    ///
231    /// # Arguments
232    /// * `node_id` - The ID of the node to process
233    pub fn push_start_node(&mut self, node_id: MastNodeId) {
234        self.stack.push(Continuation::StartNode(node_id));
235        self.push_source_node_id(None);
236    }
237
238    /// Pops the next continuation from the continuation stack, and returns it along with its
239    /// associated MAST forest.
240    pub fn pop_continuation(&mut self) -> Option<Continuation<F>> {
241        let continuation = self.stack.pop()?;
242        if let Some(source_node_ids) = &mut self.source_node_ids {
243            source_node_ids.pop();
244        }
245        Some(continuation)
246    }
247
248    pub(crate) fn pop_continuation_with_source_node_id(
249        &mut self,
250    ) -> Option<(Continuation<F>, Option<DebugSourceNodeId>)> {
251        let continuation = self.stack.pop()?;
252        let source_node_id = self.source_node_ids.as_mut().and_then(Vec::pop).flatten();
253        Some((continuation, source_node_id))
254    }
255
256    /// Consumes this stack and returns its continuations in bottom-to-top order (i.e. the order in
257    /// which they were originally pushed).
258    pub fn into_inner(self) -> Vec<Continuation<F>> {
259        self.stack
260    }
261
262    fn push_source_node_id(&mut self, source_node_id: Option<DebugSourceNodeId>) {
263        if let Some(source_node_ids) = &mut self.source_node_ids {
264            source_node_ids.push(source_node_id);
265        }
266    }
267
268    pub(crate) fn start_tracking_source_nodes(
269        &mut self,
270        next_source_node_id: Option<DebugSourceNodeId>,
271    ) {
272        let mut source_node_ids = Vec::with_capacity(self.stack.len());
273        source_node_ids.resize(self.stack.len(), None);
274        if let Some(source_node_id) = source_node_ids.last_mut() {
275            *source_node_id = next_source_node_id;
276        }
277        self.source_node_ids = Some(source_node_ids);
278    }
279
280    // PUBLIC ACCESSORS
281    // --------------------------------------------------------------------------------------------
282
283    /// Returns the number of continuations on the stack.
284    pub fn len(&self) -> usize {
285        self.stack.len()
286    }
287
288    /// Peeks at the next continuation to execute without removing it.
289    ///
290    /// Note that more than one continuation may execute in the same clock cycle. To get all
291    /// continuations that will execute in the next clock cycle, use
292    /// [`Self::iter_continuations_for_next_clock`].
293    pub fn peek_continuation(&self) -> Option<&Continuation<F>> {
294        self.stack.last()
295    }
296
297    pub(crate) fn peek_continuation_with_source_node_id(
298        &self,
299    ) -> Option<(&Continuation<F>, Option<DebugSourceNodeId>)> {
300        let continuation = self.stack.last()?;
301        let source_node_id = self
302            .source_node_ids
303            .as_ref()
304            .and_then(|source_node_ids| source_node_ids.last().copied().flatten());
305        Some((continuation, source_node_id))
306    }
307
308    pub(crate) fn tracks_source_nodes(&self) -> bool {
309        self.source_node_ids.is_some()
310    }
311
312    /// Returns an iterator over the continuations on the stack that will execute in the next clock
313    /// cycle.
314    ///
315    /// This includes all coming continuations up to and including the first continuation that
316    /// increments the clock.
317    ///
318    /// Note: for this iterator to function correctly, it must be the case that executing a
319    /// continuation that doesn't increment the clock *does not* push new continuations on the
320    /// stack. This is currently the case, and is a reasonable invariant to maintain, as
321    /// continuations that don't increment the clock can be expected to be simple (e.g. enter a new
322    /// mast forest).
323    pub fn iter_continuations_for_next_clock(&self) -> impl Iterator<Item = &Continuation<F>> {
324        let mut found_incrementing_cont = false;
325
326        self.stack.iter().rev().take_while(move |continuation| {
327            if found_incrementing_cont {
328                // We have already found the first incrementing continuation, stop here.
329                false
330            } else if continuation.increments_clk() {
331                // This is the first incrementing continuation we have found.
332                found_incrementing_cont = true;
333                true
334            } else {
335                // This continuation does not increment the clock, continue.
336                true
337            }
338        })
339    }
340
341    /// Same as [`Self::iter_continuations_for_next_clock`], but provides the set of source node
342    /// ids for each continuation.
343    pub fn iter_continuations_for_next_clock_with_source_node_ids(
344        &self,
345    ) -> impl Iterator<Item = (&Continuation<F>, Option<DebugSourceNodeId>)> {
346        let mut stack_index = self.stack.len().saturating_sub(1);
347
348        self.iter_continuations_for_next_clock().map(move |cont| {
349            let source_node_id = self
350                .source_node_ids
351                .as_deref()
352                .and_then(|ids| ids.get(stack_index).copied())
353                .flatten();
354            stack_index = stack_index.saturating_sub(1);
355            (cont, source_node_id)
356        })
357    }
358}
359
360// TESTS
361// ================================================================================================
362
363#[cfg(test)]
364mod tests {
365    use alloc::sync::Arc;
366
367    use miden_core::mast::MastForest;
368
369    use super::*;
370
371    #[test]
372    fn get_next_clock_cycle_increment_empty_stack() {
373        let stack: ContinuationStack<Arc<MastForest>> = ContinuationStack::default();
374        assert!(stack.iter_continuations_for_next_clock().next().is_none());
375    }
376
377    #[test]
378    fn get_next_clock_cycle_increment_ends_with_incrementing() {
379        let mut stack: ContinuationStack<Arc<MastForest>> = ContinuationStack::default();
380        // Push a continuation that increments the clock
381        stack.push_continuation(Continuation::StartNode(MastNodeId::new_unchecked(0)));
382
383        let result: Vec<_> = stack.iter_continuations_for_next_clock().collect();
384        assert_eq!(result.len(), 1);
385        assert!(matches!(result[0], Continuation::StartNode(_)));
386    }
387
388    #[test]
389    fn get_next_clock_cycle_increment_enter_forest_after_incrementing() {
390        let mut stack: ContinuationStack<Arc<MastForest>> = ContinuationStack::default();
391        // Push an incrementing continuation first (bottom of stack)
392        stack.push_continuation(Continuation::StartNode(MastNodeId::new_unchecked(0)));
393        // Push a non-incrementing continuation on top
394        stack.push_continuation(Continuation::EnterForest {
395            forest: Arc::new(MastForest::new()),
396            package_debug_info: None,
397        });
398
399        let result: Vec<_> = stack.iter_continuations_for_next_clock().collect();
400        // Should return: EnterForest (non-incrementing), then StartNode (first incrementing)
401        assert_eq!(result.len(), 2);
402        assert!(matches!(result[0], Continuation::EnterForest { .. }));
403        assert!(matches!(result[1], Continuation::StartNode(_)));
404    }
405
406    #[test]
407    fn get_next_clock_cycle_increment_multiple_enter_forest_after_incrementing() {
408        let mut stack: ContinuationStack<Arc<MastForest>> = ContinuationStack::default();
409        // Push an incrementing continuation first (bottom of stack)
410        stack.push_continuation(Continuation::StartNode(MastNodeId::new_unchecked(0)));
411        // Push two non-incrementing continuations on top
412        stack.push_continuation(Continuation::EnterForest {
413            forest: Arc::new(MastForest::new()),
414            package_debug_info: None,
415        });
416        stack.push_continuation(Continuation::EnterForest {
417            forest: Arc::new(MastForest::new()),
418            package_debug_info: None,
419        });
420
421        let result: Vec<_> = stack.iter_continuations_for_next_clock().collect();
422        // Should return: EnterForest, EnterForest, StartNode
423        assert_eq!(result.len(), 3);
424        assert!(matches!(result[0], Continuation::EnterForest { .. }));
425        assert!(matches!(result[1], Continuation::EnterForest { .. }));
426        assert!(matches!(result[2], Continuation::StartNode(_)));
427    }
428}
miden_processor/continuation_stack.rs

miden_processor/
continuation_stack.rs
