Implement the production SourceHandler for the LSP

Author: DavisVaughan

Cargo.lock

@@ -55,7 +55,9 @@ oak_ide.workspace = true
 oak_package_metadata.workspace = true
 oak_scan.workspace = true
 oak_semantic.workspace = true
+oak_source.workspace = true
 oak_sources.workspace = true
+oak_srcref.workspace = true
 once_cell.workspace = true
 regex.workspace = true
 reqwest.workspace = true
@@ -90,6 +92,7 @@ ark = { workspace = true, features = ["testing"] }
 ark_test.workspace = true
 assert_matches.workspace = true
 insta.workspace = true
+oak_r_process.workspace = true
 oak_semantic = { workspace = true, features = ["testing"] }
 oak_package_metadata.workspace = true
 stdext = { workspace = true, features = ["testing"] }

crates/ark/Cargo.toml

@@ -13,6 +13,7 @@ use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
+use std::sync::Arc;
 use std::sync::LazyLock;
 use std::sync::RwLock;
 
@@ -53,7 +54,9 @@ use crate::lsp::diagnostics::generate_diagnostics;
 use crate::lsp::handlers;
 use crate::lsp::indexer;
 use crate::lsp::open_file::OpenFile;
+use crate::lsp::sources::OakSourceHandler;
 use crate::lsp::sources::SourceCompleted;
+use crate::lsp::sources::SourceHandler;
 use crate::lsp::sources::SourceScheduler;
 use crate::lsp::state::WorldState;
 use crate::lsp::state_handlers;
@@ -240,7 +243,7 @@ impl GlobalState {
     ///   and auxiliary loop.
     pub(crate) fn new(
         client: Client,
-        _r_home: PathBuf,
+        r_home: PathBuf,
         console_notification_tx: TokioUnboundedSender<ConsoleNotification>,
     ) -> Self {
         // FIXME: We shouldn't call R code in the kernel to figure this out
@@ -268,7 +271,10 @@ impl GlobalState {
         Self::from_parts(
             client,
             WorldState::new(db, library),
-            LspState::new(console_notification_tx, SourceScheduler::new(None)),
+            LspState::new(
+                console_notification_tx,
+                SourceScheduler::new(source_handler(&r_home)),
+            ),
         )
     }
 
@@ -603,6 +609,33 @@ impl GlobalState {
     }
 }
 
+/// Build the LSP's [`SourceHandler`], or `None` to disable source fetching
+fn source_handler(r_home: &Path) -> Option<Arc<dyn SourceHandler>> {
+    if !cfg!(debug_assertions) {
+        // TODO!: Remove this to activate in release builds as well.
+        // Currently only active in debug builds (including unit and integration tests).
+        return None;
+    }
+
+    let Some(r) = harp::command::r_executable(r_home) else {
+        log::warn!(
+            "Can't locate an R executable under '{}', package source fetching is disabled",
+            r_home.display()
+        );
+        return None;
+    };
+
+    match OakSourceHandler::new(r) {
+        Ok(handler) => Some(Arc::new(handler)),
+        Err(err) => {
+            log::error!(
+                "Can't create package source handler, source fetching is disabled: {err:?}"
+            );
+            None
+        },
+    }
+}
+
 /// Test-only methods for driving the main loop without R or a live LSP
 /// connection. Kept here, next to the loop they exercise, so the pump uses the
 /// real `handle_event()` and the private channels rather than a reconstruction.

crates/ark/src/lsp/main_loop.rs

@@ -6,6 +6,9 @@ use std::sync::Arc;
 use aether_path::FilePath;
 use oak_db::Db;
 use oak_db::Package;
+use oak_db::Priority;
+use oak_source::SourceCache;
+use oak_srcref::SrcrefCache;
 use stdext::result::ResultExt;
 
 use crate::lsp::main_loop::Event;
@@ -19,16 +22,97 @@ pub(crate) trait SourceHandler: Send + Sync {
     fn handle(&self, request: &SourceRequest) -> SourceResponse;
 }
 
+/// The production [`SourceHandler`]
+///
+/// Recovers a package's R sources in the following order:
+/// - Base packages come from a downloaded CRAN's R source tarball
+/// - CRAN packages come from:
+///   - A local `srcref`, if available
+///   - A downloaded CRAN package tarball
+pub(crate) struct OakSourceHandler {
+    srcref: SrcrefCache,
+    source: SourceCache,
+}
+
+impl OakSourceHandler {
+    /// Build the handler, opening both caches against the shared on disk cache
+    pub(crate) fn new(r: PathBuf) -> anyhow::Result<Self> {
+        Ok(Self {
+            srcref: SrcrefCache::new(r)?,
+            source: SourceCache::new()?,
+        })
+    }
+
+    /// Build the handler with both caches rooted under `root`, so tests don't touch the
+    /// real on disk cache
+    #[cfg(test)]
+    pub(crate) fn new_in(root: &Path, r: PathBuf) -> anyhow::Result<Self> {
+        Ok(Self {
+            srcref: SrcrefCache::new_in(root.join("srcref"), r)?,
+            source: SourceCache::new_in(root.join("source"))?,
+        })
+    }
+}
+
+impl SourceHandler for OakSourceHandler {
+    fn handle(&self, request: &SourceRequest) -> SourceResponse {
+        let name = request.name();
+        let version = request.version();
+
+        // Base packages only exist in the R version source tarball, served from one
+        // download shared across all of them
+        if matches!(request.priority(), Some(Priority::Base)) {
+            return self
+                .source
+                .get_r(version)
+                .or_else(|| self.source.insert_r(version))
+                .map(|root| SourceResponse::Success(r_dir_for_base(&root, name)))
+                .unwrap_or(SourceResponse::Failure);
+        }
+
+        // Try to "get" from all sources before doing an expensive "insert"
+        if let Some(dir) = self.srcref.get(name, version, request.built()) {
+            return SourceResponse::Success(dir);
+        }
+        if let Some(root) = self.source.get_cran(name, version) {
+            return SourceResponse::Success(r_dir_for_cran(&root));
+        }
+
+        // Prefer `srcref` to CRAN download since it doesn't require internet and would
+        // be an exact match to the installed package
+        if let Some(dir) =
+            self.srcref
+                .insert(name, version, request.built(), request.library_path())
+        {
+            return SourceResponse::Success(dir);
+        }
+        if let Some(root) = self.source.insert_cran(name, version) {
+            return SourceResponse::Success(r_dir_for_cran(&root));
+        }
+
+        SourceResponse::Failure
+    }
+}
+
+/// Find `R/` from a CRAN package tarball
+fn r_dir_for_cran(root: &Path) -> PathBuf {
+    root.join("R")
+}
+
+/// Find `R/` from a R source tarball for a base package
+fn r_dir_for_base(root: &Path, name: &str) -> PathBuf {
+    root.join("src").join("library").join(name).join("R")
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct SourceRequest {
     name: String,
     version: String,
     built: String,
+    priority: Option<Priority>,
     library_path: PathBuf,
 }
 
-// TODO!: Remove when we have a production `SourceHandler`
-#[cfg_attr(not(test), expect(dead_code))]
 #[derive(Debug)]
 pub(crate) enum SourceResponse {
     Success(PathBuf),
@@ -56,7 +140,6 @@ enum SourceState {
 }
 
 pub(crate) struct SourceScheduler {
-    // TODO!: Remove the `Option<>` when we implement a production `SourceHandler`
     handler: Option<Arc<dyn SourceHandler>>,
     state: HashMap<Package, SourceState>,
 }
@@ -147,6 +230,8 @@ impl SourceRequest {
             ));
         };
 
+        let priority = package.priority(db).clone();
+
         let library_path = match package.description_path(db) {
             FilePath::File(path) => {
                 match path.as_path().as_std_path().parent().and_then(Path::parent) {
@@ -169,30 +254,27 @@ impl SourceRequest {
             name,
             version,
             built,
+            priority,
             library_path,
         })
     }
 
-    // TODO!: Remove when we have a production `SourceHandler`
-    #[cfg_attr(not(test), expect(dead_code))]
     pub(crate) fn name(&self) -> &str {
         &self.name
     }
 
-    // TODO!: Remove when we have a production `SourceHandler`
-    #[cfg_attr(not(test), expect(dead_code))]
     pub(crate) fn version(&self) -> &str {
         &self.version
     }
 
-    // TODO!: Remove when we have a production `SourceHandler`
-    #[cfg_attr(not(test), expect(dead_code))]
     pub(crate) fn built(&self) -> &str {
         &self.built
     }
 
-    // TODO!: Remove when we have a production `SourceHandler`
-    #[cfg_attr(not(test), expect(dead_code))]
+    pub(crate) fn priority(&self) -> Option<&Priority> {
+        self.priority.as_ref()
+    }
+
     pub(crate) fn library_path(&self) -> &Path {
         &self.library_path
     }

crates/ark/src/lsp/sources.rs

@@ -3,6 +3,7 @@
 
 use std::collections::HashMap;
 use std::path::Path;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::Mutex;
 
@@ -26,6 +27,7 @@ use crate::lsp::main_loop::init_aux_for_test;
 use crate::lsp::main_loop::Event;
 use crate::lsp::main_loop::GlobalState;
 use crate::lsp::main_loop::LspState;
+use crate::lsp::sources::OakSourceHandler;
 use crate::lsp::sources::SourceHandler;
 use crate::lsp::sources::SourceRequest;
 use crate::lsp::sources::SourceResponse;
@@ -121,6 +123,28 @@ fn dispatched_names(calls: &Mutex<Vec<SourceRequest>>) -> Vec<String> {
         .collect()
 }
 
+/// Find R on the `PATH`
+///
+/// On Windows, `which` (from Git) returns POSIX paths that `Command::new()` can't resolve.
+/// Use `where` which returns native paths.
+fn find_r() -> PathBuf {
+    let output = std::process::Command::new(if cfg!(windows) { "where" } else { "which" })
+        .arg("R")
+        .output()
+        .unwrap_or_else(|err| panic!("Failed to find R: {err}"));
+    assert!(output.status.success());
+
+    // `where` on Windows can return multiple matches, take the first
+    PathBuf::from(
+        String::from_utf8(output.stdout)
+            .expect("Non-UTF8 R path")
+            .trim()
+            .lines()
+            .next()
+            .expect("R should exist"),
+    )
+}
+
 /// The happy path end to end: a workspace uses an installed library package via
 /// `::`, so the revision-advance check dispatches a source request, the handler
 /// returns a directory, and the main loop ingests it into the library package.
@@ -226,3 +250,71 @@ async fn test_failed_source_is_not_retried() {
         "donor"
     )]);
 }
+
+/// End to end against real `srcref` recovery: install {generics} from source into a
+/// temporary library, point a workspace at it via `::`, inject the real
+/// [`OakSourceHandler`], and assert the recovered sources are ingested.
+///
+/// Requires R on the `PATH` and internet access. We use {generics} because it is small and
+/// easy to install from source, the same package `oak_srcref`'s own extraction test uses.
+#[tokio::test]
+async fn test_source_pipeline_ingests_real_srcref_sources() {
+    let _aux = init_aux_for_test();
+
+    let r = find_r();
+
+    // Temporary library, with {generics} installed from source so srcrefs are preserved
+    let library = tempfile::tempdir().unwrap();
+
+    // Use forward slashes so the path is safe inside R string literals on Windows
+    let library_literal = library.path().display().to_string().replace('\\', "/");
+
+    let output = oak_r_process::run_text(
+        &r,
+        &format!(
+            r#"install.packages("generics", lib = "{library_literal}", repos = "https://cran.r-project.org", type = "source", INSTALL_opts = "--with-keep.source")"#,
+        ),
+        &[],
+        &[],
+    )
+    .expect("Failed to run install.packages()");
+    assert!(output.status.success());
+
+    // The real handler, with both caches rooted in a temp dir so the test doesn't touch
+    // the shared on disk cache
+    let cache = tempfile::tempdir().unwrap();
+    let handler: Arc<dyn SourceHandler> =
+        Arc::new(OakSourceHandler::new_in(cache.path(), r).unwrap());
+
+    let mut db = OakDatabase::new();
+    db.set_library_paths(&[library.path().to_path_buf()]);
+
+    let mut state = GlobalState::from_parts(
+        test_client(),
+        WorldState::new(db, Library::new(vec![])),
+        LspState::new(
+            tokio::sync::mpsc::unbounded_channel().0,
+            SourceScheduler::new(Some(handler)),
+        ),
+    );
+
+    // A workspace package that uses {generics} via `::`
+    let workspace = tempfile::tempdir().unwrap();
+    let myproj = workspace.path().join("myproj");
+    write_description(&myproj, "myproj");
+    write_sources(&myproj.join("R"), &[("use.R", "generics::as.factor()\n")]);
+
+    state
+        .handle_event_to_quiescence(did_change_workspace_folders(workspace.path()))
+        .await;
+
+    // {generics} now carries its recovered sources, readable from disk. {generics} is a
+    // package of S3 generics, so every recovered file is full of `UseMethod()` calls.
+    let db = &state.world().db;
+    let generics = db.package_by_name("generics").unwrap();
+    let files = generics.files(db).clone();
+    assert!(!files.is_empty());
+    assert!(files
+        .iter()
+        .any(|file| file.source_text(db).contains("UseMethod")));
+}