miden-assembly 0.22.2

use alloc::vec::Vec;

use miden_assembly_syntax::{
    Word,
    ast::{InvocationTarget, InvokeKind},
    diagnostics::Report,
};
use miden_core::{
    mast::{CallNodeBuilder, DynNodeBuilder, MastForestContributor, MastNodeExt, MastNodeId},
    operations::{AssemblyOp, Operation},
};
use smallvec::SmallVec;

use crate::{
    Assembler, GlobalItemIndex, basic_block_builder::BasicBlockBuilder,
    mast_forest_builder::MastForestBuilder,
};

/// Procedure Invocation
impl Assembler {
    /// Returns the [`MastNodeId`] of the invoked procedure specified by `callee`.
    ///
    /// For example, given `exec.f`, this method would return the procedure body id of `f`. If the
    /// only representation of `f` that we have is its MAST root, then this method will also insert
    /// a [`core::mast::ExternalNode`] that wraps `f`'s MAST root and return the corresponding id.
    pub(super) fn invoke(
        &self,
        kind: InvokeKind,
        callee: &InvocationTarget,
        caller: GlobalItemIndex,
        mast_forest_builder: &mut MastForestBuilder,
        before_enter: Vec<miden_core::mast::DecoratorId>,
        asm_op: Option<AssemblyOp>,
    ) -> Result<MastNodeId, Report> {
        let resolved = self
            .resolve_target(kind, callee, caller, mast_forest_builder)?
            .expect("invocation target is not a procedure");

        match kind {
            InvokeKind::ProcRef | InvokeKind::Exec => Ok(resolved.node),
            InvokeKind::Call => mast_forest_builder.ensure_node_with_asm_op(
                CallNodeBuilder::new(resolved.node)
                    .with_before_enter(before_enter)
                    .with_after_exit(vec![]),
                asm_op.expect("call invocations must provide an AssemblyOp"),
            ),
            InvokeKind::SysCall => mast_forest_builder.ensure_node_with_asm_op(
                CallNodeBuilder::new_syscall(resolved.node)
                    .with_before_enter(before_enter)
                    .with_after_exit(vec![]),
                asm_op.expect("syscall invocations must provide an AssemblyOp"),
            ),
        }
    }

    /// Creates a new DYN block for the dynamic code execution and return.
    pub(super) fn dynexec(
        &self,
        mast_forest_builder: &mut MastForestBuilder,
        before_enter: Vec<miden_core::mast::DecoratorId>,
        asm_op: AssemblyOp,
    ) -> Result<Option<MastNodeId>, Report> {
        let dyn_node_id = mast_forest_builder.ensure_node_with_asm_op(
            DynNodeBuilder::new_dyn()
                .with_before_enter(before_enter)
                .with_after_exit(vec![]),
            asm_op,
        )?;

        Ok(Some(dyn_node_id))
    }

    /// Creates a new DYNCALL block for the dynamic function call and return.
    pub(super) fn dyncall(
        &self,
        mast_forest_builder: &mut MastForestBuilder,
        before_enter: Vec<miden_core::mast::DecoratorId>,
        asm_op: AssemblyOp,
    ) -> Result<Option<MastNodeId>, Report> {
        let dyn_call_node_id = mast_forest_builder.ensure_node_with_asm_op(
            DynNodeBuilder::new_dyncall()
                .with_before_enter(before_enter)
                .with_after_exit(vec![]),
            asm_op,
        )?;

        Ok(Some(dyn_call_node_id))
    }

    pub(super) fn procref(
        &self,
        callee: &InvocationTarget,
        caller: GlobalItemIndex,
        block_builder: &mut BasicBlockBuilder,
    ) -> Result<(), Report> {
        let mast_root = {
            let resolved = self
                .resolve_target(
                    InvokeKind::ProcRef,
                    callee,
                    caller,
                    block_builder.mast_forest_builder_mut(),
                )?
                .expect("invocation target is not a procedure");
            // Note: it's ok to `unwrap()` here since `proc_body_id` was returned from
            // `mast_forest_builder`
            block_builder
                .mast_forest_builder()
                .get_mast_node(resolved.node)
                .unwrap()
                .digest()
        };

        self.procref_mast_root(mast_root, block_builder)
    }

    fn procref_mast_root(
        &self,
        mast_root: Word,
        block_builder: &mut BasicBlockBuilder,
    ) -> Result<(), Report> {
        // Create an array with `Push` operations containing root elements.
        // Push in reverse order so that mast_root[0] ends up on top.
        let ops = mast_root
            .iter()
            .rev()
            .map(|elem| Operation::Push(*elem))
            .collect::<SmallVec<[_; 4]>>();
        block_builder.push_ops(ops);
        Ok(())
    }
}