Skip to main content

assay/
discovery.rs

1//! Module discovery and search index builder.
2//!
3//! Discovers Assay modules from three sources (in priority order):
4//! 1. Project — `./modules/` relative to CWD
5//! 2. Global  — `$ASSAY_MODULES_PATH` or `~/.assay/modules/`
6//! 3. BuiltIn — embedded stdlib + hardcoded Rust builtins
7
8use include_dir::{include_dir, Dir};
9use crate::search::{SearchEngine, SearchResult};
10
11use crate::metadata::{self, ModuleMetadata};
12#[cfg(not(feature = "db"))]
13use crate::search::BM25Index;
14#[cfg(feature = "db")]
15use crate::search_fts5::FTS5Index;
16
17static STDLIB_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/stdlib");
18
19/// Where a discovered module originates from.
20#[derive(Debug, Clone, PartialEq)]
21pub enum ModuleSource {
22    /// Embedded in the binary via `include_dir!`
23    BuiltIn,
24    /// Found in `./modules/` relative to CWD
25    Project,
26    /// Found in `$ASSAY_MODULES_PATH` or `~/.assay/modules/`
27    Global,
28}
29
30/// A module discovered during the discovery phase.
31#[derive(Debug, Clone)]
32pub struct DiscoveredModule {
33    pub module_name: String,
34    pub source: ModuleSource,
35    pub metadata: ModuleMetadata,
36    pub lua_source: String,
37}
38
39/// Hardcoded Rust builtins with their descriptions and search keywords.
40const BUILTINS: &[(&str, &str, &[&str])] = &[
41    (
42        "http",
43        "HTTP client and server: get, post, put, patch, delete, serve",
44        &["http", "client", "server", "request", "response", "headers", "endpoint", "api", "webhook", "rest"],
45    ),
46    (
47        "json",
48        "JSON serialization: parse and encode",
49        &["json", "serialization", "deserialize", "stringify", "parse", "encode", "format"],
50    ),
51    (
52        "yaml",
53        "YAML serialization: parse and encode",
54        &["yaml", "serialization", "deserialize", "parse", "encode", "format"],
55    ),
56    (
57        "toml",
58        "TOML serialization: parse and encode",
59        &["toml", "serialization", "deserialize", "parse", "encode", "configuration"],
60    ),
61    (
62        "fs",
63        "Filesystem: read and write files",
64        &["fs", "filesystem", "file", "read", "write", "io", "path"],
65    ),
66    (
67        "crypto",
68        "Cryptography: jwt_sign, hash, hmac, random",
69        &["crypto", "jwt", "signature", "hash", "hmac", "encryption", "random", "security", "password", "signing", "rsa", "sha256"],
70    ),
71    (
72        "base64",
73        "Base64 encoding and decoding",
74        &["base64", "encoding", "decode", "encode", "binary"],
75    ),
76    (
77        "regex",
78        "Regular expressions: match, find, find_all, replace",
79        &["regex", "pattern", "match", "find", "replace", "regular-expression", "regexp"],
80    ),
81    (
82        "db",
83        "Database: connect, query, execute, close (Postgres, MySQL, SQLite)",
84        &["db", "database", "sql", "postgres", "mysql", "sqlite", "connection", "query", "execute"],
85    ),
86    (
87        "ws",
88        "WebSocket: connect, send, recv, close",
89        &["ws", "websocket", "connection", "message", "streaming", "realtime", "socket"],
90    ),
91    (
92        "template",
93        "Jinja2-compatible templates: render file or string",
94        &["template", "jinja2", "rendering", "string-template", "mustache", "render"],
95    ),
96    (
97        "async",
98        "Async tasks: spawn, spawn_interval, await, cancel",
99        &["async", "asynchronous", "task", "coroutine", "concurrent", "spawn", "interval"],
100    ),
101    (
102        "assert",
103        "Assertions: eq, gt, lt, contains, not_nil, matches",
104        &["assert", "assertion", "test", "validation", "comparison", "check", "verify"],
105    ),
106    (
107        "log",
108        "Logging: info, warn, error",
109        &["log", "logging", "output", "debug", "error", "warning", "info", "trace"],
110    ),
111    (
112        "env",
113        "Environment variables: get",
114        &["env", "environment", "variable", "configuration", "config"],
115    ),
116    (
117        "sleep",
118        "Sleep for N seconds",
119        &["sleep", "delay", "pause", "wait", "time"],
120    ),
121    (
122        "time",
123        "Unix timestamp in seconds",
124        &["time", "timestamp", "unix", "epoch", "clock", "datetime"],
125    ),
126];
127
128/// Discover all modules: embedded stdlib + `./modules/` + `~/.assay/modules/` (or `$ASSAY_MODULES_PATH`).
129///
130/// Returns modules ordered by priority: Project first, then Global, then BuiltIn.
131/// Callers can deduplicate by name, keeping the highest-priority (first) occurrence.
132pub fn discover_modules() -> Vec<DiscoveredModule> {
133    let mut modules = Vec::new();
134
135    // Priority 1: Project modules (./modules/)
136    discover_filesystem_modules(
137        std::path::Path::new("./modules"),
138        ModuleSource::Project,
139        &mut modules,
140    );
141
142    // Priority 2: Global modules ($ASSAY_MODULES_PATH or ~/.assay/modules/)
143    let global_path = resolve_global_modules_path();
144    if let Some(path) = global_path {
145        discover_filesystem_modules(&path, ModuleSource::Global, &mut modules);
146    }
147
148    // Priority 3: Embedded stdlib .lua files
149    discover_embedded_stdlib(&mut modules);
150
151    // Priority 3 (continued): Hardcoded Rust builtins
152    discover_rust_builtins(&mut modules);
153
154    modules
155}
156
157/// Build a search index from discovered modules.
158///
159/// When feature `db` is enabled: uses `FTS5Index`.
160/// When feature `db` is disabled: uses `BM25Index`.
161pub fn build_index(modules: &[DiscoveredModule]) -> Box<dyn SearchEngine> {
162    #[cfg(feature = "db")]
163    {
164        let mut idx = FTS5Index::new();
165        for m in modules {
166            idx.add_document(
167                &m.module_name,
168                &[
169                    ("keywords", &m.metadata.keywords.join(" "), 3.0),
170                    ("module_name", &m.module_name, 2.0),
171                    ("description", &m.metadata.description, 1.0),
172                    ("functions", &m.metadata.auto_functions.join(" "), 1.0),
173                ],
174            );
175        }
176        Box::new(idx)
177    }
178    #[cfg(not(feature = "db"))]
179    {
180        let mut idx = BM25Index::new();
181        for m in modules {
182            idx.add_document(
183                &m.module_name,
184                &[
185                    ("keywords", &m.metadata.keywords.join(" "), 3.0),
186                    ("module_name", &m.module_name, 2.0),
187                    ("description", &m.metadata.description, 1.0),
188                    ("functions", &m.metadata.auto_functions.join(" "), 1.0),
189                ],
190            );
191        }
192        Box::new(idx)
193    }
194}
195
196/// Convenience: discover all modules, build index, search, return results.
197pub fn search_modules(query: &str, limit: usize) -> Vec<SearchResult> {
198    let modules = discover_modules();
199    let index = build_index(&modules);
200    index.search(query, limit)
201}
202
203/// Resolve the global modules directory path.
204///
205/// Checks `$ASSAY_MODULES_PATH` first, then falls back to `~/.assay/modules/`.
206/// Returns `None` if neither is available.
207fn resolve_global_modules_path() -> Option<std::path::PathBuf> {
208    if let Ok(custom) = std::env::var(crate::lua::MODULES_PATH_ENV) {
209        return Some(std::path::PathBuf::from(custom));
210    }
211    if let Ok(home) = std::env::var("HOME") {
212        return Some(std::path::Path::new(&home).join(".assay/modules"));
213    }
214    None
215}
216
217/// Discover `.lua` files from a filesystem directory.
218///
219/// Silently skips if the directory does not exist.
220fn discover_filesystem_modules(
221    dir: &std::path::Path,
222    source: ModuleSource,
223    modules: &mut Vec<DiscoveredModule>,
224) {
225    let entries = match std::fs::read_dir(dir) {
226        Ok(entries) => entries,
227        Err(_) => return, // Directory doesn't exist or can't be read — skip silently
228    };
229
230    for entry in entries.flatten() {
231        let path = entry.path();
232        if path.extension().and_then(|e| e.to_str()) != Some("lua") {
233            continue;
234        }
235
236        let lua_source = match std::fs::read_to_string(&path) {
237            Ok(s) => s,
238            Err(_) => continue,
239        };
240
241        let stem = path
242            .file_stem()
243            .and_then(|s| s.to_str())
244            .unwrap_or_default();
245        let module_name = format!("assay.{stem}");
246        let meta = metadata::parse_metadata(&lua_source);
247
248        modules.push(DiscoveredModule {
249            module_name,
250            source: source.clone(),
251            metadata: meta,
252            lua_source,
253        });
254    }
255}
256
257/// Discover embedded stdlib `.lua` files from `include_dir!`.
258fn discover_embedded_stdlib(modules: &mut Vec<DiscoveredModule>) {
259    for file in STDLIB_DIR.files() {
260        let path = file.path();
261        if path.extension().and_then(|e| e.to_str()) != Some("lua") {
262            continue;
263        }
264
265        let lua_source = match file.contents_utf8() {
266            Some(s) => s,
267            None => continue,
268        };
269
270        let stem = path
271            .file_stem()
272            .and_then(|s| s.to_str())
273            .unwrap_or_default();
274        let module_name = format!("assay.{stem}");
275        let meta = metadata::parse_metadata(lua_source);
276
277        modules.push(DiscoveredModule {
278            module_name,
279            source: ModuleSource::BuiltIn,
280            metadata: meta,
281            lua_source: lua_source.to_string(),
282        });
283    }
284}
285
286/// Add hardcoded Rust builtins (not Lua files) to the module list.
287fn discover_rust_builtins(modules: &mut Vec<DiscoveredModule>) {
288    for &(name, description, kw) in BUILTINS {
289        modules.push(DiscoveredModule {
290            module_name: name.to_string(),
291            source: ModuleSource::BuiltIn,
292            lua_source: String::new(),
293            metadata: ModuleMetadata {
294                module_name: name.to_string(),
295                description: description.to_string(),
296                keywords: kw.iter().map(|k| k.to_string()).collect(),
297                ..Default::default()
298            },
299        });
300    }
301}