miden-assembly 0.25.8

use alloc::{collections::BTreeMap, sync::Arc, vec::Vec};
use core::fmt::Debug;

use miden_core::{
    Word,
    mast::{
        BasicBlockNode, ExecutableMastForest, MastForest, MastNode, MastNodeExt, MastNodeId,
        SubtreeIterator,
    },
};
use miden_mast_package::debug_info::{
    DebugFunctionIdx, DebugInfoTableRemapping, DebugSourceAsmOp, DebugSourceInlineCall,
    DebugSourceNodeId, DebugSourceVar, PackageDebugInfo,
};

use super::{
    MastForestBuilder, MastNodeRef, PendingMastNodeDraft, PendingMastNodeKind, SourceNodeRef,
};
use crate::diagnostics::Report;

#[derive(Clone, Copy)]
struct StaticSourceRoot {
    forest_idx: usize,
    source_root_id: MastNodeId,
    source_debug_root_id: Option<DebugSourceNodeId>,
}

struct StaticLinkedRoot {
    root_id: MastNodeId,
    source: Option<StaticSourceRoot>,
}

#[derive(Default)]
pub(super) struct StaticSourceMetadata {
    op_range: Option<(usize, usize)>,
    asm_ops: Vec<DebugSourceAsmOp>,
    debug_vars: Vec<DebugSourceVar>,
    inline_calls: Vec<DebugSourceInlineCall>,
    functions: Vec<DebugFunctionIdx>,
}

struct StaticPendingDraft {
    draft: PendingMastNodeDraft,
    source_op_range: Option<(usize, usize)>,
}

/// Result of resolving an exact static-library root provenance hint.
enum StaticRootLookup {
    /// Exactly one linked source forest maps the hinted source root to the requested digest.
    Found(StaticLinkedRoot),
    /// More than one linked source forest matches the hint, so importing by provenance would risk
    /// selecting metadata from the wrong forest.
    Ambiguous,
    /// The hint did not match any linked source forest/root pair.
    Missing,
}

impl MastForestBuilder {
    /// Creates a complete [`PendingMastNodeDraft`] for a node imported from a statically
    /// linked forest, including indexed assembly ops and debug variable metadata.
    fn pending_draft_for_statically_linked_source(
        &mut self,
        source_node: MastNode,
        child_refs: Vec<MastNodeRef>,
        source_metadata: Option<StaticSourceMetadata>,
    ) -> StaticPendingDraft {
        let digest = source_node.digest();
        let kind = PendingMastNodeKind::from_node(source_node);

        let StaticSourceMetadata {
            op_range,
            asm_ops,
            debug_vars,
            inline_calls,
            functions,
        } = source_metadata.unwrap_or_default();

        StaticPendingDraft {
            draft: PendingMastNodeDraft {
                digest,
                kind,
                child_refs,
                asm_ops,
                debug_vars,
                inline_calls,
                functions,
            },
            source_op_range: op_range,
        }
    }

    /// Copies a statically linked node into this builder while keeping source metadata in the
    /// pending record when a new node is created.
    pub(super) fn ensure_node_from_statically_linked_source_ref(
        &mut self,
        source_node: MastNode,
        child_refs: Vec<MastNodeRef>,
        source_metadata: Option<StaticSourceMetadata>,
    ) -> Result<MastNodeRef, Report> {
        let StaticPendingDraft { draft, source_op_range } = self
            .pending_draft_for_statically_linked_source(source_node, child_refs, source_metadata);
        let dedup_key = self.dedup_key_for_pending_data(&draft);
        let source_child_refs = self.source_child_refs_for_node_refs(&draft.child_refs);
        if let Some(node_ref) = self.find_reusable_node_ref_by_key(&dedup_key, &draft) {
            self.record_static_source_occurrence(
                node_ref,
                source_child_refs,
                &draft,
                source_op_range,
            )?;
            return Ok(node_ref);
        }

        let node_ref = self.insert_or_replace_pending_node_record_ref(
            dedup_key,
            draft.clone(),
            &source_child_refs,
        )?;
        self.record_static_source_occurrence(node_ref, source_child_refs, &draft, source_op_range)?;
        Ok(node_ref)
    }

    fn record_static_source_occurrence(
        &mut self,
        exec_ref: MastNodeRef,
        child_refs: Vec<SourceNodeRef>,
        draft: &PendingMastNodeDraft,
        source_op_range: Option<(usize, usize)>,
    ) -> Result<SourceNodeRef, Report> {
        let (op_start, op_end) =
            source_op_range.unwrap_or_else(|| self.source_op_range_for_draft(draft));
        self.push_source_occurrence(
            exec_ref,
            child_refs,
            op_start,
            op_end,
            draft.asm_ops.clone(),
            draft.debug_vars.clone(),
            draft.inline_calls.clone(),
            &draft.functions,
            true,
        )
    }

    fn unadjust_source_block_indices<T>(
        &self,
        source_forest: &MastForest,
        source_node_id: MastNodeId,
        mut mappings: Vec<T>,
        get_element_op_idx: impl Fn(&mut T) -> &mut u32,
    ) -> Vec<T> {
        if let Some(MastNode::Block(block)) = source_forest.get_node_by_id(source_node_id) {
            let unadjusted_indices = BasicBlockNode::unadjust_asm_op_indices(
                mappings
                    .iter_mut()
                    .map(|element| *get_element_op_idx(element) as usize)
                    .collect(),
                block.op_batches(),
            );
            for (op_idx, element) in unadjusted_indices.into_iter().zip(mappings.iter_mut()) {
                *get_element_op_idx(element) = op_idx as u32;
            }
            mappings
        } else {
            mappings
        }
    }

    /// Collects builder-local refs for a statically linked source node.
    pub(super) fn pending_refs_for_statically_linked_source(
        &self,
        node: &MastNode,
        node_refs_by_source_id: &BTreeMap<MastNodeId, MastNodeRef>,
    ) -> Vec<MastNodeRef> {
        let mut child_refs = Vec::new();
        node.for_each_child(|source_child_id| {
            let child_ref = *node_refs_by_source_id
                .get(&source_child_id)
                .expect("statically linked child must be copied before its parent");
            child_refs.push(child_ref);
        });

        child_refs
    }

    /// Adds an externally-linked procedure root and returns its builder-local [`MastNodeRef`].
    pub(crate) fn ensure_external_link_with_source_ref(
        &mut self,
        mast_root: Word,
        source_library_commitment: Option<Word>,
        source_root_id: Option<MastNodeId>,
        source_debug_root_id: Option<DebugSourceNodeId>,
    ) -> Result<MastNodeRef, Report> {
        if let Some(linked_root) = self.find_statically_linked_root(
            source_library_commitment,
            source_root_id,
            source_debug_root_id,
            mast_root,
        ) {
            return self.copy_statically_linked_subtree_ref(linked_root);
        }

        self.intern_pending_node(PendingMastNodeDraft::new(
            PendingMastNodeKind::External,
            mast_root,
            Vec::new(),
        ))
    }

    fn find_statically_linked_root(
        &self,
        source_library_commitment: Option<Word>,
        source_root_id: Option<MastNodeId>,
        source_debug_root_id: Option<DebugSourceNodeId>,
        mast_root: Word,
    ) -> Option<StaticLinkedRoot> {
        if let (Some(source_library_commitment), Some(source_root_id)) =
            (source_library_commitment, source_root_id)
        {
            match self.find_exact_statically_linked_root(
                source_library_commitment,
                source_root_id,
                source_debug_root_id,
                mast_root,
            ) {
                StaticRootLookup::Found(linked_root) => return Some(linked_root),
                // `MastForest::commitment()` does not include diagnostics metadata, so multiple
                // source forests can share a commitment while still carrying different metadata.
                // In that case we drop the source hint and fall back to digest-only linking.
                StaticRootLookup::Ambiguous => {},
                StaticRootLookup::Missing => {},
            }
        }

        self.statically_linked_mast
            .find_procedure_root(mast_root)
            .map(|root_id| StaticLinkedRoot { root_id, source: None })
    }

    fn find_exact_statically_linked_root(
        &self,
        source_library_commitment: Word,
        source_root_id: MastNodeId,
        source_debug_root_id: Option<DebugSourceNodeId>,
        mast_root: Word,
    ) -> StaticRootLookup {
        let Some(forest_indices) = self
            .statically_linked_forest_indices_by_commitment
            .get(&source_library_commitment)
        else {
            return StaticRootLookup::Missing;
        };

        let mut matching_roots = forest_indices.iter().filter_map(|forest_idx| {
            self.statically_linked_root_map.map_root(*forest_idx, &source_root_id).and_then(
                |root_id| {
                    (self.statically_linked_mast[root_id].digest() == mast_root).then_some(
                        StaticLinkedRoot {
                            root_id,
                            source: Some(StaticSourceRoot {
                                forest_idx: *forest_idx,
                                source_root_id,
                                source_debug_root_id,
                            }),
                        },
                    )
                },
            )
        });

        let Some(linked_root) = matching_roots.next() else {
            return StaticRootLookup::Missing;
        };

        if matching_roots.next().is_some() {
            StaticRootLookup::Ambiguous
        } else {
            StaticRootLookup::Found(linked_root)
        }
    }

    /// Copies a subtree from the statically linked forest into the builder's forest.
    fn copy_statically_linked_subtree_ref(
        &mut self,
        linked_root: StaticLinkedRoot,
    ) -> Result<MastNodeRef, Report> {
        if let Some(source) = linked_root.source
            && let Some(package_debug_info) = self
                .statically_linked_package_debug_info
                .get(source.forest_idx)
                .cloned()
                .flatten()
            && let Some(source_forest) =
                self.statically_linked_source_forests.get(source.forest_idx)
        {
            let source_forest = Arc::clone(source_forest);
            let source_debug_root_id =
                if let Some(source_debug_root_id) = source.source_debug_root_id {
                    Some(source_debug_root_id)
                } else {
                    package_debug_info
                    .unique_source_root_for_exec_node(source.source_root_id)
                    .map_err(|err| {
                        Report::msg(format!(
                            "ambiguous statically linked source root for {source_root_id:?}: {err}",
                            source_root_id = source.source_root_id
                        ))
                    })?
                };
            if let Some(source_debug_root_id) = source_debug_root_id {
                let source_node =
                    package_debug_info.source_node(source_debug_root_id).ok_or_else(|| {
                        Report::msg(format!(
                            "statically linked package export references missing source node {source_debug_root_id:?}"
                        ))
                    })?;
                if source_node.exec_node != source.source_root_id {
                    return Err(Report::msg(format!(
                        "statically linked package export source node {source_debug_root_id:?} maps to {:?}, expected {:?}",
                        source_node.exec_node, source.source_root_id
                    )));
                }
                let tables =
                    self.import_static_debug_tables(source.forest_idx, &package_debug_info)?;
                return self.copy_package_debug_source_subtree_ref(
                    source_forest.as_ref(),
                    &package_debug_info,
                    &tables,
                    source_debug_root_id,
                );
            }
        }

        let mut node_refs_by_source_id = BTreeMap::new();
        let source_forest = Arc::clone(&self.statically_linked_mast);
        for old_id in SubtreeIterator::new(&linked_root.root_id, source_forest.as_ref()) {
            let node = self.statically_linked_mast[old_id].clone();
            let child_refs =
                self.pending_refs_for_statically_linked_source(&node, &node_refs_by_source_id);
            let new_ref =
                self.ensure_node_from_statically_linked_source_ref(node, child_refs, None)?;
            node_refs_by_source_id.insert(old_id, new_ref);
        }
        Ok(*node_refs_by_source_id
            .get(&linked_root.root_id)
            .expect("statically linked subtree root must be copied"))
    }

    fn import_static_debug_tables(
        &mut self,
        forest_idx: usize,
        package_debug_info: &PackageDebugInfo,
    ) -> Result<Arc<DebugInfoTableRemapping>, Report> {
        let slot = self
            .statically_linked_debug_table_remappings
            .get_mut(forest_idx)
            .expect("static debug table mappings must remain parallel to static libraries");
        if let Some(tables) = slot.as_ref() {
            return Ok(Arc::clone(tables));
        }

        let tables = Arc::new(package_debug_info.merge_tables_into(&mut self.debug_info).map_err(
            |error| {
                Report::msg(format!("failed to import statically linked debug tables: {error}"))
            },
        )?);
        *slot = Some(Arc::clone(&tables));
        Ok(tables)
    }

    fn copy_package_debug_source_subtree_ref(
        &mut self,
        source_forest: &MastForest,
        package_debug_info: &PackageDebugInfo,
        tables: &DebugInfoTableRemapping,
        source_root_id: DebugSourceNodeId,
    ) -> Result<MastNodeRef, Report> {
        let mut node_refs_by_source_id = BTreeMap::new();
        self.copy_package_debug_source_node_ref(
            source_forest,
            package_debug_info,
            tables,
            source_root_id,
            &mut node_refs_by_source_id,
        )
    }

    fn copy_package_debug_source_node_ref(
        &mut self,
        source_forest: &MastForest,
        package_debug_info: &PackageDebugInfo,
        tables: &DebugInfoTableRemapping,
        source_node_id: DebugSourceNodeId,
        node_refs_by_source_id: &mut BTreeMap<DebugSourceNodeId, MastNodeRef>,
    ) -> Result<MastNodeRef, Report> {
        if let Some(node_ref) = node_refs_by_source_id.get(&source_node_id).copied() {
            return Ok(node_ref);
        }

        let source_node = package_debug_info.source_node(source_node_id).ok_or_else(|| {
            Report::msg(format!(
                "statically linked package debug graph is missing source node {source_node_id:?}"
            ))
        })?;

        let source_exec_node_id = source_node.exec_node;
        let source_exec_node = source_forest
            .get_node_by_id(source_exec_node_id)
            .ok_or_else(|| {
                Report::msg(format!(
                    "statically linked package debug graph references missing execution node {source_exec_node_id:?}"
                ))
            })?
            .clone();
        let mut exec_child_ids = Vec::new();
        source_exec_node.for_each_child(|child_id| exec_child_ids.push(child_id));
        if exec_child_ids.len() != source_node.children.len() {
            return Err(Report::msg(format!(
                "statically linked package debug source node {source_node_id:?} has {} children, expected {} from execution node {source_exec_node_id:?}",
                source_node.children.len(),
                exec_child_ids.len(),
            )));
        }

        let mut child_refs = Vec::new();
        for (child_index, child_source_node_id) in source_node.children.iter().copied().enumerate()
        {
            let child_source_node =
                package_debug_info.source_node(child_source_node_id).ok_or_else(|| {
                    Report::msg(format!(
                        "statically linked package debug graph source node {source_node_id:?} references missing child source node {child_source_node_id:?}"
                    ))
                })?;
            if child_source_node.exec_node != exec_child_ids[child_index] {
                return Err(Report::msg(format!(
                    "statically linked package debug graph source node {source_node_id:?} child {child_index} maps to {:?}, expected {:?}",
                    child_source_node.exec_node, exec_child_ids[child_index],
                )));
            }
            child_refs.push(self.copy_package_debug_source_node_ref(
                source_forest,
                package_debug_info,
                tables,
                child_source_node_id,
                node_refs_by_source_id,
            )?);
        }

        let metadata = self.package_source_metadata(
            source_forest,
            package_debug_info,
            tables,
            source_node_id,
            source_exec_node_id,
        )?;
        let node_ref = self.ensure_node_from_statically_linked_source_ref(
            source_exec_node,
            child_refs,
            Some(metadata),
        )?;
        node_refs_by_source_id.insert(source_node_id, node_ref);
        Ok(node_ref)
    }

    fn package_source_metadata(
        &mut self,
        source_forest: &MastForest,
        package_debug_info: &PackageDebugInfo,
        tables: &DebugInfoTableRemapping,
        source_node_id: DebugSourceNodeId,
        source_exec_node_id: MastNodeId,
    ) -> Result<StaticSourceMetadata, Report> {
        let asm_ops = package_debug_info
            .asm_ops_for_source_node(source_node_id)
            .map(|row| {
                let location_idx = row
                    .location_idx
                    .into_option()
                    .map(|index| remapped_debug_index(tables.location(index), index, "location"))
                    .transpose()?;
                let context_name_idx = remapped_debug_index(
                    tables.string(row.context_name_idx),
                    row.context_name_idx,
                    "string",
                )?;
                let op_name_idx = remapped_debug_index(
                    tables.string(row.op_name_idx),
                    row.op_name_idx,
                    "string",
                )?;
                Ok(DebugSourceAsmOp::new(
                    row.op_idx,
                    location_idx,
                    context_name_idx,
                    op_name_idx,
                    row.num_cycles,
                ))
            })
            .collect::<Result<Vec<_>, Report>>()?;
        let debug_vars = package_debug_info
            .debug_vars_for_source_node(source_node_id)
            .map(|row| {
                let name_idx =
                    remapped_debug_index(tables.string(row.name_idx), row.name_idx, "string")?;
                let location_idx = row
                    .location_idx
                    .map(|index| remapped_debug_index(tables.location(index), index, "location"))
                    .transpose()?;
                let type_id = row
                    .type_id
                    .map(|index| remapped_debug_index(tables.ty(index), index, "type"))
                    .transpose()?;
                Ok(DebugSourceVar {
                    op_idx: row.op_idx,
                    name_idx,
                    type_id,
                    arg_idx: row.arg_idx,
                    location_idx,
                    value_location: row.value_location.clone(),
                })
            })
            .collect::<Result<Vec<_>, Report>>()?;
        let inline_calls = package_debug_info
            .inline_calls_for_source_node(source_node_id)
            .map(|row| {
                let loc_idx =
                    remapped_debug_index(tables.location(row.loc_idx), row.loc_idx, "location")?;
                let callee_idx = remapped_debug_index(
                    tables.function(row.callee_idx),
                    row.callee_idx,
                    "function",
                )?;
                Ok(DebugSourceInlineCall { op_idx: row.op_idx, callee_idx, loc_idx })
            })
            .collect::<Result<Vec<_>, Report>>()?;
        let op_range = package_debug_info.source_node(source_node_id).map(|source_node| {
            self.unadjust_source_block_range(
                source_forest,
                source_exec_node_id,
                source_node.op_start as usize,
                source_node.op_end as usize,
            )
        });

        let mast_root = source_forest.get_digest_by_id(source_exec_node_id).unwrap();
        let functions = package_debug_info
            .functions()
            .iter()
            .enumerate()
            .filter_map(|(index, function)| {
                let function_source_node = function.source_node.into_option();
                if function_source_node.is_some_and(|snid| snid == source_node_id)
                    || function_source_node.is_none() && function.mast_root == mast_root
                {
                    let function_index = DebugFunctionIdx::from(
                        u32::try_from(index).expect("invalid function index"),
                    );
                    tables.function(function_index)
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        Ok(StaticSourceMetadata {
            op_range,
            asm_ops: self.unadjust_source_block_indices(
                source_forest,
                source_exec_node_id,
                asm_ops,
                |asm_op| &mut asm_op.op_idx,
            ),
            debug_vars: self.unadjust_source_block_indices(
                source_forest,
                source_exec_node_id,
                debug_vars,
                |debug_var| &mut debug_var.op_idx,
            ),
            inline_calls: self.unadjust_source_block_indices(
                source_forest,
                source_exec_node_id,
                inline_calls,
                |inline_call| &mut inline_call.op_idx,
            ),
            functions,
        })
    }

    fn unadjust_source_block_range(
        &self,
        source_forest: &MastForest,
        source_node_id: MastNodeId,
        op_start: usize,
        op_end: usize,
    ) -> (usize, usize) {
        if op_start == op_end {
            return (op_start, op_end);
        }

        if let Some(MastNode::Block(block)) = source_forest.get_node_by_id(source_node_id) {
            let unadjusted_indices = BasicBlockNode::unadjust_asm_op_indices(
                vec![op_start, op_end - 1],
                block.op_batches(),
            );
            (unadjusted_indices[0], unadjusted_indices[1] + 1)
        } else {
            (op_start, op_end)
        }
    }
}

fn remapped_debug_index<I: Debug, T>(
    mapped: Option<T>,
    index: I,
    table: &str,
) -> Result<T, Report> {
    mapped.ok_or_else(|| {
        Report::msg(format!(
            "statically linked debug info references missing {table} index {index:?}"
        ))
    })
}