From 1a98ea34952375e8a6faf901fb5797c6294bc342 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Fri, 5 Jul 2024 16:14:39 +0530
Subject: [PATCH] Consider the cell content during notebook sync

---
 crates/ruff_server/src/edit/notebook.rs | 34 ++++++++++++++++++++++---
 crates/ruff_server/src/session/index.rs |  8 +++---
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/crates/ruff_server/src/edit/notebook.rs b/crates/ruff_server/src/edit/notebook.rs
index 52a88d0c44c55..686449d6b54f4 100644
--- a/crates/ruff_server/src/edit/notebook.rs
+++ b/crates/ruff_server/src/edit/notebook.rs
@@ -109,24 +109,47 @@ impl NotebookDocument {
             text_content,
         }) = cells
         {
+            // The structural changes should be done first, as they may affect the cell index.
             if let Some(structure) = structure {
                 let start = structure.array.start as usize;
                 let delete = structure.array.delete_count as usize;
+
+                // First, delete the cells and remove them from the index.
                 if delete > 0 {
                     for cell in self.cells.drain(start..start + delete) {
                         self.cell_index.remove(&cell.url);
                     }
                 }
+
+                // Second, insert the new cells with the available information. This array does not
+                // provide the actual contents of the cells, so we'll initialize them with empty
+                // contents.
                 for cell in structure.array.cells.into_iter().flatten().rev() {
                     self.cells
-                        .insert(start, NotebookCell::new(cell, String::new(), version));
+                        .insert(start, NotebookCell::new(cell, String::new(), 0));
                 }
 
-                // register any new cells in the index and update existing ones that came after the insertion
-                for (i, cell) in self.cells.iter().enumerate().skip(start) {
-                    self.cell_index.insert(cell.url.clone(), i);
+                // Third, register the new cells in the index and update existing ones that came
+                // after the insertion.
+                for (index, cell) in self.cells.iter().enumerate().skip(start) {
+                    self.cell_index.insert(cell.url.clone(), index);
+                }
+
+                // Finally, update the text document that represents the cell with the actual
+                // contents. This should be done at the end so that both the `cells` and
+                // `cell_index` are updated before we start applying the changes to the cells.
+                if let Some(did_open) = structure.did_open {
+                    for cell_text_document in did_open {
+                        if let Some(cell) = self.cell_by_uri_mut(&cell_text_document.uri) {
+                            cell.document = TextDocument::new(
+                                cell_text_document.text,
+                                cell_text_document.version,
+                            );
+                        }
+                    }
                 }
             }
+
             if let Some(cell_data) = data {
                 for cell in cell_data {
                     if let Some(existing_cell) = self.cell_by_uri_mut(&cell.document) {
@@ -134,6 +157,7 @@ impl NotebookDocument {
                     }
                 }
             }
+
             if let Some(content_changes) = text_content {
                 for content_change in content_changes {
                     if let Some(cell) = self.cell_by_uri_mut(&content_change.document.uri) {
@@ -143,9 +167,11 @@ impl NotebookDocument {
                 }
             }
         }
+
         if let Some(metadata_change) = metadata_change {
             self.metadata = serde_json::from_value(serde_json::Value::Object(metadata_change))?;
         }
+
         Ok(())
     }
 
diff --git a/crates/ruff_server/src/session/index.rs b/crates/ruff_server/src/session/index.rs
index 64e6333a071c9..502356c1ac3c4 100644
--- a/crates/ruff_server/src/session/index.rs
+++ b/crates/ruff_server/src/session/index.rs
@@ -145,14 +145,16 @@ impl Index {
         encoding: PositionEncoding,
     ) -> crate::Result<()> {
         // update notebook cell index
-        if let Some(lsp_types::NotebookDocumentCellChangeStructure { did_open, .. }) =
-            cells.as_ref().and_then(|cells| cells.structure.as_ref())
+        if let Some(lsp_types::NotebookDocumentCellChangeStructure {
+            did_open: Some(did_open),
+            ..
+        }) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
         {
             let Some(path) = self.url_for_key(key).cloned() else {
                 anyhow::bail!("Tried to open unavailable document `{key}`");
             };
 
-            for opened_cell in did_open.iter().flatten() {
+            for opened_cell in did_open {
                 self.notebook_cells
                     .insert(opened_cell.uri.clone(), path.clone());
             }