miden-assembly 0.24.2

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

#[cfg(any(test, feature = "testing"))]
pub use miden_assembly_syntax::parser;
use miden_assembly_syntax::{
    Parse, Path, Word,
    ast::{Module, ModuleKind},
    debuginfo::{DefaultSourceManager, SourceManager},
    diagnostics::{
        Report,
        reporting::{ReportHandlerOpts, set_hook},
    },
};
pub use miden_assembly_syntax::{
    assert_diagnostic, assert_diagnostic_lines, parse_module, regex, source_file, testing::Pattern,
};
#[cfg(feature = "testing")]
use miden_assembly_syntax::{ast::Form, debuginfo::SourceFile};
use miden_core::program::Program;
use miden_mast_package::PackageId;
#[cfg(feature = "std")]
use miden_project::TargetType;

use crate::assembler::Assembler;
#[cfg(feature = "std")]
use crate::diagnostics::reporting::set_panic_hook;

/// A [TestContext] provides common functionality for all tests which interact with an [Assembler].
///
/// It is used by constructing it with `TestContext::default()`, which will initialize the
/// diagnostic reporting infrastructure, and construct a default [Assembler] instance for you. You
/// can then optionally customize the context, or start invoking any of its test helpers.
///
/// Some of the assertion macros defined above require a [TestContext], so be aware of that.
pub struct TestContext {
    source_manager: Arc<dyn SourceManager>,
    assembler: Assembler,
    #[cfg(feature = "std")]
    registry: TestRegistry,
}

impl Default for TestContext {
    fn default() -> Self {
        Self::new()
    }
}

impl TestContext {
    pub fn new() -> Self {
        #[cfg(feature = "logging")]
        {
            // Enable debug tracing to stderr via the MIDEN_LOG environment variable, if present
            let _ = env_logger::Builder::from_env("MIDEN_LOG").format_timestamp(None).try_init();
        }

        #[cfg(feature = "std")]
        {
            let result = set_hook(Box::new(|_| Box::new(ReportHandlerOpts::new().build())));
            #[cfg(feature = "std")]
            if result.is_ok() {
                set_panic_hook();
            }
        }

        #[cfg(not(feature = "std"))]
        {
            let _ = set_hook(Box::new(|_| Box::new(ReportHandlerOpts::new().build())));
        }
        let source_manager = Arc::new(DefaultSourceManager::default());
        let assembler = Assembler::new(source_manager.clone()).with_warnings_as_errors(true);
        #[cfg(feature = "std")]
        {
            Self {
                source_manager,
                assembler,
                registry: Default::default(),
            }
        }
        #[cfg(not(feature = "std"))]
        {
            Self { source_manager, assembler }
        }
    }

    #[cfg(feature = "std")]
    pub fn with_warnings_as_errors(self, yes: bool) -> Self {
        let Self { source_manager, assembler, registry } = self;
        Self {
            source_manager,
            assembler: assembler.with_warnings_as_errors(yes),
            registry,
        }
    }

    #[cfg(not(feature = "std"))]
    pub fn with_warnings_as_errors(self, yes: bool) -> Self {
        let Self { source_manager, assembler } = self;
        Self {
            source_manager,
            assembler: assembler.with_warnings_as_errors(yes),
        }
    }

    #[inline]
    fn assembler(&self) -> Assembler {
        self.assembler.clone()
    }

    #[inline(always)]
    pub fn source_manager(&self) -> Arc<dyn SourceManager> {
        self.source_manager.clone()
    }

    /// Parse the given source file into a vector of top-level [Form]s.
    ///
    /// This does not run semantic analysis, or construct a [Module] from the parsed
    /// forms, and is largely intended for low-level testing of the parser.
    #[cfg(feature = "testing")]
    #[track_caller]
    pub fn parse_forms(&self, source: Arc<SourceFile>) -> Result<Vec<Form>, Report> {
        parser::parse_forms(source)
    }

    /// Parse the given source file into an executable [Module].
    ///
    /// This runs semantic analysis, and the returned module is guaranteed to be syntactically
    /// valid.
    #[track_caller]
    pub fn parse_program(&self, source: Arc<SourceFile>) -> Result<Box<Module>, Report> {
        source.parse(self.assembler.warnings_as_errors(), self.source_manager())
    }

    /// Parse the given source file into a kernel [Module].
    ///
    /// This runs semantic analysis, and the returned module is guaranteed to be syntactically
    /// valid.
    #[track_caller]
    pub fn parse_kernel(&self, source: Arc<SourceFile>) -> Result<Box<Module>, Report> {
        let mut parser = Module::parser(Some(ModuleKind::Kernel));
        parser.set_warnings_as_errors(self.assembler.warnings_as_errors());
        parser.parse(Some(Path::KERNEL), source, self.source_manager())
    }

    /// Parse the given source file into an anonymous library [Module].
    ///
    /// This runs semantic analysis, and the returned module is guaranteed to be syntactically
    /// valid.
    #[track_caller]
    pub fn parse_module(&self, source: impl Parse) -> Result<Box<Module>, Report> {
        source.parse(self.assembler.warnings_as_errors(), self.source_manager())
    }

    /// Add `module` to the [Assembler] constructed by this context, making it available to
    /// other modules.
    #[track_caller]
    pub fn add_module(&mut self, module: impl Parse) -> Result<(), Report> {
        self.assembler.compile_and_statically_link(module).map(|_| ())
    }

    /// Add the modules of `library` to the [Assembler] constructed by this context.
    #[track_caller]
    pub fn add_library(&mut self, package: Arc<miden_mast_package::Package>) -> Result<(), Report> {
        self.assembler.link_package(package, miden_project::Linkage::Dynamic)
    }

    /// Compile a [Program] from `source` using the [Assembler] constructed by this context.
    ///
    /// NOTE: Any modules added by, e.g. `add_module`, will be available to the executable
    /// module represented in `source`.
    #[track_caller]
    pub fn assemble(&self, source: impl Parse) -> Result<Program, Report> {
        Ok(self.assembler().assemble_program("test", source)?.unwrap_program())
    }

    /// Compile a library package from `modules` using the [Assembler] constructed by this context.
    ///
    /// NOTE: Any modules added by, e.g. `add_module`, will be available to the library
    #[track_caller]
    pub fn assemble_library(
        &self,
        name: impl Into<PackageId>,
        version: Option<&str>,
        root: Box<Module>,
        support: impl IntoIterator<Item = Box<Module>>,
    ) -> Result<Box<miden_mast_package::Package>, Report> {
        let version = version.unwrap_or("0.0.0").parse().unwrap();
        let mut package = self.assembler().assemble_library(name, root, support)?;
        package.version = version;
        Ok(package)
    }

    /// Compile a module from `source`, with the fully-qualified name `path`, to MAST, returning
    /// the MAST roots of all the exported procedures of that module.
    #[track_caller]
    pub fn assemble_module(
        &self,
        _path: impl AsRef<Path>,
        _module: impl Parse,
    ) -> Result<Vec<Word>, Report> {
        // This API will change after we implement `Assembler::add_library()`
        unimplemented!()
    }
}

#[cfg(feature = "std")]
pub use self::package_features::TestRegistry;

#[cfg(feature = "std")]
mod package_features {
    use std::{
        collections::BTreeMap,
        string::String,
        sync::{Arc, Mutex},
    };

    use miden_mast_package::{Package, PackageId};
    use miden_package_registry::{
        PackageCache, PackageIndex, PackageProvider, PackageRecord, PackageRegistry, PackageStore,
        PackageVersions, Version, VersionRequirement,
    };

    use super::*;
    use crate::ProjectTargetSelector;

    /// A simple in-memory package index/registry
    #[derive(Default)]
    pub struct TestRegistry {
        index: BTreeMap<PackageId, PackageVersions>,
        packages: BTreeMap<(PackageId, Version), Arc<Package>>,
        loads: Mutex<Vec<String>>,
        caches: Mutex<Vec<String>>,
    }

    impl TestRegistry {
        pub fn add_package(&mut self, package: Arc<Package>) -> Version {
            let version = Version::new(package.version.clone(), package.digest());
            self.publish_package(package).expect("failed to add test package");
            version
        }

        pub fn loaded_packages(&self) -> Vec<String> {
            self.loads.lock().unwrap().clone()
        }

        pub fn cached_packages(&self) -> Vec<String> {
            self.caches.lock().unwrap().clone()
        }

        pub fn clear_loaded_packages(&self) {
            self.loads.lock().unwrap().clear();
        }

        pub fn remove_package(
            &mut self,
            package: &PackageId,
            version: &Version,
        ) -> Option<Arc<Package>> {
            self.packages.remove(&(package.clone(), version.clone()))
        }

        pub fn replace_semver_package(&mut self, package: Arc<Package>) -> Version {
            let version = Version::new(package.version.clone(), package.digest());
            self.packages.retain(|(name, existing), _| {
                name != &package.name || existing.version != package.version
            });

            let dependencies = package.manifest.dependencies().map(|dependency| {
                let version = Version::new(dependency.version.clone(), dependency.digest);
                (dependency.name.clone(), VersionRequirement::Exact(version))
            });
            let record = PackageRecord::new(version.clone(), dependencies)
                .with_description(package.description.clone().unwrap_or_default());
            self.index
                .entry(package.name.clone())
                .or_default()
                .insert(package.version.clone(), record);
            self.packages.insert((package.name.clone(), version.clone()), package);
            version
        }
    }

    impl PackageRegistry for TestRegistry {
        fn available_versions(&self, package: &PackageId) -> Option<&PackageVersions> {
            self.index.get(package)
        }
    }

    impl PackageIndex for TestRegistry {
        type Error = Report;

        fn register(&mut self, name: PackageId, record: PackageRecord) -> Result<(), Self::Error> {
            use std::collections::btree_map::Entry;

            let semver = record.semantic_version().clone();
            match self.index.entry(name.clone()).or_default().entry(semver.clone()) {
                Entry::Vacant(entry) => {
                    entry.insert(record);
                    Ok(())
                },
                Entry::Occupied(_) => Err(Report::msg(format!(
                    "package '{name}' version '{semver}' is already registered"
                ))),
            }
        }
    }

    impl PackageProvider for TestRegistry {
        fn load_package(
            &self,
            package: &PackageId,
            version: &Version,
        ) -> Result<Arc<Package>, Report> {
            self.loads.lock().unwrap().push(format!("{package}@{version}"));
            self.packages.get(&(package.clone(), version.clone())).cloned().ok_or_else(|| {
                Report::msg(format!("missing test package '{package}' at '{version}'"))
            })
        }
    }

    impl PackageCache for TestRegistry {
        type Error = Report;

        fn cache_package(&mut self, package: Arc<Package>) -> Result<Version, Self::Error> {
            let version = Version::new(package.version.clone(), package.digest());
            self.caches.lock().unwrap().push(format!("{}@{version}", package.name));
            if let Some(record) = self.get_by_semver(&package.name, &package.version) {
                if record.version() == &version {
                    self.packages.insert((package.name.clone(), version.clone()), package);
                    return Ok(version);
                }
                return Err(Report::msg(format!(
                    "package '{}' version '{}' is already registered",
                    package.name, package.version
                )));
            }
            let dependencies = package.manifest.dependencies().map(|dependency| {
                let version = Version::new(dependency.version.clone(), dependency.digest);
                (dependency.name.clone(), VersionRequirement::Exact(version))
            });
            let record = PackageRecord::new(version.clone(), dependencies)
                .with_description(package.description.clone().unwrap_or_default());
            self.register(package.name.clone(), record)?;
            self.packages.insert((package.name.clone(), version.clone()), package);
            Ok(version)
        }
    }

    impl PackageStore for TestRegistry {
        fn publish_package(&mut self, package: Arc<Package>) -> Result<Version, Self::Error> {
            let version = Version::new(package.version.clone(), package.digest());
            let dependencies = package
                .manifest
                .dependencies()
                .map(|dependency| {
                    let version = Version::new(dependency.version.clone(), dependency.digest);
                    if self.get_exact_version(&dependency.name, &version).is_none() {
                        return Err(Report::msg(format!(
                            "missing dependency '{}' at '{}'",
                            dependency.name, version
                        )));
                    }
                    Ok((dependency.name.clone(), VersionRequirement::Exact(version)))
                })
                .collect::<Result<Vec<_>, _>>()?;
            let record = PackageRecord::new(version.clone(), dependencies)
                .with_description(package.description.clone().unwrap_or_default());
            self.register(package.name.clone(), record)?;
            self.packages.insert((package.name.clone(), version.clone()), package);
            Ok(version)
        }
    }

    impl TestContext {
        pub fn registry(&self) -> &TestRegistry {
            &self.registry
        }

        pub fn registry_mut(&mut self) -> &mut TestRegistry {
            &mut self.registry
        }

        pub fn project_assembler_for_path<'a>(
            &'a mut self,
            manifest_path: impl AsRef<std::path::Path>,
        ) -> Result<crate::ProjectAssembler<'a, TestRegistry>, Report> {
            self.assembler().for_project_at_path(manifest_path, &mut self.registry)
        }

        pub fn project_assembler<'a>(
            &'a mut self,
            project: Arc<miden_project::Package>,
        ) -> Result<crate::ProjectAssembler<'a, TestRegistry>, Report> {
            self.assembler().for_project(project, &mut self.registry)
        }

        /// Assembles the library target of the package at `manifest_path`, using `profile`.
        ///
        /// If `profile` is `None`, the `dev` profile will be used.
        ///
        /// This requires that the referenced package define a library target, or an error will be
        /// raised.
        pub fn assemble_library_package(
            &mut self,
            manifest_path: impl AsRef<std::path::Path>,
            profile: Option<&str>,
        ) -> Result<Arc<Package>, Report> {
            let assembler = self.assembler();
            let mut project_assembler =
                assembler.for_project_at_path(manifest_path, &mut self.registry)?;
            project_assembler.assemble(ProjectTargetSelector::Library, profile.unwrap_or("dev"))
        }

        /// Assembles the executable target `name` of the package from `manifest_path`, using
        /// `profile`.
        ///
        /// If `name` is `None`, the default executable target name is selected.
        ///
        /// If `profile` is `None`, the `dev` profile will be used.
        ///
        /// This requires that the referenced package define the given executable target, or an
        /// error will be raised.
        pub fn assemble_executable_package(
            &mut self,
            manifest_path: impl AsRef<std::path::Path>,
            name: Option<&str>,
            profile: Option<&str>,
        ) -> Result<Arc<Package>, Report> {
            let assembler = self.assembler();
            let mut project_assembler =
                assembler.for_project_at_path(manifest_path, &mut self.registry)?;
            project_assembler.assemble(
                ProjectTargetSelector::Executable(name.unwrap_or(Path::EXEC_PATH)),
                profile.unwrap_or("dev"),
            )
        }

        /// Assembles a package named `name` with `version` and `dependencies`, containing a single
        /// export whose path is `export`.
        ///
        /// The exported procedure is defined as follows, given an `export` path of `foo::bar`:
        ///
        /// ```text,ignore
        /// pub proc bar
        ///     add
        /// end
        /// ```
        pub fn assemble_library_package_with_export<'a>(
            &self,
            name: &str,
            version: &str,
            export: &str,
            dependencies: impl IntoIterator<Item = (&'static str, &'a str, TargetType, Word)>,
        ) -> Box<Package> {
            use alloc::string::ToString;

            use miden_assembly_syntax::source_file;
            use miden_mast_package::Dependency;

            let export_path = Path::new(export);
            let (export_leaf, export_module) = export_path.split_last().unwrap();
            let source_file = source_file!(
                self,
                format!("namespace {export_module}\n\npub proc {export_leaf} add end")
            );
            let module = self.parse_module(source_file).unwrap();
            let name = PackageId::from(name);
            let mut package = self
                .assembler()
                .assemble_library(name, module, None::<Box<Module>>)
                .expect("failed to assemble library");
            package.version = version.parse().unwrap();
            for (name, version, kind, digest) in dependencies {
                package
                    .manifest
                    .add_dependency(Dependency {
                        name: name.into(),
                        kind,
                        version: version.parse().unwrap(),
                        digest,
                    })
                    .expect("failed to add dependency");
            }
            package
        }
    }
}