Add support for source() directive in symbol resolution

crates/ark/src/lsp/state.rs

@@ -3,11 +3,16 @@ use std::collections::HashSet;
 use std::path::Path;
 
 use anyhow::anyhow;
+use biome_rowan::TextRange;
 use oak_core::file::list_r_files;
 use oak_ide::FileScope;
+use oak_index::external::directive_layers;
 use oak_index::external::file_layers;
 use oak_index::external::package_root_layers;
 use oak_index::external::BindingSource;
+use oak_index::semantic_index::SemanticIndex;
+use oak_index::semantic_index_with_source_resolver;
+use oak_index::SourceResolution;
 use oak_package::collation::collation_order;
 use oak_package::library::Library;
 use stdext::result::ResultExt;
@@ -89,15 +94,39 @@ impl WorldState {
         }
     }
 
-    /// Create a scope chain for a particular file, taking into account the
-    /// current project type. For packages, this creates a scope containing
-    /// imports and top-level definitions in other files, respecting the
-    /// collation order.
-    pub(crate) fn file_scope(&self, file: &Url) -> FileScope {
-        let Some(SourceRoot::Package(ref pkg)) = self.root else {
-            return FileScope::search_path(default_search_path());
-        };
+    /// Look up a document by URL: returns an open document if available,
+    /// otherwise reads from disk.
+    ///
+    /// TODO: Replace with a proper VFS so non-opened workspace documents
+    /// are cached rather than re-read on every query.
+    fn workspace_document(&self, uri: &Url) -> Option<Document> {
+        if let Some(doc) = self.documents.get(uri) {
+            return Some(doc.clone());
+        }
+        let path = uri.to_file_path().log_err()?;
+        let contents = std::fs::read_to_string(&path).log_err()?;
+        Some(Document::new(&contents, None))
+    }
 
+    /// Create the semantic index and scope chain for a particular file.
+    ///
+    /// For scripts, the index is built with a source resolver so that
+    /// `source()` directives carry the sourced file's exports.
+    /// For packages, cross-file visibility comes from NAMESPACE imports and
+    /// collation ordering.
+    pub(crate) fn file_analysis(&self, file: &Url, doc: &Document) -> (SemanticIndex, FileScope) {
+        match self.root {
+            Some(SourceRoot::Package(ref pkg)) => self.package_file_analysis(file, doc, pkg),
+            _ => self.script_file_analysis(file, doc),
+        }
+    }
+
+    fn package_file_analysis(
+        &self,
+        file: &Url,
+        doc: &Document,
+        pkg: &oak_package::package::Package,
+    ) -> (SemanticIndex, FileScope) {
         let root_layers = package_root_layers(&pkg.namespace);
 
         // Collect R source filenames from open documents and disk. Open
@@ -149,16 +178,8 @@ impl WorldState {
             let Some(uri) = Url::from_file_path(&path).log_err() else {
                 continue;
             };
-
-            // Use the open document if available, otherwise read from disk.
-            // TODO: Store non-opened workspace documents in VFS.
-            let doc = if let Some(open) = self.documents.get(&uri) {
-                open
-            } else {
-                let Ok(contents) = std::fs::read_to_string(&path) else {
-                    continue;
-                };
-                &Document::new(&contents, None)
+            let Some(doc) = self.workspace_document(&uri) else {
+                continue;
             };
 
             let layers = file_layers(uri, &doc.semantic_index());
@@ -173,7 +194,98 @@ impl WorldState {
         lazy.extend(root_layers);
         lazy.push(BindingSource::PackageExports("base".to_string()));
 
-        FileScope::package(top_level, lazy)
+        (doc.semantic_index(), FileScope::package(top_level, lazy))
+    }
+
+    fn script_file_analysis(&self, file: &Url, doc: &Document) -> (SemanticIndex, FileScope) {
+        // Resolve `source()` paths relative to the workspace root,
+        // matching RStudio's behaviour of setting the working directory
+        // to the project root. Fall back to the file's own directory
+        // when no workspace folder is open.
+        let file_dir = file
+            .to_file_path()
+            .ok()
+            .and_then(|p| p.parent().map(|d| d.to_path_buf()));
+        let source_root = self
+            .workspace
+            .folders
+            .first()
+            .and_then(|url| url.to_file_path().ok())
+            .or(file_dir);
+
+        let mut stack = HashSet::new();
+        stack.insert(file.clone());
+
+        let index = semantic_index_with_source_resolver(&doc.parse.tree(), |path| {
+            let dir = source_root.as_ref()?;
+            self.resolve_source(dir, path, &mut stack)
+        });
+
+        let directives = directive_layers(index.file_directives());
+        (
+            index,
+            FileScope::search_path(directives, default_search_path()),
+        )
+    }
+
+    /// Resolve a `source()` call into a [`SourceResolution`] containing the
+    /// sourced file's exported definitions and `library()` package attachments.
+    ///
+    /// `stack` tracks files currently being resolved (grey set) to break
+    /// cycles. A file is added when resolution starts and removed when it
+    /// finishes, so shared dependencies (diamond patterns) are resolved
+    /// independently for each parent.
+    fn resolve_source(
+        &self,
+        base_dir: &Path,
+        path: &str,
+        stack: &mut HashSet<Url>,
+    ) -> Option<SourceResolution> {
+        let resolved = base_dir.join(path);
+        let url = Url::from_file_path(&resolved).log_err()?;
+
+        if !stack.insert(url.clone()) {
+            return None;
+        }
+
+        let sourced_doc = self.workspace_document(&url)?;
+
+        // Build the sourced file's index with a nested resolver so that
+        // transitive `source()` calls are also resolved. The base
+        // directory stays the same (workspace root) throughout the chain.
+        let index = semantic_index_with_source_resolver(&sourced_doc.parse.tree(), |nested_path| {
+            self.resolve_source(base_dir, nested_path, stack)
+        });
+
+        let mut definitions: Vec<(String, Url, TextRange)> = index
+            .file_all_definitions(&url)
+            .into_iter()
+            .map(|(name, file, range)| (name.to_string(), file, range))
+            .collect();
+
+        let mut packages = Vec::new();
+        for d in index.file_directives() {
+            match d.kind() {
+                oak_index::semantic_index::DirectiveKind::Attach(pkg) => {
+                    packages.push(pkg.clone());
+                },
+                oak_index::semantic_index::DirectiveKind::Source {
+                    file: source_file,
+                    exports,
+                } => {
+                    for (name, range) in exports {
+                        definitions.push((name.clone(), source_file.clone(), *range));
+                    }
+                },
+            }
+        }
+
+        stack.remove(&url);
+
+        Some(SourceResolution {
+            definitions,
+            packages,
+        })
     }
 }
 

crates/oak_core/src/declaration.rs

@@ -0,0 +1,134 @@
+//! Helpers for detecting `declare()` annotations in R source code.
+//!
+//! `declare()` is a no-op function in R (>= 4.5) meant to hold static
+//! annotations. The compat syntax uses `~declare(...)` (a formula, also a
+//! no-op) for older R versions.
+//!
+//! This module recognises the `declare()` wrapper and returns its arguments for
+//! the caller to interpret.
+
+use aether_syntax::AnyRExpression;
+use aether_syntax::RCall;
+use aether_syntax::RCallArguments;
+use aether_syntax::RSyntaxKind;
+
+use crate::syntax_ext::RIdentifierExt;
+
+/// If `expr` is `declare(...)` or `~declare(...)`, return the arguments
+/// of the `declare()` call. Returns `None` if the expression doesn't
+/// match either pattern.
+pub fn as_declare_args(expr: &AnyRExpression) -> Option<RCallArguments> {
+    let call = as_declare_call(expr)?;
+    call.arguments().ok()
+}
+
+/// Unwrap `declare(...)` or `~declare(...)` to get the `declare` call node.
+fn as_declare_call(expr: &AnyRExpression) -> Option<RCall> {
+    match expr {
+        AnyRExpression::RCall(call) if is_declare(call) => Some(call.clone()),
+
+        AnyRExpression::RUnaryExpression(unary) => {
+            let op = unary.operator().ok()?;
+            if op.kind() != RSyntaxKind::TILDE {
+                return None;
+            }
+            let AnyRExpression::RCall(call) = unary.argument().ok()? else {
+                return None;
+            };
+            if is_declare(&call) {
+                Some(call)
+            } else {
+                None
+            }
+        },
+
+        _ => None,
+    }
+}
+
+fn is_declare(call: &RCall) -> bool {
+    let Ok(AnyRExpression::RIdentifier(ident)) = call.function() else {
+        return false;
+    };
+    ident.name_text() == "declare"
+}
+
+#[cfg(test)]
+mod tests {
+    use aether_parser::RParserOptions;
+    use aether_syntax::AnyRExpression;
+    use biome_rowan::AstNode;
+    use biome_rowan::AstNodeList;
+    use biome_rowan::AstSeparatedList;
+
+    use super::*;
+
+    fn parse_single_expr(code: &str) -> AnyRExpression {
+        let parsed = aether_parser::parse(code, RParserOptions::default());
+        parsed.tree().expressions().iter().next().unwrap()
+    }
+
+    fn declare_arg_values(code: &str) -> Option<Vec<String>> {
+        let expr = parse_single_expr(code);
+        let args = as_declare_args(&expr)?;
+        Some(
+            args.items()
+                .iter()
+                .filter_map(|arg| {
+                    let arg = arg.ok()?;
+                    Some(arg.value()?.syntax().text_trimmed().to_string())
+                })
+                .collect(),
+        )
+    }
+
+    #[test]
+    fn test_declare_returns_arguments() {
+        let values = declare_arg_values("declare(source(\"helpers.R\"))");
+        assert_eq!(values, Some(vec!["source(\"helpers.R\")".to_string()]));
+    }
+
+    #[test]
+    fn test_tilde_declare_returns_arguments() {
+        let values = declare_arg_values("~declare(source(\"helpers.R\"))");
+        assert_eq!(values, Some(vec!["source(\"helpers.R\")".to_string()]));
+    }
+
+    #[test]
+    fn test_bare_call_not_declare() {
+        let values = declare_arg_values("source(\"helpers.R\")");
+        assert_eq!(values, None);
+    }
+
+    #[test]
+    fn test_tilde_not_declare() {
+        let values = declare_arg_values("~other(source(\"helpers.R\"))");
+        assert_eq!(values, None);
+    }
+
+    #[test]
+    fn test_declare_no_args() {
+        let values = declare_arg_values("declare()");
+        assert_eq!(values, Some(vec![]));
+    }
+
+    #[test]
+    fn test_declare_multiple_args() {
+        let values = declare_arg_values("declare(source(\"a.R\"), source(\"b.R\"))");
+        assert_eq!(
+            values,
+            Some(vec![
+                "source(\"a.R\")".to_string(),
+                "source(\"b.R\")".to_string(),
+            ])
+        );
+    }
+
+    #[test]
+    fn test_declare_preserves_named_args() {
+        let expr = parse_single_expr("declare(foo = source(\"a.R\"))");
+        let args = as_declare_args(&expr).unwrap();
+        let arg = args.items().iter().next().unwrap().unwrap();
+        assert!(arg.name_clause().is_some());
+    }
+}

crates/oak_core/src/lib.rs

@@ -1,2 +1,3 @@
+pub mod declaration;
 pub mod file;
 pub mod syntax_ext;

crates/oak_ide/src/goto_definition.rs

@@ -4,6 +4,7 @@ use oak_index::external::resolve_external_name;
 use oak_index::external::resolve_in_package;
 use oak_index::external::BindingSource;
 use oak_index::external::ExternalDefinition;
+use oak_index::semantic_index::DefinitionKind;
 use oak_index::semantic_index::SemanticIndex;
 use oak_index::DefinitionId;
 use oak_index::ScopeId;
@@ -60,7 +61,7 @@ pub fn goto_definition(
         },
         Identifier::Use { scope_id, use_id } => {
             let scope_chain = scope.at(index, offset);
-            resolve_use(scope_id, use_id, file, index, scope_chain, library)
+            resolve_use(scope_id, use_id, file, index, &scope_chain, library)
         },
         Identifier::NamespaceAccess {
             ref package,
@@ -90,8 +91,12 @@ fn resolve_use(
         defs.iter()
             .map(|&def_id| {
                 let def = &index.definitions(scope)[def_id];
+                let target_file = match def.kind() {
+                    DefinitionKind::Sourced { file: source_file } => source_file.clone(),
+                    _ => file.clone(),
+                };
                 NavigationTarget {
-                    file: file.clone(),
+                    file: target_file,
                     name: symbol_name.to_string(),
                     full_range: def.range(),
                     focus_range: def.range(),