Add symbols imported in packages to workspace

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,
@@ -130,7 +135,11 @@ impl<'a> DiagnosticContext<'a> {
     }
 }
 
-pub(crate) fn generate_diagnostics(doc: Document, state: WorldState) -> Vec<Diagnostic> {
+pub(crate) fn generate_diagnostics(
+    doc: Document,
+    state: WorldState,
+    testthat: bool,
+) -> Vec<Diagnostic> {
     let mut diagnostics = Vec::new();
 
     if !state.config.diagnostics.enable {
@@ -144,7 +153,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());
@@ -154,9 +163,45 @@ pub(crate) fn generate_diagnostics(doc: Document, state: WorldState) -> Vec<Diag
         indexer::IndexEntryData::Function { name, arguments: _ } => {
             context.workspace_symbols.insert(name.to_string());
         },
+        indexer::IndexEntryData::Variable { name } => {
+            context.workspace_symbols.insert(name.to_string());
+        },
         _ => {},
     });
 
+    // 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());
+                }
+            }
+        }
+    }
+
+    // Simple workaround to include testthat exports in test files. I think the
+    // general principle would be that (a) files in `tests/testthat/` include
+    // `testthat.R` as a preamble (note that people modify that file e.g. to add
+    // more `library()` calls), and (b) all helper files are included in a
+    // test-specific workspace (which is effectively the case currently as we
+    // don't special-case how workspace inclusion works for packages). We might
+    // want to provide a mechanism for test packages to declare this sort of
+    // test files setup.
+    if testthat {
+        if let Some(pkg) = state.library.get("testthat") {
+            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() {
@@ -1097,10 +1142,10 @@ mod tests {
 
     use harp::eval::RParseEvalOptions;
     use once_cell::sync::Lazy;
+    use tower_lsp::lsp_types;
     use tower_lsp::lsp_types::Position;
 
     use crate::interface::console_inputs;
-    use crate::lsp::diagnostics::generate_diagnostics;
     use crate::lsp::documents::Document;
     use crate::lsp::inputs::library::Library;
     use crate::lsp::inputs::package::Package;
@@ -1113,6 +1158,10 @@ mod tests {
     // Default state that includes installed packages and default scopes.
     static DEFAULT_STATE: Lazy<WorldState> = Lazy::new(|| current_state());
 
+    fn generate_diagnostics(doc: Document, state: WorldState) -> Vec<lsp_types::Diagnostic> {
+        super::generate_diagnostics(doc, state, false)
+    }
+
     fn current_state() -> WorldState {
         let inputs = console_inputs().unwrap();
 

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_from_folder(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_from_folder(&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)
+        }
+    }
 }