Skip to main content

proof_engine/scripting/
modules.rs

1//! Module system — loading, caching, dependency resolution, hot-reload.
2//!
3//! # Architecture
4//! ```text
5//! PackageManager
6//!   └─ ModuleRegistry
7//!        ├─ Module (Unloaded / Loading / Loaded / Error)
8//!        └─ StringMapLoader / FileLoader
9//! Namespace  — hierarchical name lookup
10//! HotReloadWatcher — timestamp-based change detection
11//! ```
12
13use std::collections::HashMap;
14use std::sync::Arc;
15
16use super::compiler::{Chunk, Compiler};
17use super::parser::Parser;
18use super::vm::{ScriptError, Table, Value, Vm};
19
20// ── ModuleId ──────────────────────────────────────────────────────────────────
21
22/// FNV-1a hash of a module path.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
24pub struct ModuleId(pub u64);
25
26impl ModuleId {
27    pub fn from_path(path: &str) -> Self {
28        ModuleId(fnv1a(path.as_bytes()))
29    }
30}
31
32fn fnv1a(data: &[u8]) -> u64 {
33    let mut hash: u64 = 0xcbf29ce484222325;
34    for &b in data {
35        hash ^= b as u64;
36        hash = hash.wrapping_mul(0x100000001b3);
37    }
38    hash
39}
40
41impl std::fmt::Display for ModuleId {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        write!(f, "ModuleId({:016x})", self.0)
44    }
45}
46
47// ── LoadStatus ────────────────────────────────────────────────────────────────
48
49#[derive(Debug, Clone, PartialEq)]
50pub enum LoadStatus {
51    Unloaded,
52    /// Currently being loaded — used to detect circular deps.
53    Loading,
54    Loaded,
55    Error(String),
56}
57
58// ── Module ────────────────────────────────────────────────────────────────────
59
60/// A compiled and (optionally) executed script module.
61#[derive(Debug, Clone)]
62pub struct Module {
63    pub id:           ModuleId,
64    pub name:         String,
65    pub source_path:  String,
66    pub chunk:        Option<Arc<Chunk>>,
67    pub exports:      HashMap<String, Value>,
68    pub dependencies: Vec<ModuleId>,
69    pub status:       LoadStatus,
70}
71
72impl Module {
73    pub fn new(name: impl Into<String>, source_path: impl Into<String>) -> Self {
74        let name = name.into();
75        let path = source_path.into();
76        Module {
77            id:           ModuleId::from_path(&path),
78            name,
79            source_path:  path,
80            chunk:        None,
81            exports:      HashMap::new(),
82            dependencies: Vec::new(),
83            status:       LoadStatus::Unloaded,
84        }
85    }
86
87    /// Build a Value::Table from the exports map.
88    pub fn exports_table(&self) -> Table {
89        let t = Table::new();
90        for (k, v) in &self.exports {
91            t.rawset_str(k, v.clone());
92        }
93        t
94    }
95}
96
97// ── ModuleLoader trait ────────────────────────────────────────────────────────
98
99/// Pluggable backend for loading module source code.
100pub trait ModuleLoader: Send + Sync {
101    fn load_source(&self, path: &str) -> Result<String, String>;
102}
103
104// ── StringMapLoader ───────────────────────────────────────────────────────────
105
106/// Loads modules from an in-memory map of `path -> source`.
107pub struct StringMapLoader {
108    pub map: HashMap<String, String>,
109}
110
111impl StringMapLoader {
112    pub fn new() -> Self { StringMapLoader { map: HashMap::new() } }
113
114    pub fn add(&mut self, path: impl Into<String>, source: impl Into<String>) {
115        self.map.insert(path.into(), source.into());
116    }
117}
118
119impl Default for StringMapLoader {
120    fn default() -> Self { Self::new() }
121}
122
123impl ModuleLoader for StringMapLoader {
124    fn load_source(&self, path: &str) -> Result<String, String> {
125        self.map.get(path)
126            .cloned()
127            .ok_or_else(|| format!("module not found: {}", path))
128    }
129}
130
131// ── ModuleRegistry ────────────────────────────────────────────────────────────
132
133/// Maintains the full module graph with circular-dependency detection.
134pub struct ModuleRegistry {
135    modules: HashMap<ModuleId, Module>,
136    loader:  Box<dyn ModuleLoader>,
137    /// DFS coloring for cycle detection: gray = currently on stack, black = done.
138    gray:    std::collections::HashSet<ModuleId>,
139    black:   std::collections::HashSet<ModuleId>,
140}
141
142impl ModuleRegistry {
143    pub fn new(loader: Box<dyn ModuleLoader>) -> Self {
144        ModuleRegistry {
145            modules: HashMap::new(),
146            loader,
147            gray:    std::collections::HashSet::new(),
148            black:   std::collections::HashSet::new(),
149        }
150    }
151
152    pub fn with_string_map(map: HashMap<String, String>) -> Self {
153        let mut loader = StringMapLoader::new();
154        loader.map = map;
155        Self::new(Box::new(loader))
156    }
157
158    /// `require(path, vm)` — load, compile and execute a module; cache result.
159    pub fn require(&mut self, path: &str, vm: &mut Vm) -> Result<Value, ScriptError> {
160        let id = ModuleId::from_path(path);
161
162        // Already loaded?
163        if let Some(m) = self.modules.get(&id) {
164            match &m.status {
165                LoadStatus::Loaded => {
166                    return Ok(Value::Table(m.exports_table()));
167                }
168                LoadStatus::Loading => {
169                    return Err(ScriptError::new(format!(
170                        "circular dependency detected for module '{}'", path
171                    )));
172                }
173                LoadStatus::Error(e) => {
174                    return Err(ScriptError::new(format!("module '{}' failed: {}", path, e)));
175                }
176                LoadStatus::Unloaded => {}
177            }
178        }
179
180        // Cycle check
181        if self.gray.contains(&id) {
182            return Err(ScriptError::new(format!("circular dependency: '{}'", path)));
183        }
184        self.gray.insert(id);
185
186        // Load source
187        let source = self.loader.load_source(path).map_err(|e| ScriptError::new(e))?;
188
189        // Register module as loading
190        let mut module = Module::new(path, path);
191        module.status  = LoadStatus::Loading;
192        self.modules.insert(id, module);
193
194        // Compile
195        let chunk = match Parser::from_source(path, &source) {
196            Ok(script) => Compiler::compile_script(&script),
197            Err(e) => {
198                if let Some(m) = self.modules.get_mut(&id) {
199                    m.status = LoadStatus::Error(e.to_string());
200                }
201                self.gray.remove(&id);
202                return Err(ScriptError::new(format!("parse error in '{}': {}", path, e)));
203            }
204        };
205
206        // Execute
207        let result = vm.execute(Arc::clone(&chunk));
208        self.gray.remove(&id);
209        self.black.insert(id);
210
211        match result {
212            Ok(vals) => {
213                // The module's "exports" are whatever it returns or sets in globals
214                let exports_val = vals.into_iter().next().unwrap_or(Value::Nil);
215                let mut exports_map = HashMap::new();
216                if let Value::Table(t) = &exports_val {
217                    let mut key = Value::Nil;
218                    loop {
219                        match t.next(&key) {
220                            Some((k, v)) => {
221                                if let Value::Str(ks) = &k {
222                                    exports_map.insert(ks.as_ref().clone(), v);
223                                }
224                                key = k;
225                            }
226                            None => break,
227                        }
228                    }
229                }
230                if let Some(m) = self.modules.get_mut(&id) {
231                    m.chunk   = Some(chunk);
232                    m.exports = exports_map;
233                    m.status  = LoadStatus::Loaded;
234                }
235                Ok(exports_val)
236            }
237            Err(e) => {
238                if let Some(m) = self.modules.get_mut(&id) {
239                    m.status = LoadStatus::Error(e.message.clone());
240                }
241                Err(e)
242            }
243        }
244    }
245
246    /// Reload a module by path (re-compile and re-execute).
247    pub fn reload(&mut self, path: &str, vm: &mut Vm) -> Result<Value, ScriptError> {
248        let id = ModuleId::from_path(path);
249        // Reset status to unloaded
250        if let Some(m) = self.modules.get_mut(&id) {
251            m.status = LoadStatus::Unloaded;
252        }
253        self.black.remove(&id);
254        self.require(path, vm)
255    }
256
257    /// Remove a module from the cache.
258    pub fn unload(&mut self, path: &str) {
259        let id = ModuleId::from_path(path);
260        self.modules.remove(&id);
261        self.black.remove(&id);
262    }
263
264    /// List all loaded module names.
265    pub fn loaded_modules(&self) -> Vec<String> {
266        self.modules.values()
267            .filter(|m| m.status == LoadStatus::Loaded)
268            .map(|m| m.name.clone())
269            .collect()
270    }
271
272    pub fn get_module(&self, path: &str) -> Option<&Module> {
273        self.modules.get(&ModuleId::from_path(path))
274    }
275
276    /// ASCII-art dependency tree for a module.
277    pub fn dependency_tree(&self, path: &str) -> String {
278        let mut out = String::new();
279        self.dep_tree_inner(path, 0, &mut std::collections::HashSet::new(), &mut out);
280        out
281    }
282
283    fn dep_tree_inner(
284        &self,
285        path: &str,
286        depth: usize,
287        visited: &mut std::collections::HashSet<ModuleId>,
288        out: &mut String,
289    ) {
290        let id  = ModuleId::from_path(path);
291        let prefix = if depth == 0 { String::new() } else {
292            format!("{}{}", "│  ".repeat(depth - 1), "├─ ")
293        };
294        out.push_str(&format!("{}{}\n", prefix, path));
295        if visited.contains(&id) {
296            out.push_str(&format!("{}{}  (already shown)\n", "│  ".repeat(depth), "└─"));
297            return;
298        }
299        visited.insert(id);
300        if let Some(m) = self.modules.get(&id) {
301            let deps: Vec<ModuleId> = m.dependencies.clone();
302            for dep_id in deps {
303                // Find module name by id
304                if let Some(dep_m) = self.modules.values().find(|m2| m2.id == dep_id) {
305                    let dep_path = dep_m.source_path.clone();
306                    self.dep_tree_inner(&dep_path, depth + 1, visited, out);
307                }
308            }
309        }
310    }
311}
312
313// ── PackageManager ────────────────────────────────────────────────────────────
314
315/// Tries multiple search paths and file suffixes when loading modules.
316pub struct PackageManager {
317    pub search_paths: Vec<String>,
318    pub suffixes:     Vec<String>,
319    pub registry:     ModuleRegistry,
320}
321
322impl PackageManager {
323    pub fn new(loader: Box<dyn ModuleLoader>) -> Self {
324        PackageManager {
325            search_paths: vec![String::new()],
326            suffixes:     vec![".lua".to_string(), "".to_string()],
327            registry:     ModuleRegistry::new(loader),
328        }
329    }
330
331    /// Add a search path prefix.
332    pub fn add_path(&mut self, path: impl Into<String>) {
333        self.search_paths.push(path.into());
334    }
335
336    /// require(name, vm) — tries each search_path + suffix combination.
337    pub fn require(&mut self, name: &str, vm: &mut Vm) -> Result<Value, ScriptError> {
338        // Already cached?
339        let id = ModuleId::from_path(name);
340        if let Some(m) = self.registry.modules.get(&id) {
341            if m.status == LoadStatus::Loaded {
342                return Ok(Value::Table(m.exports_table()));
343            }
344        }
345
346        // Try each candidate path
347        let paths: Vec<String> = self.search_paths.iter().flat_map(|base| {
348            self.suffixes.iter().map(move |suf| {
349                if base.is_empty() {
350                    format!("{}{}", name, suf)
351                } else {
352                    format!("{}/{}{}", base, name, suf)
353                }
354            })
355        }).collect();
356
357        for candidate in &paths {
358            // Check if loader can find this
359            if self.registry.loader.load_source(candidate).is_ok() {
360                return self.registry.require(candidate, vm);
361            }
362        }
363        Err(ScriptError::new(format!("module '{}' not found in path", name)))
364    }
365
366    pub fn loaded_modules(&self) -> Vec<String> {
367        self.registry.loaded_modules()
368    }
369}
370
371// ── Namespace ─────────────────────────────────────────────────────────────────
372
373/// Hierarchical name lookup.  E.g. `math.sin` → table field.
374pub struct Namespace {
375    pub name:     String,
376    pub children: HashMap<String, Namespace>,
377    pub values:   HashMap<String, Value>,
378}
379
380impl Namespace {
381    pub fn new(name: impl Into<String>) -> Self {
382        Namespace {
383            name:     name.into(),
384            children: HashMap::new(),
385            values:   HashMap::new(),
386        }
387    }
388
389    /// Set a value at a dotted path, e.g. `"math.sin"`.
390    pub fn set(&mut self, path: &str, value: Value) {
391        let parts: Vec<&str> = path.splitn(2, '.').collect();
392        if parts.len() == 1 {
393            self.values.insert(parts[0].to_string(), value);
394        } else {
395            self.children
396                .entry(parts[0].to_string())
397                .or_insert_with(|| Namespace::new(parts[0]))
398                .set(parts[1], value);
399        }
400    }
401
402    /// Get a value at a dotted path.
403    pub fn get(&self, path: &str) -> Option<&Value> {
404        let parts: Vec<&str> = path.splitn(2, '.').collect();
405        if parts.len() == 1 {
406            self.values.get(parts[0])
407        } else {
408            self.children.get(parts[0])?.get(parts[1])
409        }
410    }
411
412    /// Import all values from this namespace into a VM as globals.
413    pub fn import_into(&self, vm: &mut Vm, prefix: &str) {
414        for (k, v) in &self.values {
415            let name = if prefix.is_empty() { k.clone() } else { format!("{}.{}", prefix, k) };
416            vm.set_global(&name, v.clone());
417        }
418        for (child_name, child_ns) in &self.children {
419            let new_prefix = if prefix.is_empty() {
420                child_name.clone()
421            } else {
422                format!("{}.{}", prefix, child_name)
423            };
424            child_ns.import_into(vm, &new_prefix);
425        }
426    }
427
428    /// Export this namespace as a table value.
429    pub fn export_table(&self) -> Value {
430        let t = Table::new();
431        for (k, v) in &self.values {
432            t.rawset_str(k, v.clone());
433        }
434        for (child_name, child_ns) in &self.children {
435            t.rawset_str(child_name, child_ns.export_table());
436        }
437        Value::Table(t)
438    }
439
440    /// Merge another namespace into this one (other takes precedence).
441    pub fn merge_namespaces(&mut self, other: &Namespace) {
442        for (k, v) in &other.values {
443            self.values.insert(k.clone(), v.clone());
444        }
445        for (k, child) in &other.children {
446            self.children
447                .entry(k.clone())
448                .or_insert_with(|| Namespace::new(k))
449                .merge_namespaces(child);
450        }
451    }
452}
453
454// ── HotReloadWatcher ─────────────────────────────────────────────────────────
455
456/// Tracks file modification timestamps and detects changes.
457/// Uses a `HashMap<String, u64>` for timestamps (simulated or real).
458pub struct HotReloadWatcher {
459    /// path -> last seen timestamp (or content hash)
460    timestamps: HashMap<String, u64>,
461    /// The registry to reload from.
462    loader_snapshots: HashMap<String, String>,
463}
464
465impl HotReloadWatcher {
466    pub fn new() -> Self {
467        HotReloadWatcher {
468            timestamps:       HashMap::new(),
469            loader_snapshots: HashMap::new(),
470        }
471    }
472
473    /// Register a file with its current timestamp.
474    pub fn watch(&mut self, path: impl Into<String>, timestamp: u64) {
475        let p = path.into();
476        self.timestamps.insert(p, timestamp);
477    }
478
479    /// Update the watcher's snapshot of a file's source content.
480    pub fn snapshot_source(&mut self, path: impl Into<String>, source: impl Into<String>) {
481        let p = path.into();
482        let s = source.into();
483        // Use a simple hash as a proxy timestamp
484        let ts = fnv1a(s.as_bytes());
485        self.timestamps.insert(p.clone(), ts);
486        self.loader_snapshots.insert(p, s);
487    }
488
489    /// Set timestamp for a path (simulates an mtime update).
490    pub fn set_timestamp(&mut self, path: &str, ts: u64) {
491        self.timestamps.insert(path.to_string(), ts);
492    }
493
494    /// Check which files have changed since last snapshot.
495    /// Returns paths whose current timestamp differs from registered.
496    pub fn check_changes(&self, current_timestamps: &HashMap<String, u64>) -> Vec<String> {
497        let mut changed = Vec::new();
498        for (path, &last_ts) in &self.timestamps {
499            if let Some(&cur_ts) = current_timestamps.get(path) {
500                if cur_ts != last_ts {
501                    changed.push(path.clone());
502                }
503            }
504        }
505        changed
506    }
507
508    /// Reload any changed modules detected by comparing the given current timestamps.
509    pub fn reload_changed(
510        &mut self,
511        current_timestamps: &HashMap<String, u64>,
512        registry: &mut ModuleRegistry,
513        vm: &mut Vm,
514    ) -> Vec<(String, Result<(), String>)> {
515        let changed = self.check_changes(current_timestamps);
516        let mut results = Vec::new();
517        for path in &changed {
518            let res = registry.reload(path, vm)
519                .map(|_| ())
520                .map_err(|e| e.message);
521            if res.is_ok() {
522                // Update timestamp
523                if let Some(&ts) = current_timestamps.get(path) {
524                    self.timestamps.insert(path.clone(), ts);
525                }
526            }
527            results.push((path.clone(), res));
528        }
529        results
530    }
531
532    pub fn watched_paths(&self) -> Vec<String> {
533        self.timestamps.keys().cloned().collect()
534    }
535}
536
537impl Default for HotReloadWatcher {
538    fn default() -> Self { Self::new() }
539}
540
541// ── Tests ─────────────────────────────────────────────────────────────────────
542
543#[cfg(test)]
544mod tests {
545    use super::*;
546    use crate::scripting::stdlib::register_all;
547    use crate::scripting::vm::Vm;
548
549    fn make_vm() -> Vm {
550        let mut vm = Vm::new();
551        register_all(&mut vm);
552        vm
553    }
554
555    fn string_registry(entries: &[(&str, &str)]) -> ModuleRegistry {
556        let mut loader = StringMapLoader::new();
557        for (k, v) in entries {
558            loader.add(*k, *v);
559        }
560        ModuleRegistry::new(Box::new(loader))
561    }
562
563    #[test]
564    fn test_module_id_stable() {
565        let a = ModuleId::from_path("math.utils");
566        let b = ModuleId::from_path("math.utils");
567        assert_eq!(a, b);
568    }
569
570    #[test]
571    fn test_module_id_different() {
572        let a = ModuleId::from_path("a");
573        let b = ModuleId::from_path("b");
574        assert_ne!(a, b);
575    }
576
577    #[test]
578    fn test_string_map_loader() {
579        let mut loader = StringMapLoader::new();
580        loader.add("foo", "return 42");
581        assert_eq!(loader.load_source("foo").unwrap(), "return 42");
582        assert!(loader.load_source("bar").is_err());
583    }
584
585    #[test]
586    fn test_registry_require_simple() {
587        let mut vm  = make_vm();
588        let mut reg = string_registry(&[("mod", "return 99")]);
589        let v = reg.require("mod", &mut vm).unwrap();
590        assert_eq!(v, Value::Int(99));
591    }
592
593    #[test]
594    fn test_registry_require_cached() {
595        let mut vm  = make_vm();
596        let mut reg = string_registry(&[("mod", "return {x=1}")]);
597        let _  = reg.require("mod", &mut vm).unwrap();
598        let v2 = reg.require("mod", &mut vm).unwrap();
599        assert!(matches!(v2, Value::Table(_)));
600    }
601
602    #[test]
603    fn test_registry_unload() {
604        let mut vm  = make_vm();
605        let mut reg = string_registry(&[("mod", "return 1")]);
606        reg.require("mod", &mut vm).unwrap();
607        reg.unload("mod");
608        assert!(reg.get_module("mod").is_none());
609    }
610
611    #[test]
612    fn test_registry_loaded_modules() {
613        let mut vm  = make_vm();
614        let mut reg = string_registry(&[("a", "return 1"), ("b", "return 2")]);
615        reg.require("a", &mut vm).unwrap();
616        reg.require("b", &mut vm).unwrap();
617        let mods = reg.loaded_modules();
618        assert_eq!(mods.len(), 2);
619    }
620
621    #[test]
622    fn test_namespace_set_get() {
623        let mut ns = Namespace::new("root");
624        ns.set("x", Value::Int(10));
625        ns.set("math.pi", Value::Float(3.14));
626        assert_eq!(ns.get("x"), Some(&Value::Int(10)));
627        assert!(ns.get("math.pi").is_some());
628    }
629
630    #[test]
631    fn test_namespace_export_table() {
632        let mut ns = Namespace::new("root");
633        ns.set("a", Value::Int(1));
634        ns.set("b", Value::Int(2));
635        let t = ns.export_table();
636        if let Value::Table(tbl) = &t {
637            assert_eq!(tbl.rawget_str("a"), Value::Int(1));
638        } else {
639            panic!("expected table");
640        }
641    }
642
643    #[test]
644    fn test_namespace_merge() {
645        let mut a = Namespace::new("a");
646        a.set("x", Value::Int(1));
647        let mut b = Namespace::new("b");
648        b.set("x", Value::Int(99));
649        b.set("y", Value::Int(2));
650        a.merge_namespaces(&b);
651        assert_eq!(a.get("x"), Some(&Value::Int(99)));
652        assert_eq!(a.get("y"), Some(&Value::Int(2)));
653    }
654
655    #[test]
656    fn test_hot_reload_detect_change() {
657        let mut watcher = HotReloadWatcher::new();
658        watcher.watch("file.lua", 100);
659        let mut current = HashMap::new();
660        current.insert("file.lua".to_string(), 100u64);
661        assert!(watcher.check_changes(&current).is_empty());
662        current.insert("file.lua".to_string(), 200u64);
663        let changed = watcher.check_changes(&current);
664        assert_eq!(changed, vec!["file.lua".to_string()]);
665    }
666
667    #[test]
668    fn test_hot_reload_no_change() {
669        let mut watcher = HotReloadWatcher::new();
670        watcher.watch("a.lua", 42);
671        watcher.watch("b.lua", 43);
672        let mut current = HashMap::new();
673        current.insert("a.lua".to_string(), 42u64);
674        current.insert("b.lua".to_string(), 43u64);
675        assert!(watcher.check_changes(&current).is_empty());
676    }
677
678    #[test]
679    fn test_fnv1a_hash() {
680        // Verify determinism
681        assert_eq!(fnv1a(b"hello"), fnv1a(b"hello"));
682        assert_ne!(fnv1a(b"hello"), fnv1a(b"world"));
683    }
684}