[red-knot] Support textDocument/didChange notification (#13042)

## Summary

This PR adds support for `textDocument/didChange` notification.

There seems to be a bug (probably in Salsa) where it panics with:
```
2024-08-22 15:33:38.802 [info] panicked at /Users/dhruv/.cargo/git/checkouts/salsa-61760caba2b17ca5/f608ff8/src/tracked_struct.rs:377:9:
two concurrent writers to Id(4800), should not be possible
```

## Test Plan


https://github.com/user-attachments/assets/81055feb-ba8e-4acf-ad2f-94084a3efead

crates/red_knot_server/src/server.rs

@@ -6,7 +6,8 @@ use std::panic::PanicInfo;
 use lsp_server::Message;
 use lsp_types::{
     ClientCapabilities, DiagnosticOptions, DiagnosticServerCapabilities, MessageType,
-    ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncOptions, Url,
+    ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
+    Url,
 };
 
 use self::connection::{Connection, ConnectionInitializer};
@@ -220,6 +221,7 @@ impl Server {
             text_document_sync: Some(TextDocumentSyncCapability::Options(
                 TextDocumentSyncOptions {
                     open_close: Some(true),
+                    change: Some(TextDocumentSyncKind::INCREMENTAL),
                     ..Default::default()
                 },
             )),

crates/red_knot_server/src/server/api.rs

@@ -42,6 +42,7 @@ pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> {
     match notif.method.as_str() {
         notification::DidCloseTextDocumentHandler::METHOD => local_notification_task::<notification::DidCloseTextDocumentHandler>(notif),
         notification::DidOpenTextDocumentHandler::METHOD => local_notification_task::<notification::DidOpenTextDocumentHandler>(notif),
+        notification::DidChangeTextDocumentHandler::METHOD => local_notification_task::<notification::DidChangeTextDocumentHandler>(notif),
         notification::DidOpenNotebookHandler::METHOD => {
             local_notification_task::<notification::DidOpenNotebookHandler>(notif)
         }

crates/red_knot_server/src/server/api/notifications.rs

@@ -1,9 +1,11 @@
+mod did_change;
 mod did_close;
 mod did_close_notebook;
 mod did_open;
 mod did_open_notebook;
 mod set_trace;
 
+pub(super) use did_change::DidChangeTextDocumentHandler;
 pub(super) use did_close::DidCloseTextDocumentHandler;
 pub(super) use did_close_notebook::DidCloseNotebookHandler;
 pub(super) use did_open::DidOpenTextDocumentHandler;

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

@@ -0,0 +1,47 @@
+use lsp_server::ErrorCode;
+use lsp_types::notification::DidChangeTextDocument;
+use lsp_types::DidChangeTextDocumentParams;
+
+use red_knot_workspace::watch::ChangeEvent;
+
+use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
+use crate::server::api::LSPResult;
+use crate::server::client::{Notifier, Requester};
+use crate::server::Result;
+use crate::session::Session;
+use crate::system::url_to_system_path;
+
+pub(crate) struct DidChangeTextDocumentHandler;
+
+impl NotificationHandler for DidChangeTextDocumentHandler {
+    type NotificationType = DidChangeTextDocument;
+}
+
+impl SyncNotificationHandler for DidChangeTextDocumentHandler {
+    fn run(
+        session: &mut Session,
+        _notifier: Notifier,
+        _requester: &mut Requester,
+        params: DidChangeTextDocumentParams,
+    ) -> Result<()> {
+        let Ok(path) = url_to_system_path(&params.text_document.uri) else {
+            return Ok(());
+        };
+
+        let key = session.key_from_url(params.text_document.uri);
+
+        session
+            .update_text_document(&key, params.content_changes, params.text_document.version)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        let db = match session.workspace_db_for_path_mut(path.as_std_path()) {
+            Some(db) => db,
+            None => session.default_workspace_db_mut(),
+        };
+        db.apply_changes(vec![ChangeEvent::file_content_changed(path)], None);
+
+        // TODO(dhruvmanila): Publish diagnostics if the client doesnt support pull diagnostics
+
+        Ok(())
+    }
+}

crates/red_knot_server/src/server/api/requests/diagnostic.rs

@@ -46,10 +46,19 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
 
 fn compute_diagnostics(snapshot: &DocumentSnapshot, db: &RootDatabase) -> Vec<Diagnostic> {
     let Some(file) = snapshot.file(db) else {
+        tracing::info!(
+            "No file found for snapshot for '{}'",
+            snapshot.query().file_url()
+        );
         return vec![];
     };
-    let Ok(diagnostics) = db.check_file(file) else {
+
+    let diagnostics = match db.check_file(file) {
+        Ok(diagnostics) => diagnostics,
+        Err(cancelled) => {
+            tracing::info!("Diagnostics computation {cancelled}");
             return vec![];
+        }
     };
 
     diagnostics

crates/red_knot_server/src/session.rs

@@ -6,14 +6,14 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
 use anyhow::anyhow;
-use lsp_types::{ClientCapabilities, Url};
+use lsp_types::{ClientCapabilities, TextDocumentContentChangeEvent, Url};
 
 use red_knot_workspace::db::RootDatabase;
 use red_knot_workspace::workspace::WorkspaceMetadata;
 use ruff_db::files::{system_path_to_file, File};
 use ruff_db::system::SystemPath;
 
-use crate::edit::{DocumentKey, NotebookDocument};
+use crate::edit::{DocumentKey, DocumentVersion, NotebookDocument};
 use crate::system::{url_to_system_path, LSPSystem};
 use crate::{PositionEncoding, TextDocument};
 
@@ -146,6 +146,20 @@ impl Session {
         self.index_mut().open_text_document(url, document);
     }
 
+    /// Updates a text document at the associated `key`.
+    ///
+    /// The document key must point to a text document, or this will throw an error.
+    pub(crate) fn update_text_document(
+        &mut self,
+        key: &DocumentKey,
+        content_changes: Vec<TextDocumentContentChangeEvent>,
+        new_version: DocumentVersion,
+    ) -> crate::Result<()> {
+        let position_encoding = self.position_encoding;
+        self.index_mut()
+            .update_text_document(key, content_changes, new_version, position_encoding)
+    }
+
     /// De-registers a document, specified by its key.
     /// Calling this multiple times for the same document is a logic error.
     pub(crate) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {

crates/red_knot_workspace/src/watch.rs

@@ -49,6 +49,16 @@ pub enum ChangeEvent {
 }
 
 impl ChangeEvent {
+    /// Creates a new [`Changed`] event for the file content at the given path.
+    ///
+    /// [`Changed`]: ChangeEvent::Changed
+    pub fn file_content_changed(path: SystemPathBuf) -> ChangeEvent {
+        ChangeEvent::Changed {
+            path,
+            kind: ChangedKind::FileContent,
+        }
+    }
+
     pub fn file_name(&self) -> Option<&str> {
         self.path().and_then(|path| path.file_name())
     }