diff --git a/crates/ark/src/lsp/completions/completion_context.rs b/crates/ark/src/lsp/completions/completion_context.rs index 2b122b5666..cc0c233525 100644 --- a/crates/ark/src/lsp/completions/completion_context.rs +++ b/crates/ark/src/lsp/completions/completion_context.rs @@ -13,18 +13,18 @@ use crate::lsp::completions::function_context::FunctionContext; use crate::lsp::completions::sources::composite::pipe::find_pipe_root; use crate::lsp::completions::sources::composite::pipe::PipeRoot; use crate::lsp::document_context::DocumentContext; -use crate::lsp::state::WorldState; +use crate::lsp::state::WorldSnapshot; use crate::treesitter::node_find_containing_call; pub(crate) struct CompletionContext<'a> { pub(crate) document_context: &'a DocumentContext<'a>, - pub(crate) state: &'a WorldState, + pub(crate) state: &'a WorldSnapshot, pipe_root_cell: OnceCell>, containing_call_cell: OnceCell>>, function_context_cell: OnceCell>, } impl<'a> CompletionContext<'a> { - pub fn new(document_context: &'a DocumentContext, state: &'a WorldState) -> Self { + pub fn new(document_context: &'a DocumentContext, state: &'a WorldSnapshot) -> Self { Self { document_context, state, diff --git a/crates/ark/src/lsp/completions/provide.rs b/crates/ark/src/lsp/completions/provide.rs index 6006d1efad..331c6f4748 100644 --- a/crates/ark/src/lsp/completions/provide.rs +++ b/crates/ark/src/lsp/completions/provide.rs @@ -11,7 +11,7 @@ use crate::lsp::completions::completion_context::CompletionContext; use crate::lsp::completions::sources::composite; use crate::lsp::completions::sources::unique; use crate::lsp::document_context::DocumentContext; -use crate::lsp::state::WorldState; +use crate::lsp::state::WorldSnapshot; use crate::lsp::traits::node::NodeExt; use crate::treesitter::NodeTypeExt; @@ -19,7 +19,7 @@ use crate::treesitter::NodeTypeExt; // Must be within an `r_task()`. pub(crate) fn provide_completions( document_context: &DocumentContext, - state: &WorldState, + state: &WorldSnapshot, ) -> anyhow::Result> { log::info!( "provide_completions() - Completion node text: '{node_text}', Node type: '{node_type:?}'", diff --git a/crates/ark/src/lsp/completions/sources/composite.rs b/crates/ark/src/lsp/completions/sources/composite.rs index 016e39965f..319e0ba2ac 100644 --- a/crates/ark/src/lsp/completions/sources/composite.rs +++ b/crates/ark/src/lsp/completions/sources/composite.rs @@ -252,7 +252,7 @@ mod tests { let (text, point) = point_from_cursor("@"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); assert!(context.document_context.node.is_program()); @@ -270,7 +270,7 @@ mod tests { let (text, point) = point_from_cursor(code); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); assert!(context.document_context.node.is_program()); diff --git a/crates/ark/src/lsp/completions/sources/composite/call.rs b/crates/ark/src/lsp/completions/sources/composite/call.rs index df44b00251..2a8cec643d 100644 --- a/crates/ark/src/lsp/completions/sources/composite/call.rs +++ b/crates/ark/src/lsp/completions/sources/composite/call.rs @@ -94,7 +94,7 @@ fn completions_from_call( }, }; - completions_from_arguments(&context.state.db, document_context, callee, object) + completions_from_arguments(context.state.db.read(), document_context, callee, object) } fn get_first_argument(context: &DocumentContext, node: &Node) -> anyhow::Result> { @@ -294,7 +294,7 @@ mod tests { let (text, point) = point_from_cursor("match(tab@)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); @@ -307,7 +307,7 @@ mod tests { let (text, point) = point_from_cursor("match(1, tab@)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); @@ -327,7 +327,7 @@ mod tests { let (text, point) = point_from_cursor("not_a_known_function(@)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap(); assert!(completions.is_none()); @@ -347,7 +347,7 @@ mod tests { let (text, point) = point_from_cursor("my_fun(@)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); @@ -364,7 +364,7 @@ mod tests { let (text, point) = point_from_cursor("my_fun@()"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap(); assert!(completions.is_none()); @@ -373,7 +373,7 @@ mod tests { let (text, point) = point_from_cursor("my_fun()@"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap(); assert!(completions.is_none()); @@ -396,7 +396,7 @@ mod tests { let (text, point) = point_from_cursor("my_fun(@)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); assert_eq!(completions.len(), 0); @@ -413,7 +413,7 @@ mod tests { let (text, point) = point_from_cursor("match(\n @\n)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); @@ -425,7 +425,7 @@ mod tests { let (text, point) = point_from_cursor("match(\n tab@\n)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); @@ -437,7 +437,7 @@ mod tests { let (text, point) = point_from_cursor("match(\n 1,\n tab@\n)"); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap().unwrap(); @@ -454,7 +454,7 @@ mod tests { let (text, point) = point_from_cursor(code_with_cursor); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); let completions = completions_from_call(&context).unwrap(); diff --git a/crates/ark/src/lsp/completions/sources/composite/workspace.rs b/crates/ark/src/lsp/completions/sources/composite/workspace.rs index 18c3d3cfc5..d3d8e8a044 100644 --- a/crates/ark/src/lsp/completions/sources/composite/workspace.rs +++ b/crates/ark/src/lsp/completions/sources/composite/workspace.rs @@ -69,7 +69,7 @@ fn completions_from_workspace( let token = token.as_str(); // get entries from the index - indexer::map(&state.db, |file, symbol, entry| { + indexer::map(state.db.read(), |file, symbol, entry| { if !symbol.fuzzy_matches(token) { return; } diff --git a/crates/ark/src/lsp/completions/sources/unique/custom.rs b/crates/ark/src/lsp/completions/sources/unique/custom.rs index 4bb2bde44e..bf5ecaabf9 100644 --- a/crates/ark/src/lsp/completions/sources/unique/custom.rs +++ b/crates/ark/src/lsp/completions/sources/unique/custom.rs @@ -217,7 +217,7 @@ mod tests { // Helper functions for testing custom completions fn assert_has_completion(code_with_cursor: &str, name: &str, expected_insert_text: &str) { let (text, point) = point_from_cursor(code_with_cursor); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let doc = TestDocument::new(&text); let document_context = doc.context(point); let context = CompletionContext::new(&document_context, &state); @@ -235,7 +235,7 @@ mod tests { fn assert_no_completions(code_with_cursor: &str) { let (text, point) = point_from_cursor(code_with_cursor); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let doc = TestDocument::new(&text); let document_context = doc.context(point); let context = CompletionContext::new(&document_context, &state); @@ -254,7 +254,7 @@ mod tests { }; let (text, point) = point_from_cursor("library(@)"); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let doc = TestDocument::new(&text); let document_context = doc.context(point); let context = CompletionContext::new(&document_context, &state); @@ -268,7 +268,7 @@ mod tests { assert_eq!(n_compls, n_packages); let (text, point) = point_from_cursor("library(uti@)"); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let doc = TestDocument::new(&text); let document_context = doc.context(point); let context = CompletionContext::new(&document_context, &state); diff --git a/crates/ark/src/lsp/completions/sources/unique/namespace.rs b/crates/ark/src/lsp/completions/sources/unique/namespace.rs index f89f0ff54b..dea4106da3 100644 --- a/crates/ark/src/lsp/completions/sources/unique/namespace.rs +++ b/crates/ark/src/lsp/completions/sources/unique/namespace.rs @@ -256,7 +256,7 @@ mod tests { let (text, point) = point_from_cursor(cursor_text); let doc = TestDocument::new(&text); let document_context = doc.context(point); - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let context = CompletionContext::new(&document_context, &state); completions_from_namespace(&context) diff --git a/crates/ark/src/lsp/completions/sources/unique/string.rs b/crates/ark/src/lsp/completions/sources/unique/string.rs index 9f19bd2c32..478e2d5dcd 100644 --- a/crates/ark/src/lsp/completions/sources/unique/string.rs +++ b/crates/ark/src/lsp/completions/sources/unique/string.rs @@ -132,7 +132,7 @@ mod tests { assert_match!(res, Some(items) => { assert!(items.is_empty()) }); // Check for same result when consulting (potentially all) unique sources - let state = WorldState::default(); + let state = WorldState::default().snapshot(); let completion_context = CompletionContext::new(&context, &state); let res = unique::get_completions(&completion_context).unwrap(); diff --git a/crates/ark/src/lsp/completions/tests/utils.rs b/crates/ark/src/lsp/completions/tests/utils.rs index b09a12de1f..36953c0765 100644 --- a/crates/ark/src/lsp/completions/tests/utils.rs +++ b/crates/ark/src/lsp/completions/tests/utils.rs @@ -18,7 +18,7 @@ pub(crate) fn get_completions_at_cursor(cursor_text: &str) -> anyhow::Result Ok(completions), diff --git a/crates/ark/src/lsp/db.rs b/crates/ark/src/lsp/db.rs index be474f36ba..cc42eca66e 100644 --- a/crates/ark/src/lsp/db.rs +++ b/crates/ark/src/lsp/db.rs @@ -17,6 +17,37 @@ pub(crate) trait ArkDb: oak_db::Db {} #[salsa::db] impl ArkDb for OakDatabase {} +/// Read-only handle to an `OakDatabase` that acts as a read-only snapshot. +/// +/// An `Analysis` snapshot holds the DB clone in a private field and lends a +/// shared ref via `read()`. This prevents a read-only background workers +/// from reaching DB setters and deadock with the main loop. This is oak's +/// version of rust-analyzer's `Analysis` facade. +#[derive(Clone, Debug)] +pub(crate) struct Analysis { + db: OakDatabase, +} + +impl Analysis { + pub(crate) fn new(db: OakDatabase) -> Self { + Self { db } + } + + /// The database as a shared `&dyn ArkDb`. Only lends `&`, so a reader can + /// query but can't reach a setter through it. + pub(crate) fn read(&self) -> &dyn ArkDb { + &self.db + } + + /// The database's salsa cancellation token. Read-side only: it observes and + /// arms cancellation, it doesn't mutate any input, so it's safe to expose + /// on the read-only facade. Only cancellation tests arm it by hand. + #[cfg(test)] + pub(crate) fn cancellation_token(&self) -> salsa::CancellationToken { + salsa::Database::cancellation_token(&self.db) + } +} + /// Extension trait that adds ark-side query methods to `oak_db::File`. pub(crate) trait FileArkExt { fn tree_sitter(self, db: &dyn ArkDb) -> &tree_sitter::Tree; diff --git a/crates/ark/src/lsp/diagnostics.rs b/crates/ark/src/lsp/diagnostics.rs index 00e6185fa7..79b03baf1f 100644 --- a/crates/ark/src/lsp/diagnostics.rs +++ b/crates/ark/src/lsp/diagnostics.rs @@ -33,7 +33,7 @@ use crate::lsp::diagnostics_syntax::syntax_diagnostics; use crate::lsp::indexer; use crate::lsp::inputs::source_root::SourceRoot; use crate::lsp::open_file::lsp_range_from_tree_sitter_range; -use crate::lsp::state::WorldState; +use crate::lsp::state::WorldSnapshot; use crate::lsp::traits::node::NodeExt; use crate::treesitter::node_has_error_or_missing; use crate::treesitter::BinaryOperatorType; @@ -151,7 +151,7 @@ impl<'a> DiagnosticContext<'a> { pub(crate) fn generate_diagnostics( file: File, - state: WorldState, + state: WorldSnapshot, testthat: bool, ) -> Vec { let mut diagnostics = Vec::new(); @@ -160,7 +160,7 @@ pub(crate) fn generate_diagnostics( return diagnostics; } - let db = &state.db; + let db = state.db.read(); // Check that diagnostics are not disabled in top-level declarations for // this document @@ -1193,7 +1193,7 @@ mod tests { Some(code.to_string()), None, ); - super::generate_diagnostics(file, state, false) + super::generate_diagnostics(file, state.snapshot(), false) } fn current_state() -> WorldState { diff --git a/crates/ark/src/lsp/handlers.rs b/crates/ark/src/lsp/handlers.rs index d50ce37243..88d38f880e 100644 --- a/crates/ark/src/lsp/handlers.rs +++ b/crates/ark/src/lsp/handlers.rs @@ -233,10 +233,15 @@ pub(crate) fn handle_completion( ); lsp::log_info!("Completion context: {:#?}", context); - // TODO(oak/completions): Clone so the closure captures by value. `r_task()` - // sends the closure across threads, and `&WorldState` isn't `Send` because - // `OakDatabase`'s salsa storage keeps thread-local query state. - let state = state.clone(); + // Snapshot so the closure captures by value. `r_task()` sends the closure + // across threads, and `&WorldState` isn't `Send` because `OakDatabase`'s + // salsa storage keeps thread-local query state. `snapshot()` hands the + // reader a `WorldSnapshot`, so the background thread can query oak but + // can't call a setter. + // TODO(oak/completions): We don't really need a snapshot here since + // completions are serviced from the main loop, it's only needed for the + // `r_task()`. + let state = state.snapshot(); let completions = r_task(move || provide_completions(&context, &state))?; if !completions.is_empty() { diff --git a/crates/ark/src/lsp/main_loop.rs b/crates/ark/src/lsp/main_loop.rs index 2f8e6c0d97..81595eb54a 100644 --- a/crates/ark/src/lsp/main_loop.rs +++ b/crates/ark/src/lsp/main_loop.rs @@ -50,6 +50,7 @@ use crate::lsp::backend::LspRequest; use crate::lsp::backend::LspResponse; use crate::lsp::backend::LspResult; use crate::lsp::capabilities::Capabilities; +use crate::lsp::db::Analysis; use crate::lsp::diagnostics::generate_diagnostics; use crate::lsp::handlers; use crate::lsp::indexer; @@ -58,6 +59,7 @@ use crate::lsp::sources::OakSourceHandler; use crate::lsp::sources::SourceCompleted; use crate::lsp::sources::SourceHandler; use crate::lsp::sources::SourceScheduler; +use crate::lsp::state::WorldSnapshot; use crate::lsp::state::WorldState; use crate::lsp::state_handlers; use crate::lsp::state_handlers::ConsoleInputs; @@ -559,7 +561,7 @@ impl GlobalState { // and the diagnostics passes they trigger force the same // memos. if !self.lsp_state.oak_scheduler.has_pending_scans() { - warm_workspace_index(self.world.db.clone()); + warm_workspace_index(Analysis::new(self.world.db.clone())); } }, @@ -1005,7 +1007,7 @@ impl std::fmt::Debug for TraceKernelNotification<'_> { pub(crate) struct RefreshDiagnosticsTask { /// Snapshot carrying the live oak plus the session context the diagnostics /// walk reads. See [`WorldState::diagnostics_snapshot`]. - state: WorldState, + state: WorldSnapshot, /// The file to diagnose, built against the live oak at enqueue time. file: OpenFile, } @@ -1144,11 +1146,11 @@ pub(crate) fn diagnostics_refresh_all(state: &WorldState) { /// cancelled (`spawn_blocking()` swallows the unwind). A cancelling write can /// only come from an editor buffer, so a document is open, and the diagnostics /// passes spawned by that same write force the same memos and finish the job. -fn warm_workspace_index(db: OakDatabase) { +fn warm_workspace_index(analysis: Analysis) { spawn_blocking(move || { let now = std::time::Instant::now(); lsp::log_info!("Starting workspace index warmup"); - indexer::warm(&db); + indexer::warm(analysis.read()); lsp::log_info!("Finished workspace index warmup ({:.0?})", now.elapsed()); Ok(None) }) @@ -1158,7 +1160,6 @@ fn warm_workspace_index(db: OakDatabase) { mod tests { use aether_path::FilePath; use oak_scan::DbScan; - use salsa::Database; use tower_lsp::jsonrpc; use url::Url; @@ -1222,7 +1223,7 @@ mod tests { respond( response_tx, || { - let _ = file.tree_sitter(&snapshot.db); + let _ = file.tree_sitter(snapshot.db.read()); Ok(LspResponse::Hover(None)) }, |response| response, diff --git a/crates/ark/src/lsp/state.rs b/crates/ark/src/lsp/state.rs index 251c0ec6c6..09a9a6ea74 100644 --- a/crates/ark/src/lsp/state.rs +++ b/crates/ark/src/lsp/state.rs @@ -8,6 +8,7 @@ use oak_semantic::library::Library; use url::Url; use crate::lsp::config::LspConfig; +use crate::lsp::db::Analysis; use crate::lsp::inputs::source_root::SourceRoot; use crate::lsp::open_file::OpenFile; @@ -15,6 +16,13 @@ use crate::lsp::open_file::OpenFile; /// The world state, i.e. all the inputs necessary for analysing or refactoring /// code. This is a pure value. There is no interior mutability in this data /// structure. It can be cloned and safely sent to other threads. +/// +/// The main loop owns and mutates this. Background readers get a +/// [`WorldSnapshot`] instead, which holds a read-only [`Analysis`] in place of +/// the writable `OakDatabase`. This prevents background threads from reaching a +/// Salsa input setter. See [`Self::diagnostics_snapshot`] and +/// [`Self::snapshot`]. This split mirrors rust-analyzer's `GlobalState` and +/// `GlobalStateSnapshot`. pub(crate) struct WorldState { /// Salsa input tree for Oak queries. pub(crate) db: OakDatabase, @@ -81,30 +89,45 @@ impl WorldState { } } - pub(crate) fn open_file_mut(&mut self, uri: &Url) -> anyhow::Result<&mut OpenFile> { - let key = FilePath::from_url(uri); - if let Some(open_file) = self.open_files.get_mut(&key) { - Ok(open_file) - } else { - Err(anyhow!("Can't find document for URI {uri}")) + /// Full read-only snapshot for a background reader that needs the whole + /// world (completions read `open_files` and the workspace). Same shape as + /// `self.clone()`, but the db becomes a read-only [`Analysis`]. + pub(crate) fn snapshot(&self) -> WorldSnapshot { + WorldSnapshot { + db: Analysis::new(self.db.clone()), + console_scopes: self.console_scopes.clone(), + installed_packages: self.installed_packages.clone(), + root: self.root.clone(), + library: self.library.clone(), + config: self.config.clone(), + open_files: self.open_files.clone(), + workspace: self.workspace.clone(), } } - /// Snapshot for the diagnostics worker, which runs off the main loop and - /// queries oak. - pub(crate) fn diagnostics_snapshot(&self) -> WorldState { - WorldState { - db: self.db.clone(), + /// Trimmed read-only snapshot for the diagnostics worker, which runs off + /// the main loop and queries oak. Drops the open-file / workspace maps the + /// diagnostics pass doesn't read. + pub(crate) fn diagnostics_snapshot(&self) -> WorldSnapshot { + WorldSnapshot { + db: Analysis::new(self.db.clone()), console_scopes: self.console_scopes.clone(), installed_packages: self.installed_packages.clone(), root: self.root.clone(), library: self.library.clone(), config: self.config.clone(), open_files: HashMap::new(), - virtual_documents: HashMap::new(), workspace: Workspace::default(), } } + pub(crate) fn open_file_mut(&mut self, uri: &Url) -> anyhow::Result<&mut OpenFile> { + let key = FilePath::from_url(uri); + if let Some(open_file) = self.open_files.get_mut(&key) { + Ok(open_file) + } else { + Err(anyhow!("Can't find document for URI {uri}")) + } + } /// The stored [`OpenFile`] for a request. pub(crate) fn open_file(&self, uri: &Url) -> anyhow::Result<&OpenFile> { @@ -140,6 +163,35 @@ impl WorldState { } } +/// Read-only snapshot of [`WorldState`] handed to background readers (e.g. +/// diagnostics). Holds a read-only [`Analysis`] in place of the writable +/// `OakDatabase`, so a reader thread can't reach Salsa input setters. Carries +/// only the fields readers actually use. Mirrors rust-analyzer's +/// `GlobalStateSnapshot`. +#[derive(Clone, Debug)] +pub(crate) struct WorldSnapshot { + pub(crate) db: Analysis, + pub(crate) open_files: HashMap, + pub(crate) workspace: Workspace, + pub(crate) console_scopes: Vec>, + pub(crate) installed_packages: Vec, + pub(crate) root: Option, + pub(crate) library: Library, + pub(crate) config: LspConfig, +} + +impl WorldSnapshot { + /// URL to put on the wire for `file`. Same rule as [`WorldState::wire_url`], + /// reading through the [`Analysis`] handle. + pub(crate) fn wire_url(&self, file: File) -> Url { + let path = file.path(self.db.read()); + self.open_files + .get(path) + .map(|open_file| open_file.wire_url().clone()) + .unwrap_or_else(|| path.to_url()) + } +} + pub(crate) fn open_file_wire_urls(state: &WorldState) -> Vec { state .open_files