[ty] Add `ty_server::Db` trait (#21241)

crates/ty/src/lib.rs

@@ -450,12 +450,12 @@ impl ty_project::ProgressReporter for IndicatifReporter {
         self.bar.set_draw_target(self.printer.progress_target());
     }
 
-    fn report_checked_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]) {
+    fn report_checked_file(&self, db: &ProjectDatabase, file: File, diagnostics: &[Diagnostic]) {
         self.collector.report_checked_file(db, file, diagnostics);
         self.bar.inc(1);
     }
 
-    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+    fn report_diagnostics(&mut self, db: &ProjectDatabase, diagnostics: Vec<Diagnostic>) {
         self.collector.report_diagnostics(db, diagnostics);
     }
 }

crates/ty_project/src/lib.rs

@@ -124,12 +124,12 @@ pub trait ProgressReporter: Send + Sync {
     fn set_files(&mut self, files: usize);
 
     /// Report the completion of checking a given file along with its diagnostics.
-    fn report_checked_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]);
+    fn report_checked_file(&self, db: &ProjectDatabase, file: File, diagnostics: &[Diagnostic]);
 
     /// Reports settings or IO related diagnostics. The diagnostics
     /// can belong to different files or no file at all.
     /// But it's never a file for which [`Self::report_checked_file`] gets called.
-    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>);
+    fn report_diagnostics(&mut self, db: &ProjectDatabase, diagnostics: Vec<Diagnostic>);
 }
 
 /// Reporter that collects all diagnostics into a `Vec`.
@@ -149,7 +149,7 @@ impl CollectReporter {
 
 impl ProgressReporter for CollectReporter {
     fn set_files(&mut self, _files: usize) {}
-    fn report_checked_file(&self, _db: &dyn Db, _file: File, diagnostics: &[Diagnostic]) {
+    fn report_checked_file(&self, _db: &ProjectDatabase, _file: File, diagnostics: &[Diagnostic]) {
         if diagnostics.is_empty() {
             return;
         }
@@ -160,7 +160,7 @@ impl ProgressReporter for CollectReporter {
             .extend(diagnostics.iter().map(Clone::clone));
     }
 
-    fn report_diagnostics(&mut self, _db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+    fn report_diagnostics(&mut self, _db: &ProjectDatabase, diagnostics: Vec<Diagnostic>) {
         self.0.get_mut().unwrap().extend(diagnostics);
     }
 }

crates/ty_server/src/db.rs

@@ -0,0 +1,33 @@
+use crate::NotebookDocument;
+use crate::session::index::Document;
+use crate::system::LSPSystem;
+use ruff_db::Db as _;
+use ruff_db::files::{File, FilePath};
+use ty_project::{Db as ProjectDb, ProjectDatabase};
+
+#[salsa::db]
+pub(crate) trait Db: ProjectDb {
+    /// Returns the LSP [`Document`] corresponding to `File` or
+    /// `None` if the file isn't open in the editor.
+    fn document(&self, file: File) -> Option<&Document>;
+
+    /// Returns the LSP [`NotebookDocument`] corresponding to `File` or
+    /// `None` if the file isn't open in the editor or if it isn't a notebook.
+    fn notebook_document(&self, file: File) -> Option<&NotebookDocument> {
+        self.document(file)?.as_notebook()
+    }
+}
+
+#[salsa::db]
+impl Db for ProjectDatabase {
+    fn document(&self, file: File) -> Option<&Document> {
+        self.system()
+            .as_any()
+            .downcast_ref::<LSPSystem>()
+            .and_then(|system| match file.path(self) {
+                FilePath::System(path) => system.system_path_to_document(path),
+                FilePath::SystemVirtual(path) => system.system_virtual_path_to_document(path),
+                FilePath::Vendored(_) => None,
+            })
+    }
+}

crates/ty_server/src/document/location.rs

@@ -1,9 +1,10 @@
-use crate::PositionEncoding;
-use crate::document::{FileRangeExt, ToRangeExt};
 use lsp_types::Location;
 use ruff_db::files::FileRange;
 use ty_ide::{NavigationTarget, ReferenceTarget};
-use ty_python_semantic::Db;
+
+use crate::Db;
+use crate::PositionEncoding;
+use crate::document::{FileRangeExt, ToRangeExt};
 
 pub(crate) trait ToLink {
     fn to_location(&self, db: &dyn Db, encoding: PositionEncoding) -> Option<Location>;

crates/ty_server/src/document/range.rs

@@ -1,6 +1,6 @@
 use super::PositionEncoding;
+use crate::Db;
 use crate::system::file_to_url;
-use ty_python_semantic::Db;
 
 use lsp_types as types;
 use lsp_types::{Location, Position, Url};

crates/ty_server/src/lib.rs

@@ -4,6 +4,7 @@ use anyhow::Context;
 use lsp_server::Connection;
 use ruff_db::system::{OsSystem, SystemPathBuf};
 
+use crate::db::Db;
 pub use crate::logging::{LogLevel, init_logging};
 pub use crate::server::{PartialWorkspaceProgress, PartialWorkspaceProgressParams, Server};
 pub use crate::session::{ClientOptions, DiagnosticMode};
@@ -11,6 +12,7 @@ pub use document::{NotebookDocument, PositionEncoding, TextDocument};
 pub(crate) use session::Session;
 
 mod capabilities;
+mod db;
 mod document;
 mod logging;
 mod server;

crates/ty_server/src/server/api/diagnostics.rs

@@ -10,8 +10,9 @@ use rustc_hash::FxHashMap;
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
 use ruff_db::files::FileRange;
 use ruff_db::system::SystemPathBuf;
-use ty_project::{Db, ProjectDatabase};
+use ty_project::{Db as _, ProjectDatabase};
 
+use crate::Db;
 use crate::document::{FileRangeExt, ToRangeExt};
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;

crates/ty_server/src/server/api/notifications/did_change_watched_files.rs

@@ -8,7 +8,7 @@ use crate::system::AnySystemPath;
 use lsp_types as types;
 use lsp_types::{FileChangeType, notification as notif};
 use rustc_hash::FxHashMap;
-use ty_project::Db;
+use ty_project::Db as _;
 use ty_project::watch::{ChangeEvent, ChangedKind, CreatedKind, DeletedKind};
 
 pub(crate) struct DidChangeWatchedFiles;

crates/ty_server/src/server/api/requests/document_symbols.rs

@@ -4,8 +4,9 @@ use lsp_types::request::DocumentSymbolRequest;
 use lsp_types::{DocumentSymbol, DocumentSymbolParams, SymbolInformation, Url};
 use ruff_db::files::File;
 use ty_ide::{HierarchicalSymbols, SymbolId, SymbolInfo, document_symbols};
-use ty_project::{Db, ProjectDatabase};
+use ty_project::ProjectDatabase;
 
+use crate::Db;
 use crate::document::{PositionEncoding, ToRangeExt};
 use crate::server::api::symbols::{convert_symbol_kind, convert_to_lsp_symbol_information};
 use crate::server::api::traits::{

crates/ty_server/src/server/api/requests/execute_command.rs

@@ -9,7 +9,7 @@ use lsp_server::ErrorCode;
 use lsp_types::{self as types, request as req};
 use std::fmt::Write;
 use std::str::FromStr;
-use ty_project::Db;
+use ty_project::Db as _;
 
 pub(crate) struct ExecuteCommand;
 

crates/ty_server/src/server/api/requests/workspace_diagnostic.rs

@@ -26,7 +26,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::sync::Mutex;
 use std::time::{Duration, Instant};
-use ty_project::{Db, ProgressReporter};
+use ty_project::{ProgressReporter, ProjectDatabase};
 
 /// Handler for [Workspace diagnostics](workspace-diagnostics)
 ///
@@ -230,7 +230,7 @@ impl ProgressReporter for WorkspaceDiagnosticsProgressReporter<'_> {
         state.report_progress(&self.work_done);
     }
 
-    fn report_checked_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]) {
+    fn report_checked_file(&self, db: &ProjectDatabase, file: File, diagnostics: &[Diagnostic]) {
         // Another thread might have panicked at this point because of a salsa cancellation which
         // poisoned the result. If the response is poisoned, just don't report and wait for our thread
         // to unwind with a salsa cancellation next.
@@ -260,7 +260,7 @@ impl ProgressReporter for WorkspaceDiagnosticsProgressReporter<'_> {
         state.response.maybe_flush();
     }
 
-    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+    fn report_diagnostics(&mut self, db: &ProjectDatabase, diagnostics: Vec<Diagnostic>) {
         let mut by_file: BTreeMap<File, Vec<Diagnostic>> = BTreeMap::new();
 
         for diagnostic in diagnostics {
@@ -358,7 +358,12 @@ impl<'a> ResponseWriter<'a> {
         }
     }
 
-    fn write_diagnostics_for_file(&mut self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]) {
+    fn write_diagnostics_for_file(
+        &mut self,
+        db: &ProjectDatabase,
+        file: File,
+        diagnostics: &[Diagnostic],
+    ) {
         let Some(url) = file_to_url(db, file) else {
             tracing::debug!("Failed to convert file path to URL at {}", file.path(db));
             return;

crates/ty_server/src/server/api/symbols.rs

@@ -3,8 +3,8 @@
 
 use lsp_types::{SymbolInformation, SymbolKind};
 use ty_ide::SymbolInfo;
-use ty_project::Db;
 
+use crate::Db;
 use crate::document::{PositionEncoding, ToRangeExt};
 
 /// Convert `ty_ide` `SymbolKind` to LSP `SymbolKind`

crates/ty_server/src/system.rs

@@ -4,6 +4,7 @@ use std::fmt::Display;
 use std::panic::RefUnwindSafe;
 use std::sync::Arc;
 
+use crate::Db;
 use crate::document::DocumentKey;
 use crate::session::index::{Document, Index};
 use lsp_types::Url;
@@ -16,7 +17,6 @@ use ruff_db::system::{
 };
 use ruff_notebook::{Notebook, NotebookError};
 use ty_ide::cached_vendored_path;
-use ty_python_semantic::Db;
 
 /// Returns a [`Url`] for the given [`File`].
 pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
@@ -112,25 +112,28 @@ impl LSPSystem {
         self.index.as_ref().unwrap()
     }
 
-    fn make_document_ref(&self, path: AnySystemPath) -> Option<&Document> {
+    fn document(&self, path: AnySystemPath) -> Option<&Document> {
         let index = self.index();
         index.document(&DocumentKey::from(path)).ok()
     }
 
-    fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<&Document> {
+    pub(crate) fn system_path_to_document(&self, path: &SystemPath) -> Option<&Document> {
         let any_path = AnySystemPath::System(path.to_path_buf());
-        self.make_document_ref(any_path)
+        self.document(any_path)
     }
 
-    fn system_virtual_path_to_document_ref(&self, path: &SystemVirtualPath) -> Option<&Document> {
+    pub(crate) fn system_virtual_path_to_document(
+        &self,
+        path: &SystemVirtualPath,
+    ) -> Option<&Document> {
         let any_path = AnySystemPath::SystemVirtual(path.to_path_buf());
-        self.make_document_ref(any_path)
+        self.document(any_path)
     }
 }
 
 impl System for LSPSystem {
     fn path_metadata(&self, path: &SystemPath) -> Result<Metadata> {
-        let document = self.system_path_to_document_ref(path);
+        let document = self.system_path_to_document(path);
 
         if let Some(document) = document {
             Ok(Metadata::new(
@@ -152,7 +155,7 @@ impl System for LSPSystem {
     }
 
     fn read_to_string(&self, path: &SystemPath) -> Result<String> {
-        let document = self.system_path_to_document_ref(path);
+        let document = self.system_path_to_document(path);
 
         match document {
             Some(Document::Text(document)) => Ok(document.contents().to_string()),
@@ -161,7 +164,7 @@ impl System for LSPSystem {
     }
 
     fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError> {
-        let document = self.system_path_to_document_ref(path);
+        let document = self.system_path_to_document(path);
 
         match document {
             Some(Document::Text(document)) => Notebook::from_source_code(document.contents()),
@@ -172,7 +175,7 @@ impl System for LSPSystem {
 
     fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
         let document = self
-            .system_virtual_path_to_document_ref(path)
+            .system_virtual_path_to_document(path)
             .ok_or_else(|| virtual_path_not_found(path))?;
 
         if let Document::Text(document) = &document {
@@ -187,7 +190,7 @@ impl System for LSPSystem {
         path: &SystemVirtualPath,
     ) -> std::result::Result<Notebook, NotebookError> {
         let document = self
-            .system_virtual_path_to_document_ref(path)
+            .system_virtual_path_to_document(path)
             .ok_or_else(|| virtual_path_not_found(path))?;
 
         match document {