(4) Wire up the production SourceHandler in the LSP (#1301)

Part of #1234
Branched from #1300 

This PR is most of the point of this series

The main point is that the source downloading is still off by default in
release builds, but is on in debug builds (including unit and
integration tests).

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,29 +6,113 @@ 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;
 use crate::lsp::main_loop::TokioUnboundedSender;
 
 /// Handle source requests to provide R source files for an installed package
 ///
-/// Implementations live outside of oak, oak is only in charge of ingesting the
-/// returned directory.
+/// Implemented as a trait so we can have a production level [OakSourceHandler]
+/// and a test level `TestSourceHandler`.
 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")));
+}

crates/ark/src/lsp/tests/sources.rs

@@ -39,6 +39,7 @@ pub use inputs::RootKind;
 pub use inputs::StaleRoot;
 pub use inputs::WorkspaceRoots;
 pub use name::Name;
+pub use oak_package_metadata::description::Priority;
 pub use package::Package;
 pub use package_resolve::PackageVisibility;
 pub use storage::OakDatabase;