[ty] Add virtual files to the only project database

Author: dhruvmanila

crates/ty_server/src/server/api.rs

@@ -236,14 +236,7 @@ where
             });
         };
 
-        let db = match &path {
-            AnySystemPath::System(path) => match session.project_db_for_path(path) {
-                Some(db) => db.clone(),
-                None => session.default_project_db().clone(),
-            },
-            AnySystemPath::SystemVirtual(_) => session.default_project_db().clone(),
-        };
-
+        let db = session.project_db(&path).clone();
         let snapshot = session.take_document_snapshot(url);
 
         Box::new(move |client| {

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

@@ -82,7 +82,7 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client:
         }
     };
 
-    let db = session.project_db_or_default(key.path());
+    let db = session.project_db(key.path());
 
     let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
         return;

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

@@ -40,16 +40,14 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             .update_text_document(&key, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        match key.path() {
+        let path = key.path();
+        let db = session.project_db_mut(path);
+
+        match path {
             AnySystemPath::System(path) => {
-                let db = match session.project_db_for_path_mut(path) {
-                    Some(db) => db,
-                    None => session.default_project_db_mut(),
-                };
                 db.apply_changes(vec![ChangeEvent::file_content_changed(path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.default_project_db_mut();
                 db.apply_changes(
                     vec![ChangeEvent::ChangedVirtual(virtual_path.clone())],
                     None,

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

@@ -38,8 +38,10 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
             .close_document(&key)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        if let AnySystemPath::SystemVirtual(virtual_path) = key.path() {
-            let db = session.default_project_db_mut();
+        let path = key.path();
+        let db = session.project_db_mut(path);
+
+        if let AnySystemPath::SystemVirtual(virtual_path) = path {
             db.apply_changes(
                 vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
                 None,

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

@@ -38,8 +38,10 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
             .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
-        if let AnySystemPath::SystemVirtual(virtual_path) = key.path() {
-            let db = session.default_project_db_mut();
+        let path = key.path();
+        let db = session.project_db_mut(path);
+
+        if let AnySystemPath::SystemVirtual(virtual_path) = path {
             db.apply_changes(
                 vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
                 None,

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

@@ -44,16 +44,14 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         let document = TextDocument::new(text, version).with_language_id(&language_id);
         session.open_text_document(key.path(), document);
 
-        match key.path() {
+        let path = key.path();
+        let db = session.project_db_mut(path);
+
+        match path {
             AnySystemPath::System(system_path) => {
-                let db = match session.project_db_for_path_mut(system_path) {
-                    Some(db) => db,
-                    None => session.default_project_db_mut(),
-                };
                 db.apply_changes(vec![ChangeEvent::Opened(system_path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.default_project_db_mut();
                 db.files().virtual_file(db, virtual_path);
             }
         }

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

@@ -38,16 +38,13 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
         .with_failure_code(ErrorCode::InternalError)?;
         session.open_notebook_document(&path, notebook);
 
+        let db = session.project_db_mut(&path);
+
         match &path {
             AnySystemPath::System(system_path) => {
-                let db = match session.project_db_for_path_mut(system_path) {
-                    Some(db) => db,
-                    None => session.default_project_db_mut(),
-                };
                 db.apply_changes(vec![ChangeEvent::Opened(system_path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.default_project_db_mut();
                 db.files().virtual_file(db, virtual_path);
             }
         }

crates/ty_server/src/session.rs

@@ -168,16 +168,47 @@ impl Session {
         &self.workspaces
     }
 
-    // TODO(dhruvmanila): Ideally, we should have a single method for `workspace_db_for_path_mut`
-    // and `default_workspace_db_mut` but the borrow checker doesn't allow that.
-    // https://github.com/astral-sh/ruff/pull/13041#discussion_r1726725437
+    /// Returns a reference to the project's [`ProjectDatabase`] in which the given `path` belongs.
+    ///
+    /// If the path is a system path, it will return the project database that is closest to the
+    /// given path, or the default project if no project is found for the path.
+    ///
+    /// If the path is a virtual path, it will return the first project database in the session.
+    pub(crate) fn project_db(&self, path: &AnySystemPath) -> &ProjectDatabase {
+        match path {
+            AnySystemPath::System(system_path) => self
+                .project_db_for_path(system_path)
+                .unwrap_or(&self.default_project),
+            AnySystemPath::SystemVirtual(_virtual_path) => {
+                // TODO: Currently, ty only supports single workspace but we need to figure out
+                // which project should this virtual path belong to when there are multiple
+                // projects: https://github.com/astral-sh/ty/issues/794
+                self.projects.iter().next().map(|(_, db)| db).unwrap()
+            }
+        }
+    }
 
-    /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path,
-    /// or the default project if no project is found for the path.
-    pub(crate) fn project_db_or_default(&self, path: &AnySystemPath) -> &ProjectDatabase {
-        path.as_system()
-            .and_then(|path| self.project_db_for_path(path))
-            .unwrap_or_else(|| self.default_project_db())
+    /// Returns a mutable reference to the project's [`ProjectDatabase`] in which the given `path`
+    /// belongs.
+    ///
+    /// Refer to [`project_db`] for more details on how the project is selected.
+    ///
+    /// [`project_db`]: Session::project_db
+    pub(crate) fn project_db_mut(&mut self, path: &AnySystemPath) -> &mut ProjectDatabase {
+        match path {
+            AnySystemPath::System(system_path) => self
+                .projects
+                .range_mut(..=system_path.to_path_buf())
+                .next_back()
+                .map(|(_, db)| db)
+                .unwrap_or(&mut self.default_project),
+            AnySystemPath::SystemVirtual(_virtual_path) => {
+                // TODO: Currently, ty only supports single workspace but we need to figure out
+                // which project should this virtual path belong to when there are multiple
+                // projects: https://github.com/astral-sh/ty/issues/794
+                self.projects.iter_mut().next().map(|(_, db)| db).unwrap()
+            }
+        }
     }
 
     /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path, if
@@ -204,17 +235,6 @@ impl Session {
             .map(|(_, db)| db)
     }
 
-    /// Returns a reference to the default project [`ProjectDatabase`]. The default project is the
-    /// minimum root path in the project map.
-    pub(crate) fn default_project_db(&self) -> &ProjectDatabase {
-        &self.default_project
-    }
-
-    /// Returns a mutable reference to the default project [`ProjectDatabase`].
-    pub(crate) fn default_project_db_mut(&mut self) -> &mut ProjectDatabase {
-        &mut self.default_project
-    }
-
     fn projects_mut(&mut self) -> impl Iterator<Item = &'_ mut ProjectDatabase> + '_ {
         self.projects
             .values_mut()

crates/ty_server/src/system.rs

@@ -72,13 +72,6 @@ impl AnySystemPath {
         }
     }
 
-    pub(crate) const fn as_system(&self) -> Option<&SystemPathBuf> {
-        match self {
-            AnySystemPath::System(system_path_buf) => Some(system_path_buf),
-            AnySystemPath::SystemVirtual(_) => None,
-        }
-    }
-
     /// Returns the extension of the path, if any.
     pub(crate) fn extension(&self) -> Option<&str> {
         match self {