Add symbols imported in packages to workspace

Author: lionel-

crates/ark/src/lsp/diagnostics.rs

@@ -30,6 +30,7 @@ use crate::lsp::encoding::convert_tree_sitter_range_to_lsp_range;
 use crate::lsp::indexer;
 use crate::lsp::inputs::library::Library;
 use crate::lsp::inputs::package::Package;
+use crate::lsp::inputs::source_root::SourceRoot;
 use crate::lsp::state::WorldState;
 use crate::lsp::traits::node::NodeExt;
 use crate::lsp::traits::rope::RopeExt;
@@ -62,6 +63,9 @@ pub struct DiagnosticContext<'a> {
     // The set of packages that are currently installed.
     pub installed_packages: HashSet<String>,
 
+    /// Reference to source root, if any.
+    pub root: &'a Option<SourceRoot>,
+
     /// Reference to the library for looking up package exports.
     pub library: &'a Library,
 
@@ -83,13 +87,14 @@ impl Default for DiagnosticsConfig {
 }
 
 impl<'a> DiagnosticContext<'a> {
-    pub fn new(contents: &'a Rope, library: &'a Library) -> Self {
+    pub fn new(contents: &'a Rope, root: &'a Option<SourceRoot>, library: &'a Library) -> Self {
         Self {
             contents,
             document_symbols: Vec::new(),
             session_symbols: HashSet::new(),
             workspace_symbols: HashSet::new(),
             installed_packages: HashSet::new(),
+            root,
             library,
             library_symbols: BTreeMap::new(),
             in_formula: false,
@@ -144,7 +149,7 @@ pub(crate) fn generate_diagnostics(doc: Document, state: WorldState) -> Vec<Diag
         return diagnostics;
     }
 
-    let mut context = DiagnosticContext::new(&doc.contents, &state.library);
+    let mut context = DiagnosticContext::new(&doc.contents, &state.root, &state.library);
 
     // Add a 'root' context for the document.
     context.document_symbols.push(HashMap::new());
@@ -157,6 +162,23 @@ pub(crate) fn generate_diagnostics(doc: Document, state: WorldState) -> Vec<Diag
         _ => {},
     });
 
+    // If this is a package, add imported symbols to workspace
+    if let Some(SourceRoot::Package(root)) = &state.root {
+        // Add symbols from `importFrom()` directives
+        for import in &root.namespace.imports {
+            context.workspace_symbols.insert(import.clone());
+        }
+
+        // Add symbols from `import()` directives
+        for package_import in &root.namespace.package_imports {
+            if let Some(pkg) = state.library.get(package_import) {
+                for export in &pkg.namespace.exports {
+                    context.workspace_symbols.insert(export.clone());
+                }
+            }
+        }
+    }
+
     // Add per-environment session symbols
     for scope in state.console_scopes.iter() {
         for name in scope.iter() {

crates/ark/src/lsp/diagnostics_syntax.rs

@@ -314,7 +314,7 @@ mod tests {
     fn text_diagnostics(text: &str) -> Vec<Diagnostic> {
         let document = Document::new(text, None);
         let library = Library::default();
-        let context = DiagnosticContext::new(&document.contents, &library);
+        let context = DiagnosticContext::new(&document.contents, &None, &library);
         let diagnostics = syntax_diagnostics(document.ast.root_node(), &context).unwrap();
         diagnostics
     }

crates/ark/src/lsp/inputs/library.rs

@@ -68,7 +68,7 @@ impl Library {
 
     fn load_package(&self, name: &str) -> anyhow::Result<Option<Package>> {
         for lib_path in self.library_paths.iter() {
-            match Package::load(&lib_path, name)? {
+            match Package::load_from_library(&lib_path, name)? {
                 Some(pkg) => return Ok(Some(pkg)),
                 None => (),
             }

crates/ark/src/lsp/inputs/package.rs

@@ -23,15 +23,12 @@ pub struct Package {
 }
 
 impl Package {
-    /// Attempts to load a package from the given path and name.
-    pub fn load(lib_path: &std::path::Path, name: &str) -> anyhow::Result<Option<Self>> {
-        let package_path = lib_path.join(name);
-
+    /// Load a package from a given path.
+    pub fn load(package_path: &std::path::Path) -> anyhow::Result<Option<Self>> {
         let description_path = package_path.join("DESCRIPTION");
         let namespace_path = package_path.join("NAMESPACE");
 
-        // Only consider libraries that have a folder named after the
-        // requested package and that contains a description file
+        // Only consider directories that contain a description file
         if !description_path.is_file() {
             return Ok(None);
         }
@@ -41,24 +38,42 @@ impl Package {
         let description_contents = fs::read_to_string(&description_path)?;
         let description = Description::parse(&description_contents)?;
 
-        if description.name != name {
-            return Err(anyhow::anyhow!(
-                "`Package` field in `DESCRIPTION` doesn't match folder name '{name}'"
-            ));
-        }
-
         let namespace = if namespace_path.is_file() {
             let namespace_contents = fs::read_to_string(&namespace_path)?;
             Namespace::parse(&namespace_contents)?
         } else {
-            tracing::info!("Package `{name}` doesn't contain a NAMESPACE file, using defaults");
+            tracing::info!(
+                "Package `{name}` doesn't contain a NAMESPACE file, using defaults",
+                name = description.name
+            );
             Namespace::default()
         };
 
         Ok(Some(Package {
-            path: package_path,
+            path: package_path.to_path_buf(),
             description,
             namespace,
         }))
     }
+
+    /// Load a package from the given library path and name.
+    pub fn load_from_library(
+        lib_path: &std::path::Path,
+        name: &str,
+    ) -> anyhow::Result<Option<Self>> {
+        let package_path = lib_path.join(name);
+
+        // For library packages, ensure the invariant that the package name
+        // matches the folder name
+        if let Some(pkg) = Self::load(&package_path)? {
+            if pkg.description.name != name {
+                return Err(anyhow::anyhow!(
+                    "`Package` field in `DESCRIPTION` doesn't match folder name '{name}'"
+                ));
+            }
+            Ok(Some(pkg))
+        } else {
+            Ok(None)
+        }
+    }
 }

crates/ark/src/lsp/state.rs

@@ -51,7 +51,7 @@ pub(crate) struct WorldState {
     pub(crate) installed_packages: Vec<String>,
 
     /// The root of the source tree (e.g., a package).
-    pub(crate) _root: Option<SourceRoot>,
+    pub(crate) root: Option<SourceRoot>,
 
     /// Map of package name to package metadata for installed libraries. Lazily populated.
     pub(crate) library: Library,

crates/ark/src/lsp/state_handlers.rs

@@ -45,6 +45,8 @@ use crate::lsp::config::GLOBAL_SETTINGS;
 use crate::lsp::documents::Document;
 use crate::lsp::encoding::get_position_encoding_kind;
 use crate::lsp::indexer;
+use crate::lsp::inputs::package::Package;
+use crate::lsp::inputs::source_root::SourceRoot;
 use crate::lsp::main_loop::DidCloseVirtualDocumentParams;
 use crate::lsp::main_loop::DidOpenVirtualDocumentParams;
 use crate::lsp::main_loop::LspState;
@@ -84,8 +86,35 @@ pub(crate) fn initialize(
         for folder in workspace_folders.iter() {
             state.workspace.folders.push(folder.uri.clone());
             if let Ok(path) = folder.uri.to_file_path() {
-                if let Some(path) = path.to_str() {
-                    folders.push(path.to_string());
+                // Try to load package from this workspace folder and set as
+                // root if found. This means we're dealing with a package
+                // source.
+                if state.root.is_none() {
+                    match Package::load(&path) {
+                        Ok(Some(pkg)) => {
+                            log::info!(
+                                "Root: Loaded package `{pkg}` from {path} as project root",
+                                pkg = pkg.description.name,
+                                path = path.display()
+                            );
+                            state.root = Some(SourceRoot::Package(pkg));
+                        },
+                        Ok(None) => {
+                            log::info!(
+                                "Root: No package found at {path}, treating as folder of scripts",
+                                path = path.display()
+                            );
+                        },
+                        Err(err) => {
+                            log::warn!(
+                                "Root: Error loading package at {path}: {err}",
+                                path = path.display()
+                            );
+                        },
+                    }
+                }
+                if let Some(path_str) = path.to_str() {
+                    folders.push(path_str.to_string());
                 }
             }
         }