1use dashmap::DashMap;
7use shape_ast::ast::{Program, Span};
8#[cfg(test)]
9use shape_ast::parser::parse_program;
10use shape_runtime::extension_context::ExtensionModuleSchemaCache;
11use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use std::sync::{Arc, OnceLock};
14
15pub(crate) fn lsp_extension_schema_cache() -> &'static Arc<ExtensionModuleSchemaCache> {
24 static CACHE: OnceLock<Arc<ExtensionModuleSchemaCache>> = OnceLock::new();
25 CACHE.get_or_init(|| Arc::new(ExtensionModuleSchemaCache::new()))
26}
27
28fn module_path_segments(path: &str) -> Vec<&str> {
29 if path.contains("::") {
30 path.split("::")
31 .filter(|segment| !segment.is_empty())
32 .collect()
33 } else {
34 path.split('.')
35 .filter(|segment| !segment.is_empty())
36 .collect()
37 }
38}
39
40fn is_std_module_path(path: &str) -> bool {
41 module_path_segments(path)
42 .first()
43 .is_some_and(|segment| *segment == "std")
44}
45
46#[derive(Debug, Clone, PartialEq, Eq)]
48pub enum SymbolKind {
49 Function,
50 Pattern,
51 Variable,
52 TypeAlias,
53 Trait,
54 Enum,
55 Annotation,
56}
57
58#[derive(Debug, Clone)]
60pub struct ExportedSymbol {
61 pub name: String,
63 pub alias: Option<String>,
65 pub kind: SymbolKind,
67 pub span: Span,
69}
70
71impl ExportedSymbol {
72 pub fn exported_name(&self) -> &str {
74 self.alias.as_ref().unwrap_or(&self.name)
75 }
76}
77
78#[derive(Debug, Clone)]
80pub struct ModuleInfo {
81 pub path: PathBuf,
83 pub program: Arc<Program>,
85 pub exports: Vec<ExportedSymbol>,
87}
88
89#[derive(Debug, Default)]
91pub struct ModuleCache {
92 modules: DashMap<PathBuf, ModuleInfo>,
94}
95
96impl ModuleCache {
97 pub fn new() -> Self {
99 Self {
100 modules: DashMap::new(),
101 }
102 }
103
104 fn loader_for_context(
105 current_file: &Path,
106 workspace_root: Option<&Path>,
107 current_source: Option<&str>,
108 ) -> shape_runtime::module_loader::ModuleLoader {
109 let mut loader = shape_runtime::module_loader::ModuleLoader::new();
110 loader.configure_for_context_with_source(
111 current_file,
112 workspace_root,
113 current_source,
114 lsp_extension_schema_cache(),
115 );
116 loader
117 }
118
119 pub fn resolve_import(
121 &self,
122 import_path: &str,
123 current_file: &Path,
124 workspace_root: Option<&Path>,
125 ) -> Option<PathBuf> {
126 let loader = Self::loader_for_context(current_file, workspace_root, None);
127
128 let context_dir = current_file.parent().map(Path::to_path_buf);
129 let resolved = loader.resolve_module_path_with_context(import_path, context_dir.as_ref());
130 if let Ok(path) = resolved {
131 return Some(path);
132 }
133
134 if import_path.contains("::")
136 || import_path.starts_with("./")
137 || import_path.starts_with("../")
138 || import_path.starts_with('/')
139 {
140 return None;
141 }
142
143 let canonical = import_path.replace('.', "::");
144 loader
145 .resolve_module_path_with_context(&canonical, context_dir.as_ref())
146 .ok()
147 }
148
149 pub fn load_module(&self, path: &Path) -> Option<ModuleInfo> {
154 self.load_module_with_context(path, path, None)
155 }
156
157 pub fn load_module_with_context(
159 &self,
160 path: &Path,
161 current_file: &Path,
162 workspace_root: Option<&Path>,
163 ) -> Option<ModuleInfo> {
164 if let Some(cached) = self.modules.get(path) {
166 return Some(cached.clone());
167 }
168
169 let mut loader = Self::loader_for_context(current_file, workspace_root, None);
171 let module = loader.load_module_from_file(path).ok()?;
172 let program = Arc::new(module.ast.clone());
173
174 let exports = extract_exports(&program);
176
177 let module_info = ModuleInfo {
178 path: path.to_path_buf(),
179 program: program.clone(),
180 exports,
181 };
182
183 self.modules.insert(path.to_path_buf(), module_info.clone());
185
186 Some(module_info)
187 }
188
189 pub fn load_module_by_import_with_context_and_source(
193 &self,
194 import_path: &str,
195 current_file: &Path,
196 workspace_root: Option<&Path>,
197 current_source: Option<&str>,
198 ) -> Option<ModuleInfo> {
199 let mut loader = Self::loader_for_context(current_file, workspace_root, current_source);
200 let context_dir = current_file.parent().map(Path::to_path_buf);
201 let module = loader
202 .load_module_with_context(import_path, context_dir.as_ref())
203 .ok()?;
204
205 let cache_path = PathBuf::from(format!(
206 "__shape_lsp_virtual__/{}.shape",
207 import_path.replace("::", "/").replace('.', "/")
208 ));
209 let program = Arc::new(module.ast.clone());
210 let exports = extract_exports(&program);
211 let module_info = ModuleInfo {
212 path: cache_path.clone(),
213 program: program.clone(),
214 exports,
215 };
216 self.modules.insert(cache_path, module_info.clone());
217 Some(module_info)
218 }
219
220 pub fn get_module(&self, path: &Path) -> Option<ModuleInfo> {
222 self.modules.get(path).map(|entry| entry.clone())
223 }
224
225 pub fn invalidate(&self, path: &Path) {
227 self.modules.remove(path);
228 }
229
230 pub fn clear(&self) {
232 self.modules.clear();
233 }
234
235 pub fn list_importable_modules_with_context(
242 &self,
243 current_file: &Path,
244 workspace_root: Option<&Path>,
245 ) -> Vec<String> {
246 self.list_importable_modules_with_context_and_source(current_file, workspace_root, None)
247 }
248
249 pub fn list_importable_modules_with_context_and_source(
251 &self,
252 current_file: &Path,
253 workspace_root: Option<&Path>,
254 current_source: Option<&str>,
255 ) -> Vec<String> {
256 let mut loader = shape_runtime::module_loader::ModuleLoader::new();
257 loader.configure_for_context_with_source(
258 current_file,
259 workspace_root,
260 current_source,
261 lsp_extension_schema_cache(),
262 );
263 loader.list_importable_modules_with_context(current_file, workspace_root)
264 }
265
266 pub fn list_importable_modules(&self) -> Vec<String> {
268 let current_file = std::env::current_dir()
269 .unwrap_or_else(|_| PathBuf::from("."))
270 .join("__shape_lsp__.shape");
271 self.list_importable_modules_with_context(¤t_file, None)
272 }
273
274 pub fn list_stdlib_modules(&self) -> Vec<String> {
276 self.list_importable_modules()
277 .into_iter()
278 .filter(|module_path| is_std_module_path(module_path))
279 .collect()
280 }
281
282 pub fn list_stdlib_children(&self, prefix: &str) -> Vec<ModuleChild> {
288 let effective_prefix = if prefix.is_empty() { "std" } else { prefix };
289 if !is_std_module_path(effective_prefix) {
290 return Vec::new();
291 }
292
293 self.list_module_children(effective_prefix)
294 }
295
296 pub fn list_module_children_with_context(
298 &self,
299 prefix: &str,
300 current_file: &Path,
301 workspace_root: Option<&Path>,
302 ) -> Vec<ModuleChild> {
303 let base = if prefix.is_empty() {
304 "std".to_string()
305 } else {
306 prefix.to_string()
307 };
308
309 let mut children: HashMap<String, ModuleChild> = HashMap::new();
310 let base_segments = module_path_segments(&base);
311 let base_len = base_segments.len();
312 for module_path in self.list_importable_modules_with_context(current_file, workspace_root) {
313 let module_segments = module_path_segments(&module_path);
314 if module_segments.len() <= base_len {
315 continue;
316 }
317 if module_segments[..base_len] != base_segments[..] {
318 continue;
319 }
320
321 let child = module_segments[base_len];
322 let has_children = module_segments.len() > base_len + 1;
323
324 let entry = children.entry(child.to_string()).or_insert(ModuleChild {
325 name: child.to_string(),
326 has_leaf_module: false,
327 has_children: false,
328 });
329 if has_children {
330 entry.has_children = true;
331 } else {
332 entry.has_leaf_module = true;
333 }
334 }
335
336 let mut out: Vec<ModuleChild> = children.into_values().collect();
337 out.sort_by(|a, b| a.name.cmp(&b.name));
338 out
339 }
340
341 pub fn list_module_children(&self, prefix: &str) -> Vec<ModuleChild> {
343 let current_file = std::env::current_dir()
344 .unwrap_or_else(|_| PathBuf::from("."))
345 .join("__shape_lsp__.shape");
346 self.list_module_children_with_context(prefix, ¤t_file, None)
347 }
348
349 pub fn find_exported_symbol_with_context(
352 &self,
353 name: &str,
354 current_file: &Path,
355 workspace_root: Option<&Path>,
356 ) -> Vec<(String, ExportedSymbol)> {
357 let mut results = Vec::new();
358
359 for import_path in self.list_importable_modules_with_context(current_file, workspace_root) {
360 let Some(resolved) = self.resolve_import(&import_path, current_file, workspace_root)
361 else {
362 continue;
363 };
364 let Some(module_info) =
365 self.load_module_with_context(&resolved, current_file, workspace_root)
366 else {
367 continue;
368 };
369
370 for export in &module_info.exports {
371 if export.exported_name() == name {
372 results.push((import_path.clone(), export.clone()));
373 }
374 }
375 }
376
377 results
378 }
379
380 pub fn find_exported_symbol(&self, name: &str) -> Vec<(String, ExportedSymbol)> {
382 let current_file = std::env::current_dir()
383 .unwrap_or_else(|_| PathBuf::from("."))
384 .join("__shape_lsp__.shape");
385 self.find_exported_symbol_with_context(name, ¤t_file, None)
386 }
387
388 pub fn enumerate_workspace_shape_files(&self, workspace_root: &Path) -> Vec<PathBuf> {
405 let mut files = Vec::new();
406 let mut stack = vec![workspace_root.to_path_buf()];
407 while let Some(dir) = stack.pop() {
408 if files.len() >= MAX_WORKSPACE_FILES {
409 break;
410 }
411 let Ok(entries) = std::fs::read_dir(&dir) else {
412 continue;
413 };
414 for entry in entries.flatten() {
415 let path = entry.path();
416 let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
417 continue;
418 };
419 if name.starts_with('.')
421 || matches!(name, "target" | "node_modules" | "dist" | "build")
422 {
423 continue;
424 }
425 if path.is_dir() {
426 stack.push(path);
427 } else if path.extension().and_then(|e| e.to_str()) == Some("shape") {
428 files.push(path);
429 if files.len() >= MAX_WORKSPACE_FILES {
430 break;
431 }
432 }
433 }
434 }
435 files
436 }
437}
438
439pub const MAX_WORKSPACE_FILES: usize = 4096;
445
446pub fn path_to_module_path(file: &Path, workspace_root: &Path) -> Option<String> {
462 let rel = file.strip_prefix(workspace_root).ok()?;
463 let stem = rel.file_stem().and_then(|s| s.to_str())?;
464 if rel.extension().and_then(|e| e.to_str()) != Some("shape") {
465 return None;
466 }
467 let mut segments: Vec<String> = Vec::new();
468 for component in rel.parent()?.components() {
469 let std::path::Component::Normal(seg) = component else {
470 continue;
471 };
472 let seg = seg.to_str()?;
473 if !is_valid_shape_identifier(seg) {
474 return None;
475 }
476 segments.push(seg.to_string());
477 }
478 if !is_valid_shape_identifier(stem) {
479 return None;
480 }
481 segments.push(stem.to_string());
482 if segments.is_empty() {
483 return None;
484 }
485 Some(segments.join("::"))
486}
487
488fn is_valid_shape_identifier(s: &str) -> bool {
489 let mut chars = s.chars();
490 let Some(first) = chars.next() else {
491 return false;
492 };
493 if !(first.is_ascii_alphabetic() || first == '_') {
494 return false;
495 }
496 chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
497}
498
499#[derive(Debug, Clone)]
501pub struct ModuleChild {
502 pub name: String,
503 pub has_leaf_module: bool,
504 pub has_children: bool,
505}
506
507fn map_module_export_kind(kind: shape_runtime::module_loader::ModuleExportKind) -> SymbolKind {
508 use shape_runtime::module_loader::ModuleExportKind as RuntimeKind;
509 match kind {
510 RuntimeKind::Function => SymbolKind::Function,
511 RuntimeKind::BuiltinFunction => SymbolKind::Function,
512 RuntimeKind::TypeAlias => SymbolKind::TypeAlias,
513 RuntimeKind::BuiltinType => SymbolKind::TypeAlias,
514 RuntimeKind::Trait => SymbolKind::Trait,
515 RuntimeKind::Enum => SymbolKind::Enum,
516 RuntimeKind::Annotation => SymbolKind::Annotation,
517 RuntimeKind::Value => SymbolKind::Variable,
518 }
519}
520
521fn extract_exports(program: &Program) -> Vec<ExportedSymbol> {
523 shape_runtime::module_loader::collect_exported_symbols(program)
524 .unwrap_or_default()
525 .into_iter()
526 .map(|sym| ExportedSymbol {
527 name: sym.name,
528 alias: sym.alias,
529 kind: map_module_export_kind(sym.kind),
530 span: sym.span,
531 })
532 .collect()
533}
534
535#[cfg(test)]
536mod tests {
537 use super::*;
538
539 #[test]
540 fn test_resolve_stdlib_import() {
541 let cache = ModuleCache::new();
542 let current_file =
543 PathBuf::from("/home/dev/dev/finance/analysis-suite/shape/examples/test.shape");
544
545 let resolved = cache.resolve_import("std::core::math", ¤t_file, None);
546
547 assert!(resolved.is_some());
548 let path = resolved.unwrap();
549 let path_str = path.to_string_lossy();
550 assert!(
551 path_str.contains("stdlib/core/math.shape")
552 || path_str.contains("stdlib-src/core/math.shape"),
553 "Expected stdlib math path, got: {}",
554 path_str
555 );
556 }
557
558 #[test]
559 fn test_relative_import_is_supported() {
560 let tmp = tempfile::tempdir().unwrap();
561 let current_file = tmp.path().join("main.shape");
562 let util = tmp.path().join("utils.shape");
563 std::fs::write(¤t_file, "from ./utils use { helper }").unwrap();
564 std::fs::write(&util, "pub fn helper() { 1 }").unwrap();
565
566 let cache = ModuleCache::new();
567
568 let resolved = cache.resolve_import("./utils", ¤t_file, None);
569 assert_eq!(resolved.as_deref(), Some(util.as_path()));
570 }
571
572 #[test]
573 fn test_non_std_import_returns_none() {
574 let cache = ModuleCache::new();
575 let current_file = PathBuf::from("/home/user/project/src/main.shape");
576
577 let resolved = cache.resolve_import("finance::indicators", ¤t_file, None);
579 assert!(resolved.is_none());
580 }
581
582 #[test]
583 fn test_exported_symbol_name() {
584 let symbol = ExportedSymbol {
585 name: "originalName".to_string(),
586 alias: Some("aliasName".to_string()),
587 kind: SymbolKind::Function,
588 span: Span::default(),
589 };
590
591 assert_eq!(symbol.exported_name(), "aliasName");
592
593 let symbol_no_alias = ExportedSymbol {
594 name: "originalName".to_string(),
595 alias: None,
596 kind: SymbolKind::Function,
597 span: Span::default(),
598 };
599
600 assert_eq!(symbol_no_alias.exported_name(), "originalName");
601 }
602
603 #[test]
604 fn test_extract_exports() {
605 let source = r#"
606pub fn myFunc(x) {
607 return x + 1;
608}
609
610fn localFunc() {
611 return 42;
612}
613"#;
614
615 let program = parse_program(source).unwrap();
616 let exports = extract_exports(&program);
617
618 assert_eq!(exports.len(), 1);
619 assert_eq!(exports[0].name, "myFunc");
620 assert_eq!(exports[0].kind, SymbolKind::Function);
621 }
622
623 #[test]
624 fn test_list_stdlib_modules_not_empty() {
625 let cache = ModuleCache::new();
626 let modules = cache.list_stdlib_modules();
627 assert!(
628 !modules.is_empty(),
629 "expected stdlib module list to be non-empty"
630 );
631 assert!(
632 modules.iter().all(|m| m.starts_with("std::")),
633 "all stdlib modules should be std::-prefixed: {:?}",
634 modules
635 );
636 }
637
638 #[test]
639 fn test_list_stdlib_children_for_std_prefix() {
640 let cache = ModuleCache::new();
641 let children = cache.list_stdlib_children("std");
642 assert!(
643 !children.is_empty(),
644 "expected stdlib root to have child modules"
645 );
646 assert!(
647 children.iter().any(|c| c.name == "core"),
648 "expected std.core child in stdlib tree"
649 );
650 }
651
652 #[test]
653 fn test_list_importable_modules_with_project_modules_and_deps() {
654 let tmp = tempfile::tempdir().unwrap();
655 let root = tmp.path();
656 std::fs::write(
657 root.join("shape.toml"),
658 r#"
659[modules]
660paths = ["lib"]
661
662[dependencies]
663mydep = { path = "deps/mydep" }
664"#,
665 )
666 .unwrap();
667
668 std::fs::create_dir_all(root.join("src")).unwrap();
669 std::fs::create_dir_all(root.join("lib")).unwrap();
670 std::fs::create_dir_all(root.join("deps/mydep")).unwrap();
671
672 std::fs::write(root.join("src/main.shape"), "let x = 1").unwrap();
673 std::fs::write(root.join("lib/tools.shape"), "pub fn tool() { 1 }").unwrap();
674 std::fs::write(root.join("deps/mydep/index.shape"), "pub fn root() { 1 }").unwrap();
675 std::fs::write(root.join("deps/mydep/util.shape"), "pub fn util() { 1 }").unwrap();
676
677 let cache = ModuleCache::new();
678 let modules =
679 cache.list_importable_modules_with_context(&root.join("src/main.shape"), None);
680
681 assert!(
682 modules.iter().any(|m| m == "tools"),
683 "expected module path from [modules].paths, got: {:?}",
684 modules
685 );
686 assert!(
687 modules.iter().any(|m| m == "mydep"),
688 "expected dependency index module path, got: {:?}",
689 modules
690 );
691 assert!(
692 modules.iter().any(|m| m == "mydep::util"),
693 "expected dependency submodule path, got: {:?}",
694 modules
695 );
696 }
697
698 #[test]
699 fn test_enumerate_workspace_shape_files() {
700 let tmp = tempfile::tempdir().unwrap();
703 let root = tmp.path();
704 std::fs::write(root.join("a.shape"), "let x = 1").unwrap();
705 std::fs::create_dir_all(root.join("sub")).unwrap();
706 std::fs::write(root.join("sub/b.shape"), "fn foo() { 2 }").unwrap();
707 std::fs::create_dir_all(root.join(".git")).unwrap();
709 std::fs::write(root.join(".git/c.shape"), "ignored").unwrap();
710 std::fs::create_dir_all(root.join("target")).unwrap();
712 std::fs::write(root.join("target/d.shape"), "ignored").unwrap();
713 std::fs::write(root.join("e.txt"), "ignored").unwrap();
715
716 let cache = ModuleCache::new();
717 let files = cache.enumerate_workspace_shape_files(root);
718 let names: Vec<String> = files
719 .iter()
720 .filter_map(|p| p.file_name().and_then(|n| n.to_str()).map(String::from))
721 .collect();
722 assert!(names.contains(&"a.shape".to_string()), "expected a.shape, got {:?}", names);
723 assert!(names.contains(&"b.shape".to_string()), "expected b.shape, got {:?}", names);
724 assert!(
725 !names.contains(&"c.shape".to_string()),
726 "hidden dir .git must be excluded, got {:?}",
727 names
728 );
729 assert!(
730 !names.contains(&"d.shape".to_string()),
731 "build dir target must be excluded, got {:?}",
732 names
733 );
734 assert!(
735 !names.contains(&"e.txt".to_string()),
736 "non-.shape file must be excluded, got {:?}",
737 names
738 );
739 }
740
741 #[test]
742 fn test_path_to_module_path_w27() {
743 let root = PathBuf::from("/ws");
745 assert_eq!(
746 path_to_module_path(&root.join("foo.shape"), &root).as_deref(),
747 Some("foo")
748 );
749 assert_eq!(
750 path_to_module_path(&root.join("lib/util.shape"), &root).as_deref(),
751 Some("lib::util")
752 );
753 assert_eq!(
754 path_to_module_path(&root.join("src/sub/mod.shape"), &root).as_deref(),
755 Some("src::sub::mod")
756 );
757 assert!(path_to_module_path(&root.join("foo.txt"), &root).is_none());
759 assert!(path_to_module_path(&PathBuf::from("/other/x.shape"), &root).is_none());
761 assert!(path_to_module_path(&root.join("2bad.shape"), &root).is_none());
763 assert!(path_to_module_path(&root.join("a-b/x.shape"), &root).is_none());
764 assert_eq!(
766 path_to_module_path(&root.join("_priv/_x.shape"), &root).as_deref(),
767 Some("_priv::_x")
768 );
769 }
770
771 #[test]
772 fn test_module_cache_invalidation() {
773 let cache = ModuleCache::new();
774 let path = PathBuf::from("/test/module.shape");
775
776 let program = Arc::new(Program {
778 items: vec![],
779 docs: shape_ast::ast::ProgramDocs::default(),
780 });
781 let module_info = ModuleInfo {
782 path: path.clone(),
783 program,
784 exports: vec![],
785 };
786
787 cache.modules.insert(path.clone(), module_info.clone());
789 assert!(cache.get_module(&path).is_some());
790
791 cache.invalidate(&path);
793 assert!(cache.get_module(&path).is_none());
794 }
795}