miden-debug-engine 0.19.0

use alloc::{sync::Arc, vec::Vec};

use miden_assembly_syntax::ast::DebugVarInfo;
use miden_core::{mast::MastNode, operations::AssemblyOp};
use miden_mast_package::debug_info::{DebugSourceNode, DebugSourceNodeId, PackageDebugInfo};
use miden_processor::{Continuation, ResumeContext, operation::Operation};

use crate::debug::{ControlFlowOp, InlineCallFrame, inline_frames_for_operation};

/// Source metadata for one VM cycle, selected once for every debugger consumer.
pub(crate) struct SourceContext {
    pub cycle: CurrentCycleInfo,
    debug_info: Option<Arc<PackageDebugInfo>>,
    inline_frames: Vec<InlineCallFrame>,
}

impl SourceContext {
    pub fn new(context: &ResumeContext) -> Self {
        let debug_info = context.debug_info();
        let cycle = extract_cycle(context, debug_info.as_deref());
        let inline_frames = inline_frames_for_operation(
            debug_info
                .as_deref()
                .zip(cycle.source_node_id)
                .map(|(info, source)| (info, source, cycle.op_idx.unwrap_or_default() as u32)),
            context.inherited_inline_call_contexts(),
        );
        Self {
            cycle,
            debug_info,
            inline_frames,
        }
    }

    pub fn debug_info(&self) -> Option<&Arc<PackageDebugInfo>> {
        self.debug_info.as_ref()
    }

    pub fn source_node(&self) -> Option<&DebugSourceNode> {
        let info = self.debug_info.as_deref()?;
        info.source_node(self.cycle.source_node_id?)
    }

    pub fn source_operation(&self) -> Option<(DebugSourceNodeId, u32)> {
        Some((self.cycle.source_node_id?, self.cycle.op_idx.unwrap_or_default() as u32))
    }

    /// The current source location, including the first operation on control-flow cycles.
    pub fn asmop(&self) -> Option<AssemblyOp> {
        self.source_node()?
            .asm_op_for_operation(self.cycle.op_idx.unwrap_or_default() as u32)?
            .to_assembly_op(self.debug_info.as_deref()?)
    }

    /// Operation-indexed metadata used by DAP stepping. SPAN/END cycles have no such operation;
    /// RESPAN retains its operation index, matching variable-event timing.
    #[cfg(any(feature = "dap", test))]
    pub fn executed_asmop(&self) -> Option<AssemblyOp> {
        self.cycle.op_idx?;
        self.asmop()
    }

    pub fn inline_frames(&self) -> &[InlineCallFrame] {
        &self.inline_frames
    }

    pub fn debug_vars(&self) -> Vec<DebugVarInfo> {
        let Some(((node, operation), info)) =
            self.source_node().zip(self.cycle.op_idx).zip(self.debug_info.as_deref())
        else {
            return Vec::new();
        };
        node.debug_infos_for_operation(operation as u32, info).collect()
    }
}

#[derive(Clone, Copy)]
pub(crate) struct CurrentCycleInfo {
    pub op: Option<Operation>,
    pub source_node_id: Option<DebugSourceNodeId>,
    pub op_idx: Option<usize>,
    pub control_flow_kind: Option<ControlFlowOp>,
}

/// Extract the current operation and assembly info from the continuation stack
/// before a step is executed. This lets us know what operation will run next.
pub(crate) fn extract_current_op(ctx: &ResumeContext) -> CurrentCycleInfo {
    extract_cycle(ctx, ctx.debug_info().as_deref())
}

fn extract_cycle(ctx: &ResumeContext, debug_info: Option<&PackageDebugInfo>) -> CurrentCycleInfo {
    let forest = ctx.current_forest();
    let continuation_stack = ctx.continuation_stack();
    for (cont, source_node_id) in
        continuation_stack.iter_continuations_for_next_clock_with_source_node_ids()
    {
        let exec_node = cont.exec_node();
        let source_node_id = source_node_id.or_else(|| {
            exec_node.zip(debug_info).and_then(|(exec_node, di)| {
                di.unique_source_root_for_exec_node(exec_node).ok().flatten()
            })
        });
        match cont {
            Continuation::ResumeBasicBlock {
                node_id,
                batch_index,
                op_idx_in_batch,
            } => {
                let MastNode::Block(block) = &forest[*node_id] else {
                    unreachable!()
                };
                // Compute global op index within the basic block
                let mut global_idx = 0;
                for batch in &block.op_batches()[..*batch_index] {
                    global_idx += batch.ops().len();
                }
                global_idx += op_idx_in_batch;
                let op = block.op_batches()[*batch_index].ops().get(*op_idx_in_batch).copied();
                return CurrentCycleInfo {
                    op,
                    source_node_id,
                    op_idx: Some(global_idx),
                    control_flow_kind: None,
                };
            }
            Continuation::Respan {
                node_id,
                batch_index,
            } => {
                let node = &forest[*node_id];
                if let MastNode::Block(block) = node {
                    let mut global_idx = 0;
                    for batch in &block.op_batches()[..*batch_index] {
                        global_idx += batch.ops().len();
                    }
                    return CurrentCycleInfo {
                        op: None,
                        source_node_id,
                        op_idx: Some(global_idx),
                        control_flow_kind: Some(ControlFlowOp::Respan),
                    };
                }
            }
            Continuation::StartNode(node_id) => {
                let control_flow_kind = match &forest[*node_id] {
                    MastNode::Block(_) => Some(ControlFlowOp::Span),
                    MastNode::Join(_) => Some(ControlFlowOp::Join),
                    MastNode::Split(_) => Some(ControlFlowOp::Split),
                    _ => None,
                };
                return CurrentCycleInfo {
                    op: None,
                    source_node_id,
                    op_idx: None,
                    control_flow_kind,
                };
            }
            Continuation::FinishBasicBlock(_)
            | Continuation::FinishJoin(_)
            | Continuation::FinishSplit(_)
            | Continuation::FinishLoop { .. }
            | Continuation::FinishCall(_)
            | Continuation::FinishDyn(_) => {
                return CurrentCycleInfo {
                    op: None,
                    source_node_id,
                    op_idx: None,
                    control_flow_kind: Some(ControlFlowOp::End),
                };
            }
            Continuation::EnterForest { .. } => {
                return CurrentCycleInfo {
                    op: None,
                    source_node_id,
                    op_idx: None,
                    control_flow_kind: None,
                };
            }
        }
    }
    CurrentCycleInfo {
        op: None,
        source_node_id: None,
        op_idx: None,
        control_flow_kind: None,
    }
}

#[cfg(test)]
mod tests;