harn-vm 0.8.2

use std::cell::RefCell;
use std::collections::{BTreeMap, HashSet};
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::rc::Rc;
use std::sync::{Arc, Mutex, OnceLock};

use crate::chunk::{CachedChunk, CachedCompiledFunction, Chunk, CompiledFunction};
use crate::value::{ModuleFunctionRegistry, VmClosure, VmEnv, VmError, VmValue};

use super::{ScopeSpan, Vm};

#[derive(Clone)]
struct ModuleImportSpec {
    path: String,
    selected_names: Option<Vec<String>>,
    is_pub: bool,
}

#[derive(Clone)]
struct CompiledStdlibModule {
    imports: Vec<ModuleImportSpec>,
    init_chunk: Option<CachedChunk>,
    functions: BTreeMap<String, CachedCompiledFunction>,
    public_names: HashSet<String>,
}

static STDLIB_MODULE_ARTIFACT_CACHE: OnceLock<Mutex<BTreeMap<String, Arc<CompiledStdlibModule>>>> =
    OnceLock::new();

fn stdlib_module_artifact_cache() -> &'static Mutex<BTreeMap<String, Arc<CompiledStdlibModule>>> {
    STDLIB_MODULE_ARTIFACT_CACHE.get_or_init(|| Mutex::new(BTreeMap::new()))
}

#[cfg(test)]
fn reset_stdlib_module_artifact_cache() {
    stdlib_module_artifact_cache().lock().unwrap().clear();
}

#[cfg(test)]
fn stdlib_module_artifact_cache_len() -> usize {
    stdlib_module_artifact_cache().lock().unwrap().len()
}

#[derive(Clone)]
pub(crate) struct LoadedModule {
    pub(crate) functions: BTreeMap<String, Rc<VmClosure>>,
    pub(crate) public_names: HashSet<String>,
}

pub fn resolve_module_import_path(base: &Path, path: &str) -> PathBuf {
    let synthetic_current_file = base.join("__harn_import_base__.harn");
    if let Some(resolved) = harn_modules::resolve_import_path(&synthetic_current_file, path) {
        return resolved;
    }

    let mut file_path = base.join(path);

    if !file_path.exists() && file_path.extension().is_none() {
        file_path.set_extension("harn");
    }

    file_path
}

fn stdlib_artifact_cache_key(module: &str, source: &str) -> String {
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    module.hash(&mut hasher);
    source.hash(&mut hasher);
    format!("{module}:{:016x}", hasher.finish())
}

fn stdlib_module_artifact(
    module: &str,
    synthetic: &Path,
    source: &str,
) -> Result<Arc<CompiledStdlibModule>, VmError> {
    let key = stdlib_artifact_cache_key(module, source);
    {
        let cache = stdlib_module_artifact_cache().lock().unwrap();
        if let Some(cached) = cache.get(&key) {
            return Ok(Arc::clone(cached));
        }
    }

    let compiled = Arc::new(compile_stdlib_module_artifact(synthetic, source)?);
    let mut cache = stdlib_module_artifact_cache().lock().unwrap();
    if let Some(cached) = cache.get(&key) {
        return Ok(Arc::clone(cached));
    }
    cache.insert(key, Arc::clone(&compiled));
    Ok(compiled)
}

fn compile_stdlib_module_artifact(
    synthetic: &Path,
    source: &str,
) -> Result<CompiledStdlibModule, VmError> {
    let mut lexer = harn_lexer::Lexer::new(source);
    let tokens = lexer.tokenize().map_err(|e| {
        VmError::Runtime(format!("Import lex error in {}: {e}", synthetic.display()))
    })?;
    let mut parser = harn_parser::Parser::new(tokens);
    let program = parser.parse().map_err(|e| {
        VmError::Runtime(format!(
            "Import parse error in {}: {e}",
            synthetic.display()
        ))
    })?;

    let imports = program
        .iter()
        .filter_map(|node| match &node.node {
            harn_parser::Node::ImportDecl { path, is_pub } => Some(ModuleImportSpec {
                path: path.clone(),
                selected_names: None,
                is_pub: *is_pub,
            }),
            harn_parser::Node::SelectiveImport {
                names,
                path,
                is_pub,
            } => Some(ModuleImportSpec {
                path: path.clone(),
                selected_names: Some(names.clone()),
                is_pub: *is_pub,
            }),
            _ => None,
        })
        .collect();

    let init_nodes: Vec<harn_parser::SNode> = program
        .iter()
        .filter(|sn| {
            matches!(
                &sn.node,
                harn_parser::Node::VarBinding { .. } | harn_parser::Node::LetBinding { .. }
            )
        })
        .cloned()
        .collect();
    let init_chunk = if init_nodes.is_empty() {
        None
    } else {
        Some(
            crate::Compiler::new()
                .compile(&init_nodes)
                .map_err(|e| VmError::Runtime(format!("Import init compile error: {e}")))?
                .freeze_for_cache(),
        )
    };

    let mut functions = BTreeMap::new();
    let mut public_names = HashSet::new();
    let module_source_file = Some(synthetic.display().to_string());
    for node in &program {
        let inner = match &node.node {
            harn_parser::Node::AttributedDecl { inner, .. } => inner.as_ref(),
            _ => node,
        };
        let harn_parser::Node::FnDecl {
            name,
            type_params,
            params,
            body,
            is_pub,
            ..
        } = &inner.node
        else {
            continue;
        };

        let mut compiler = crate::Compiler::new();
        let func_chunk = compiler
            .compile_fn_body(type_params, params, body, module_source_file.clone())
            .map_err(|e| VmError::Runtime(format!("Import compile error: {e}")))?;
        functions.insert(name.clone(), func_chunk.freeze_for_cache());
        if *is_pub {
            public_names.insert(name.clone());
        }
    }

    Ok(CompiledStdlibModule {
        imports,
        init_chunk,
        functions,
        public_names,
    })
}

impl Vm {
    async fn load_module_from_source(
        &mut self,
        synthetic: PathBuf,
        source: &str,
    ) -> Result<LoadedModule, VmError> {
        if let Some(loaded) = self.module_cache.get(&synthetic).cloned() {
            return Ok(loaded);
        }
        Rc::make_mut(&mut self.source_cache).insert(synthetic.clone(), source.to_string());

        let mut lexer = harn_lexer::Lexer::new(source);
        let tokens = lexer.tokenize().map_err(|e| {
            VmError::Runtime(format!("Import lex error in {}: {e}", synthetic.display()))
        })?;
        let mut parser = harn_parser::Parser::new(tokens);
        let program = parser.parse().map_err(|e| {
            VmError::Runtime(format!(
                "Import parse error in {}: {e}",
                synthetic.display()
            ))
        })?;

        self.imported_paths.push(synthetic.clone());
        let loaded = self
            .import_declarations(&program, None, Some(&synthetic))
            .await?;
        self.imported_paths.pop();
        Rc::make_mut(&mut self.module_cache).insert(synthetic, loaded.clone());
        Ok(loaded)
    }

    async fn load_stdlib_module_from_source(
        &mut self,
        module: &str,
        synthetic: PathBuf,
        source: &'static str,
    ) -> Result<LoadedModule, VmError> {
        if let Some(loaded) = self.module_cache.get(&synthetic).cloned() {
            return Ok(loaded);
        }
        Rc::make_mut(&mut self.source_cache).insert(synthetic.clone(), source.to_string());

        let artifact = stdlib_module_artifact(module, &synthetic, source)?;
        self.imported_paths.push(synthetic.clone());
        let loaded = self
            .instantiate_stdlib_module(&synthetic, artifact.as_ref())
            .await?;
        self.imported_paths.pop();
        Rc::make_mut(&mut self.module_cache).insert(synthetic, loaded.clone());
        Ok(loaded)
    }

    async fn instantiate_stdlib_module(
        &mut self,
        _synthetic: &Path,
        artifact: &CompiledStdlibModule,
    ) -> Result<LoadedModule, VmError> {
        let caller_env = self.env.clone();
        let old_source_dir = self.source_dir.clone();
        self.env = VmEnv::new();
        self.source_dir = None;

        for import in &artifact.imports {
            self.execute_import(&import.path, import.selected_names.as_deref())
                .await?;
        }

        let module_state: crate::value::ModuleState = {
            let mut init_env = self.env.clone();
            if let Some(init_chunk) = &artifact.init_chunk {
                let fresh_init_chunk = Chunk::from_cached(init_chunk);
                let saved_env = std::mem::replace(&mut self.env, init_env);
                let saved_frames = std::mem::take(&mut self.frames);
                let saved_handlers = std::mem::take(&mut self.exception_handlers);
                let saved_iterators = std::mem::take(&mut self.iterators);
                let saved_deadlines = std::mem::take(&mut self.deadlines);
                let init_result = self.run_chunk(&fresh_init_chunk).await;
                init_env = std::mem::replace(&mut self.env, saved_env);
                self.frames = saved_frames;
                self.exception_handlers = saved_handlers;
                self.iterators = saved_iterators;
                self.deadlines = saved_deadlines;
                init_result?;
            }
            Rc::new(RefCell::new(init_env))
        };

        let module_env = self.env.clone();
        let registry: ModuleFunctionRegistry = Rc::new(RefCell::new(BTreeMap::new()));
        let mut functions: BTreeMap<String, Rc<VmClosure>> = BTreeMap::new();
        let mut public_names = artifact.public_names.clone();

        for (name, compiled) in &artifact.functions {
            let closure = Rc::new(VmClosure {
                func: Rc::new(CompiledFunction::from_cached(compiled)),
                env: module_env.clone(),
                source_dir: None,
                module_functions: Some(Rc::clone(&registry)),
                module_state: Some(Rc::clone(&module_state)),
            });
            registry
                .borrow_mut()
                .insert(name.clone(), Rc::clone(&closure));
            self.env
                .define(name, VmValue::Closure(Rc::clone(&closure)), false)?;
            module_state
                .borrow_mut()
                .define(name, VmValue::Closure(Rc::clone(&closure)), false)?;
            functions.insert(name.clone(), Rc::clone(&closure));
        }

        for import in artifact.imports.iter().filter(|import| import.is_pub) {
            let cache_key = self.cache_key_for_import(&import.path);
            let Some(loaded) = self.module_cache.get(&cache_key).cloned() else {
                return Err(VmError::Runtime(format!(
                    "Re-export error: imported module '{}' was not loaded",
                    import.path
                )));
            };
            let names_to_reexport: Vec<String> = match &import.selected_names {
                Some(names) => names.clone(),
                None => {
                    if loaded.public_names.is_empty() {
                        loaded.functions.keys().cloned().collect()
                    } else {
                        loaded.public_names.iter().cloned().collect()
                    }
                }
            };
            for name in names_to_reexport {
                let Some(closure) = loaded.functions.get(&name) else {
                    return Err(VmError::Runtime(format!(
                        "Re-export error: '{name}' is not exported by '{}'",
                        import.path
                    )));
                };
                if let Some(existing) = functions.get(&name) {
                    if !Rc::ptr_eq(existing, closure) {
                        return Err(VmError::Runtime(format!(
                            "Re-export collision: '{name}' is defined here and also \
                             re-exported from '{}'",
                            import.path
                        )));
                    }
                }
                functions.insert(name.clone(), Rc::clone(closure));
                public_names.insert(name);
            }
        }

        self.env = caller_env;
        self.source_dir = old_source_dir;

        Ok(LoadedModule {
            functions,
            public_names,
        })
    }

    fn export_loaded_module(
        &mut self,
        module_path: &Path,
        loaded: &LoadedModule,
        selected_names: Option<&[String]>,
    ) -> Result<(), VmError> {
        let export_names: Vec<String> = if let Some(names) = selected_names {
            names.to_vec()
        } else if !loaded.public_names.is_empty() {
            loaded.public_names.iter().cloned().collect()
        } else {
            loaded.functions.keys().cloned().collect()
        };

        let module_name = module_path.display().to_string();
        for name in export_names {
            let Some(closure) = loaded.functions.get(&name) else {
                return Err(VmError::Runtime(format!(
                    "Import error: '{name}' is not defined in {module_name}"
                )));
            };
            if let Some(VmValue::Closure(_)) = self.env.get(&name) {
                return Err(VmError::Runtime(format!(
                    "Import collision: '{name}' is already defined when importing {module_name}. \
                     Use selective imports to disambiguate: import {{ {name} }} from \"...\""
                )));
            }
            self.env
                .define(&name, VmValue::Closure(Rc::clone(closure)), false)?;
        }
        Ok(())
    }

    /// Execute an import, reading and running the file's declarations.
    pub(super) fn execute_import<'a>(
        &'a mut self,
        path: &'a str,
        selected_names: Option<&'a [String]>,
    ) -> Pin<Box<dyn Future<Output = Result<(), VmError>> + 'a>> {
        Box::pin(async move {
            let _import_span = ScopeSpan::new(crate::tracing::SpanKind::Import, path.to_string());

            if let Some(module) = path.strip_prefix("std/") {
                if let Some(source) = crate::stdlib_modules::get_stdlib_source(module) {
                    let synthetic = PathBuf::from(format!("<stdlib>/{module}.harn"));
                    if self.imported_paths.contains(&synthetic) {
                        return Ok(());
                    }
                    let loaded = self
                        .load_stdlib_module_from_source(module, synthetic.clone(), source)
                        .await?;
                    self.export_loaded_module(&synthetic, &loaded, selected_names)?;
                    return Ok(());
                }
                return Err(VmError::Runtime(format!(
                    "Unknown stdlib module: std/{module}"
                )));
            }

            let base = self
                .source_dir
                .clone()
                .unwrap_or_else(|| PathBuf::from("."));
            let file_path = resolve_module_import_path(&base, path);

            let canonical = file_path
                .canonicalize()
                .unwrap_or_else(|_| file_path.clone());
            if self.imported_paths.contains(&canonical) {
                return Ok(());
            }
            if let Some(loaded) = self.module_cache.get(&canonical).cloned() {
                return self.export_loaded_module(&canonical, &loaded, selected_names);
            }
            self.imported_paths.push(canonical.clone());

            let source = std::fs::read_to_string(&file_path).map_err(|e| {
                VmError::Runtime(format!(
                    "Import error: cannot read '{}': {e}",
                    file_path.display()
                ))
            })?;
            Rc::make_mut(&mut self.source_cache).insert(canonical.clone(), source.clone());
            Rc::make_mut(&mut self.source_cache).insert(file_path.clone(), source.clone());

            let mut lexer = harn_lexer::Lexer::new(&source);
            let tokens = lexer
                .tokenize()
                .map_err(|e| VmError::Runtime(format!("Import lex error: {e}")))?;
            let mut parser = harn_parser::Parser::new(tokens);
            let program = parser
                .parse()
                .map_err(|e| VmError::Runtime(format!("Import parse error: {e}")))?;

            let loaded = self
                .import_declarations(&program, Some(&file_path), Some(&file_path))
                .await?;
            self.imported_paths.pop();
            Rc::make_mut(&mut self.module_cache).insert(canonical.clone(), loaded.clone());
            self.export_loaded_module(&canonical, &loaded, selected_names)?;

            Ok(())
        })
    }

    /// Process top-level declarations from an imported module.
    fn import_declarations<'a>(
        &'a mut self,
        program: &'a [harn_parser::SNode],
        file_path: Option<&'a Path>,
        debug_source_file: Option<&'a Path>,
    ) -> Pin<Box<dyn Future<Output = Result<LoadedModule, VmError>> + 'a>> {
        Box::pin(async move {
            let caller_env = self.env.clone();
            let old_source_dir = self.source_dir.clone();
            self.env = VmEnv::new();
            if let Some(fp) = file_path {
                if let Some(parent) = fp.parent() {
                    self.source_dir = Some(parent.to_path_buf());
                }
            }

            for node in program {
                match &node.node {
                    harn_parser::Node::ImportDecl { path: sub_path, .. } => {
                        self.execute_import(sub_path, None).await?;
                    }
                    harn_parser::Node::SelectiveImport {
                        names,
                        path: sub_path,
                        ..
                    } => {
                        self.execute_import(sub_path, Some(names)).await?;
                    }
                    _ => {}
                }
            }

            // Route top-level `var`/`let` bindings into a shared
            // `module_state` rather than `module_env`. If they appeared in
            // `module_env` (captured by each closure's lexical snapshot),
            // every call's per-invocation env clone would shadow them and
            // writes would land in a per-call copy discarded on return.
            let module_state: crate::value::ModuleState = {
                let mut init_env = self.env.clone();
                let init_nodes: Vec<harn_parser::SNode> = program
                    .iter()
                    .filter(|sn| {
                        matches!(
                            &sn.node,
                            harn_parser::Node::VarBinding { .. }
                                | harn_parser::Node::LetBinding { .. }
                        )
                    })
                    .cloned()
                    .collect();
                if !init_nodes.is_empty() {
                    let init_compiler = crate::Compiler::new();
                    let init_chunk = init_compiler
                        .compile(&init_nodes)
                        .map_err(|e| VmError::Runtime(format!("Import init compile error: {e}")))?;
                    // Save frame state so run_chunk_entry's top-level
                    // frame-pop doesn't restore self.env.
                    let saved_env = std::mem::replace(&mut self.env, init_env);
                    let saved_frames = std::mem::take(&mut self.frames);
                    let saved_handlers = std::mem::take(&mut self.exception_handlers);
                    let saved_iterators = std::mem::take(&mut self.iterators);
                    let saved_deadlines = std::mem::take(&mut self.deadlines);
                    let init_result = self.run_chunk(&init_chunk).await;
                    init_env = std::mem::replace(&mut self.env, saved_env);
                    self.frames = saved_frames;
                    self.exception_handlers = saved_handlers;
                    self.iterators = saved_iterators;
                    self.deadlines = saved_deadlines;
                    init_result?;
                }
                Rc::new(RefCell::new(init_env))
            };

            let module_env = self.env.clone();
            let registry: ModuleFunctionRegistry = Rc::new(RefCell::new(BTreeMap::new()));
            let source_dir = file_path.and_then(|fp| fp.parent().map(|p| p.to_path_buf()));
            let mut functions: BTreeMap<String, Rc<VmClosure>> = BTreeMap::new();
            let mut public_names: HashSet<String> = HashSet::new();

            for node in program {
                // Imports may carry `@deprecated` / `@test` etc. on top-level
                // fn decls; transparently peel the wrapper before pattern
                // matching the FnDecl shape.
                let inner = match &node.node {
                    harn_parser::Node::AttributedDecl { inner, .. } => inner.as_ref(),
                    _ => node,
                };
                let harn_parser::Node::FnDecl {
                    name,
                    type_params,
                    params,
                    body,
                    is_pub,
                    ..
                } = &inner.node
                else {
                    continue;
                };

                let mut compiler = crate::Compiler::new();
                let module_source_file = debug_source_file.map(|p| p.display().to_string());
                let func_chunk = compiler
                    .compile_fn_body(type_params, params, body, module_source_file)
                    .map_err(|e| VmError::Runtime(format!("Import compile error: {e}")))?;
                let closure = Rc::new(VmClosure {
                    func: Rc::new(func_chunk),
                    env: module_env.clone(),
                    source_dir: source_dir.clone(),
                    module_functions: Some(Rc::clone(&registry)),
                    module_state: Some(Rc::clone(&module_state)),
                });
                registry
                    .borrow_mut()
                    .insert(name.clone(), Rc::clone(&closure));
                self.env
                    .define(name, VmValue::Closure(Rc::clone(&closure)), false)?;
                // Publish into module_state so sibling fns can be read
                // as VALUES (e.g. `{handler: other_fn}` or as callbacks).
                // Closures captured module_env BEFORE fn decls were added,
                // so their static env alone can't resolve sibling fns.
                // Direct calls use the module_functions late-binding path;
                // value reads rely on this module_state entry.
                module_state.borrow_mut().define(
                    name,
                    VmValue::Closure(Rc::clone(&closure)),
                    false,
                )?;
                functions.insert(name.clone(), Rc::clone(&closure));
                if *is_pub {
                    public_names.insert(name.clone());
                }
            }

            // Re-export pass: for every `pub import ...` declaration in
            // the module, surface the imported closures in this module's
            // `functions`/`public_names` so callers that import this
            // module see the re-exported names.
            for node in program {
                let (sub_path, selective_names, is_pub_import) = match &node.node {
                    harn_parser::Node::ImportDecl {
                        path: sub_path,
                        is_pub,
                    } => (sub_path.clone(), None, *is_pub),
                    harn_parser::Node::SelectiveImport {
                        names,
                        path: sub_path,
                        is_pub,
                    } => (sub_path.clone(), Some(names.clone()), *is_pub),
                    _ => continue,
                };
                if !is_pub_import {
                    continue;
                }
                let cache_key = self.cache_key_for_import(&sub_path);
                let Some(loaded) = self.module_cache.get(&cache_key).cloned() else {
                    return Err(VmError::Runtime(format!(
                        "Re-export error: imported module '{sub_path}' was not loaded"
                    )));
                };
                let names_to_reexport: Vec<String> = match selective_names {
                    Some(names) => names,
                    None => {
                        if loaded.public_names.is_empty() {
                            loaded.functions.keys().cloned().collect()
                        } else {
                            loaded.public_names.iter().cloned().collect()
                        }
                    }
                };
                for name in names_to_reexport {
                    let Some(closure) = loaded.functions.get(&name) else {
                        return Err(VmError::Runtime(format!(
                            "Re-export error: '{name}' is not exported by '{sub_path}'"
                        )));
                    };
                    if let Some(existing) = functions.get(&name) {
                        if !Rc::ptr_eq(existing, closure) {
                            return Err(VmError::Runtime(format!(
                                "Re-export collision: '{name}' is defined here and also \
                                 re-exported from '{sub_path}'"
                            )));
                        }
                    }
                    functions.insert(name.clone(), Rc::clone(closure));
                    public_names.insert(name);
                }
            }

            self.env = caller_env;
            self.source_dir = old_source_dir;

            Ok(LoadedModule {
                functions,
                public_names,
            })
        })
    }

    /// Return the path key that `execute_import` would use to cache the
    /// LoadedModule for this import string. Used by the re-export pass to
    /// look up the already-loaded source module after `execute_import`
    /// has populated [`Vm::module_cache`].
    fn cache_key_for_import(&self, path: &str) -> PathBuf {
        if let Some(module) = path.strip_prefix("std/") {
            return PathBuf::from(format!("<stdlib>/{module}.harn"));
        }
        let base = self
            .source_dir
            .clone()
            .unwrap_or_else(|| PathBuf::from("."));
        let file_path = resolve_module_import_path(&base, path);
        file_path.canonicalize().unwrap_or(file_path)
    }

    /// Load a module file and return the exported function closures that
    /// would be visible to a wildcard import.
    pub async fn load_module_exports(
        &mut self,
        path: &Path,
    ) -> Result<BTreeMap<String, Rc<VmClosure>>, VmError> {
        let path_str = path.to_string_lossy().into_owned();
        self.execute_import(&path_str, None).await?;

        let mut file_path = if path.is_absolute() {
            path.to_path_buf()
        } else {
            self.source_dir
                .clone()
                .unwrap_or_else(|| PathBuf::from("."))
                .join(path)
        };
        if !file_path.exists() && file_path.extension().is_none() {
            file_path.set_extension("harn");
        }

        let canonical = file_path
            .canonicalize()
            .unwrap_or_else(|_| file_path.clone());
        let loaded = self.module_cache.get(&canonical).cloned().ok_or_else(|| {
            VmError::Runtime(format!(
                "Import error: failed to cache loaded module '{}'",
                canonical.display()
            ))
        })?;

        let export_names: Vec<String> = if loaded.public_names.is_empty() {
            loaded.functions.keys().cloned().collect()
        } else {
            loaded.public_names.iter().cloned().collect()
        };

        let mut exports = BTreeMap::new();
        for name in export_names {
            let Some(closure) = loaded.functions.get(&name) else {
                return Err(VmError::Runtime(format!(
                    "Import error: exported function '{name}' is missing from {}",
                    canonical.display()
                )));
            };
            exports.insert(name, Rc::clone(closure));
        }

        Ok(exports)
    }

    /// Load synthetic source keyed by a synthetic module path and return
    /// the exported function closures that a wildcard import would expose.
    pub async fn load_module_exports_from_source(
        &mut self,
        source_key: impl Into<PathBuf>,
        source: &str,
    ) -> Result<BTreeMap<String, Rc<VmClosure>>, VmError> {
        let synthetic = source_key.into();
        let loaded = self
            .load_module_from_source(synthetic.clone(), source)
            .await?;
        let export_names: Vec<String> = if loaded.public_names.is_empty() {
            loaded.functions.keys().cloned().collect()
        } else {
            loaded.public_names.iter().cloned().collect()
        };

        let mut exports = BTreeMap::new();
        for name in export_names {
            let Some(closure) = loaded.functions.get(&name) else {
                return Err(VmError::Runtime(format!(
                    "Import error: exported function '{name}' is missing from {}",
                    synthetic.display()
                )));
            };
            exports.insert(name, Rc::clone(closure));
        }

        Ok(exports)
    }

    /// Load a module by import path (`std/foo`, relative module path, or
    /// package import) and return the exported function closures that a
    /// wildcard import would expose.
    pub async fn load_module_exports_from_import(
        &mut self,
        import_path: &str,
    ) -> Result<BTreeMap<String, Rc<VmClosure>>, VmError> {
        self.execute_import(import_path, None).await?;

        if let Some(module) = import_path.strip_prefix("std/") {
            let synthetic = PathBuf::from(format!("<stdlib>/{module}.harn"));
            let loaded = self.module_cache.get(&synthetic).cloned().ok_or_else(|| {
                VmError::Runtime(format!(
                    "Import error: failed to cache loaded module '{}'",
                    synthetic.display()
                ))
            })?;
            let mut exports = BTreeMap::new();
            let export_names: Vec<String> = if loaded.public_names.is_empty() {
                loaded.functions.keys().cloned().collect()
            } else {
                loaded.public_names.iter().cloned().collect()
            };
            for name in export_names {
                let Some(closure) = loaded.functions.get(&name) else {
                    return Err(VmError::Runtime(format!(
                        "Import error: exported function '{name}' is missing from {}",
                        synthetic.display()
                    )));
                };
                exports.insert(name, Rc::clone(closure));
            }
            return Ok(exports);
        }

        let base = self
            .source_dir
            .clone()
            .unwrap_or_else(|| PathBuf::from("."));
        let file_path = resolve_module_import_path(&base, import_path);
        self.load_module_exports(&file_path).await
    }
}

#[cfg(test)]
mod tests {
    use std::rc::Rc;
    use std::sync::{Mutex, MutexGuard, OnceLock};

    use super::*;

    static CACHE_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();

    fn cache_test_guard() -> MutexGuard<'static, ()> {
        CACHE_TEST_LOCK
            .get_or_init(|| Mutex::new(()))
            .lock()
            .unwrap()
    }

    #[test]
    fn stdlib_artifact_cache_reuses_compilation_with_fresh_vm_state() {
        let _guard = cache_test_guard();
        reset_stdlib_module_artifact_cache();
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("runtime builds");

        let first_exports = runtime.block_on(async {
            let mut first_vm = Vm::new();
            first_vm
                .load_module_exports_from_import("std/agent/prompts")
                .await
                .expect("first stdlib import succeeds")
        });
        assert_eq!(stdlib_module_artifact_cache_len(), 1);

        let second_exports = runtime.block_on(async {
            let mut second_vm = Vm::new();
            second_vm
                .load_module_exports_from_import("std/agent/prompts")
                .await
                .expect("second stdlib import succeeds")
        });
        assert_eq!(stdlib_module_artifact_cache_len(), 1);

        let first = first_exports
            .get("render_agent_prompt")
            .expect("first export exists");
        let second = second_exports
            .get("render_agent_prompt")
            .expect("second export exists");

        assert!(!Rc::ptr_eq(first, second));
        assert!(!Rc::ptr_eq(&first.func, &second.func));
        assert!(!Rc::ptr_eq(&first.func.chunk, &second.func.chunk));
        assert!(!Rc::ptr_eq(
            first.module_state.as_ref().expect("first module state"),
            second.module_state.as_ref().expect("second module state")
        ));
    }

    #[test]
    fn stdlib_artifact_cache_is_process_wide_across_threads() {
        let _guard = cache_test_guard();
        reset_stdlib_module_artifact_cache();

        let handle = std::thread::spawn(|| {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .expect("runtime builds");
            runtime.block_on(async {
                let mut vm = Vm::new();
                vm.load_module_exports_from_import("std/agent/prompts")
                    .await
                    .expect("thread stdlib import succeeds");
            });
        });
        handle.join().expect("thread joins");
        assert_eq!(stdlib_module_artifact_cache_len(), 1);

        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("runtime builds");
        runtime.block_on(async {
            let mut vm = Vm::new();
            vm.load_module_exports_from_import("std/agent/prompts")
                .await
                .expect("main-thread stdlib import succeeds");
        });
        assert_eq!(stdlib_module_artifact_cache_len(), 1);
    }
}