use std::collections::{BTreeMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};
use harn_modules::{public_declarations, DefKind};
use serde::{Deserialize, Serialize};
use crate::chunk::{CachedChunk, CachedCompiledFunction};
use crate::value::VmError;
type ImportedEnumCache = BTreeMap<PathBuf, ([u8; 32], Vec<String>)>;
fn imported_enum_cache() -> &'static Mutex<ImportedEnumCache> {
static CACHE: OnceLock<Mutex<ImportedEnumCache>> = OnceLock::new();
CACHE.get_or_init(|| Mutex::new(BTreeMap::new()))
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ModuleImportSpec {
pub path: String,
pub selected_names: Option<Vec<String>>,
pub is_pub: bool,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ModuleArtifact {
pub imports: Vec<ModuleImportSpec>,
pub type_schema_init_chunk: Option<CachedChunk>,
pub init_chunk: Option<CachedChunk>,
pub functions: BTreeMap<String, CachedCompiledFunction>,
pub public_exports: BTreeMap<String, DefKind>,
pub public_value_names: HashSet<String>,
pub public_type_names: HashSet<String>,
}
impl ModuleArtifact {
pub(crate) fn bind_source_file(&mut self, source_path: &Path) {
let source_file = source_path.display().to_string();
if let Some(chunk) = &mut self.type_schema_init_chunk {
bind_chunk_source_file(chunk, &source_file);
}
if let Some(chunk) = &mut self.init_chunk {
bind_chunk_source_file(chunk, &source_file);
}
for function in self.functions.values_mut() {
bind_chunk_source_file(&mut function.chunk, &source_file);
}
}
}
fn bind_chunk_source_file(chunk: &mut CachedChunk, source_file: &str) {
chunk.source_file = Some(source_file.to_string());
for function in &mut chunk.functions {
bind_chunk_source_file(&mut function.chunk, source_file);
}
}
pub fn compile_module_artifact(
program: &[harn_parser::SNode],
module_source_file: Option<String>,
) -> Result<ModuleArtifact, VmError> {
let imported_enum_candidates = module_source_file
.as_deref()
.filter(|_| needs_imported_enum_candidates(program))
.and_then(|path| {
harn_modules::build(&[Path::new(path).to_path_buf()])
.imported_names_by_kind_for_file(Path::new(path), DefKind::Enum)
})
.unwrap_or_default();
compile_module_artifact_with_imported_enums(
program,
module_source_file,
&imported_enum_candidates.into_iter().collect::<Vec<_>>(),
)
}
fn compile_module_artifact_with_imported_enums(
program: &[harn_parser::SNode],
module_source_file: Option<String>,
imported_enum_candidates: &[String],
) -> Result<ModuleArtifact, VmError> {
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| {
let inner = match &sn.node {
harn_parser::Node::AttributedDecl { inner, .. } => inner.as_ref(),
_ => sn,
};
matches!(
&inner.node,
harn_parser::Node::LetBinding { .. }
| harn_parser::Node::ConstBinding { .. }
| harn_parser::Node::EnumDecl { is_pub: true, .. }
| harn_parser::Node::ToolDecl { .. }
| harn_parser::Node::SkillDecl { .. }
| harn_parser::Node::EvalPackDecl { .. }
)
})
.cloned()
.collect();
let init_chunk = if init_nodes.is_empty() {
None
} else {
let compiler = crate::Compiler::new();
Some(
compiler
.compile_module_init(program, &init_nodes, imported_enum_candidates)
.map_err(|e| VmError::Runtime(format!("Import init compile error: {e}")))?
.freeze_for_cache(),
)
};
let public_exports: BTreeMap<String, DefKind> = program
.iter()
.flat_map(public_declarations)
.map(|export| (export.name, export.kind))
.collect();
let public_value_names = public_exports
.iter()
.filter(|(_, kind)| {
matches!(
kind,
DefKind::Variable
| DefKind::Enum
| DefKind::Tool
| DefKind::Skill
| DefKind::EvalPack
)
})
.map(|(name, _)| name.clone())
.collect();
let public_type_names = public_exports
.iter()
.filter(|(_, kind)| !kind.has_runtime_value())
.map(|(name, _)| name.clone())
.collect();
let mut functions = BTreeMap::new();
for node in program {
let inner = match &node.node {
harn_parser::Node::AttributedDecl { inner, .. } => inner.as_ref(),
_ => node,
};
if let harn_parser::Node::StructDecl { name, fields, .. } = &inner.node {
let constructor = crate::Compiler::new()
.compile_struct_constructor(name, fields)
.map_err(|error| VmError::Runtime(format!("Import compile error: {error}")))?;
functions.insert(name.clone(), constructor.freeze_for_cache());
continue;
}
if let harn_parser::Node::Pipeline {
name,
params,
body,
extends,
..
} = &inner.node
{
let mut compiler = crate::Compiler::new();
compiler.add_imported_enum_candidates(imported_enum_candidates.iter().cloned());
let pipeline = compiler
.compile_pipeline_callable(program, name, params, body, extends.as_deref())
.map_err(|error| VmError::Runtime(format!("Import compile error: {error}")))?;
functions.insert(name.clone(), pipeline.freeze_for_cache());
continue;
}
let harn_parser::Node::FnDecl {
name,
type_params,
params,
body,
..
} = &inner.node
else {
continue;
};
let mut compiler = crate::Compiler::new();
compiler.add_imported_enum_candidates(imported_enum_candidates.iter().cloned());
compiler.prepare_module_context(program);
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());
}
let type_schema_init_chunk =
crate::Compiler::compile_public_type_schema_initializers(program, module_source_file)
.map_err(|error| VmError::Runtime(format!("Import schema compile error: {error}")))?
.map(|chunk| chunk.freeze_for_cache());
Ok(ModuleArtifact {
imports,
type_schema_init_chunk,
init_chunk,
functions,
public_exports,
public_value_names,
public_type_names,
})
}
pub fn compile_module_artifact_from_source(
source_path: &Path,
source: &str,
) -> Result<ModuleArtifact, VmError> {
let program = parse_module_source(source_path, source)?;
let imported_enum_candidates =
imported_enum_candidates_for_program(source_path, source, &program);
compile_module_artifact_with_imported_enums(
&program,
Some(source_path.display().to_string()),
&imported_enum_candidates,
)
}
fn imported_enum_candidates_for_program(
source_path: &Path,
source: &str,
program: &[harn_parser::SNode],
) -> Vec<String> {
if !needs_imported_enum_candidates(program) {
return Vec::new();
}
let source_hash = *blake3::hash(source.as_bytes()).as_bytes();
let cache_key = harn_modules::canonical_path(source_path);
let cacheable = is_immutable_stdlib_path(source_path);
if cacheable {
if let Some((_cached_hash, candidates)) = imported_enum_cache()
.lock()
.expect("imported enum cache lock poisoned")
.get(&cache_key)
.filter(|(cached_hash, _)| *cached_hash == source_hash)
{
return candidates.clone();
}
}
let graph = harn_modules::build_with_source(source_path, source);
if !cacheable {
return sorted_imported_enum_candidates(&graph, source_path);
}
let mut projections = Vec::new();
for path in graph.module_paths() {
let module_source = if path == cache_key {
Some(source.to_string())
} else {
harn_modules::read_module_source(&path).or_else(|| std::fs::read_to_string(&path).ok())
};
let Some(module_source) = module_source else {
continue;
};
let candidates = sorted_imported_enum_candidates(&graph, &path);
projections.push((
path,
(
*blake3::hash(module_source.as_bytes()).as_bytes(),
candidates,
),
));
}
let mut cache = imported_enum_cache()
.lock()
.expect("imported enum cache lock poisoned");
for (path, projection) in projections {
if is_immutable_stdlib_path(&path) {
cache.insert(path, projection);
}
}
cache
.get(&cache_key)
.filter(|(cached_hash, _)| *cached_hash == source_hash)
.map(|(_, candidates)| candidates.clone())
.unwrap_or_default()
}
fn sorted_imported_enum_candidates(
graph: &harn_modules::ModuleGraph,
source_path: &Path,
) -> Vec<String> {
let mut candidates = graph
.imported_names_by_kind_for_file(source_path, DefKind::Enum)
.unwrap_or_default()
.into_iter()
.collect::<Vec<_>>();
candidates.sort_unstable();
candidates
}
fn is_immutable_stdlib_path(path: &Path) -> bool {
path.to_str()
.is_some_and(|path| path.starts_with("<stdlib>/") || path.starts_with("<std>/"))
}
fn needs_imported_enum_candidates(program: &[harn_parser::SNode]) -> bool {
harn_parser::visit::contains_identifier_enum_pattern(program)
}
fn parse_module_source(
source_path: &Path,
source: &str,
) -> Result<Vec<harn_parser::SNode>, VmError> {
let mut lexer = harn_lexer::Lexer::new(source);
let tokens = lexer.tokenize().map_err(|e| {
VmError::Runtime(format!(
"Import lex error in {}: {e}",
source_path.display()
))
})?;
let mut parser = harn_parser::Parser::new(tokens);
parser.parse().map_err(|e| {
VmError::Runtime(format!(
"Import parse error in {}: {e}",
source_path.display()
))
})
}
pub fn compile_module_artifact_from_source_with_imported_enums(
source_path: &Path,
source: &str,
imported_enum_candidates: impl IntoIterator<Item = String>,
) -> Result<ModuleArtifact, VmError> {
let program = parse_module_source(source_path, source)?;
let imported_enum_candidates = imported_enum_candidates.into_iter().collect::<Vec<_>>();
compile_module_artifact_with_imported_enums(
&program,
Some(source_path.display().to_string()),
&imported_enum_candidates,
)
}
#[cfg(test)]
mod tests {
use std::path::Path;
use harn_lexer::Lexer;
use harn_parser::Parser;
use super::{
compile_module_artifact, compile_module_artifact_from_source,
needs_imported_enum_candidates, parse_module_source,
};
use crate::chunk::Constant;
#[test]
fn module_init_schema_of_uses_full_program_aliases() {
let source = r"
pub type Item = {id: string}
const ITEM_SCHEMA: Schema<Item> = schema_of(Item)
";
let mut lexer = Lexer::new(source);
let tokens = lexer.tokenize().unwrap();
let mut parser = Parser::new(tokens);
let program = parser.parse().unwrap();
let artifact = compile_module_artifact(&program, None).unwrap();
let constants = &artifact.init_chunk.expect("init chunk").constants;
let strings = constants
.iter()
.filter_map(|constant| match constant {
Constant::String(value) => Some(value.as_str()),
_ => None,
})
.collect::<Vec<_>>();
assert!(strings.contains(&"id"), "{strings:?}");
assert!(!strings.contains(&"Item"), "{strings:?}");
}
#[test]
fn type_only_modules_use_a_separate_schema_initializer() {
let source = r"
pub type UserShape = {name: string, active?: bool}
pub type UserList = list<UserShape>
";
let artifact =
compile_module_artifact_from_source(Path::new("<test>/schemas.harn"), source)
.expect("module compiles");
assert!(
artifact.init_chunk.is_none(),
"erased type aliases must not inflate module init bytecode"
);
assert!(artifact.public_type_names.contains("UserShape"));
assert!(artifact.public_type_names.contains("UserList"));
assert!(artifact.type_schema_init_chunk.is_some());
}
#[test]
fn schema_initializer_keeps_imported_alias_lookup_and_source() {
let source = r#"
import { External } from "./external"
pub type Wrapped = {value: External}
"#;
let source_path = Path::new("<test>/wrapped.harn");
let artifact =
compile_module_artifact_from_source(source_path, source).expect("module compiles");
let chunk = artifact.type_schema_init_chunk.expect("schema initializer");
assert_eq!(chunk.source_file.as_deref(), Some("<test>/wrapped.harn"));
assert!(chunk
.constants
.iter()
.any(|constant| matches!(constant, Constant::String(value) if value == "External")));
}
#[test]
fn imported_enum_graph_lookup_is_lazy_for_plain_modules() {
let plain = parse_module_source(
Path::new("<test>/plain.harn"),
r#"
import { helper } from "./support"
pub fn run() -> int { return helper(1) }
"#,
)
.expect("plain module parses");
assert!(!needs_imported_enum_candidates(&plain));
let qualified = parse_module_source(
Path::new("<test>/qualified.harn"),
r#"
import { Status } from "./status"
pub fn run(value: Status) {
match value {
Status.Ready -> { return 1 }
_ -> { return 0 }
}
}
"#,
)
.expect("qualified module parses");
assert!(needs_imported_enum_candidates(&qualified));
}
#[test]
fn private_declarations_do_not_expand_module_init() {
let artifact = compile_module_artifact_from_source(
Path::new("<test>/private-declarations.harn"),
r"
enum PrivateStatus { Ready }
struct PrivateConfig { value: int }
pub fn run() { return PrivateStatus.Ready }
",
)
.expect("private declarations compile");
assert!(artifact.init_chunk.is_none());
assert!(artifact.functions.contains_key("PrivateConfig"));
assert!(!artifact.public_exports.contains_key("PrivateStatus"));
assert!(!artifact.public_exports.contains_key("PrivateConfig"));
}
}