From 05a22493597b4638137c545e177ac3b80efbc011 Mon Sep 17 00:00:00 2001
From: Piotr Magiera <56825108+piotmag769@users.noreply.github.com>
Date: Thu, 12 Dec 2024 14:14:49 +0100
Subject: [PATCH] Make updating diagnostics atomic (#36)

---
 src/lang/diagnostics/file_diagnostics.rs | 19 +++++++---------
 src/lang/diagnostics/mod.rs              |  2 +-
 src/lang/diagnostics/refresh.rs          | 29 ++++++++++++++++--------
 3 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/src/lang/diagnostics/file_diagnostics.rs b/src/lang/diagnostics/file_diagnostics.rs
index 5636fd61..2ea1d98f 100644
--- a/src/lang/diagnostics/file_diagnostics.rs
+++ b/src/lang/diagnostics/file_diagnostics.rs
@@ -113,40 +113,37 @@ impl FileDiagnostics {
     }
 
     /// Constructs a new [`lsp_types::PublishDiagnosticsParams`] from this `FileDiagnostics`.
+    ///
+    /// NOTE: `file_id` must correspond to the url at the current time slice.
     pub fn to_lsp(
         &self,
         db: &AnalysisDatabase,
+        file_id: FileId,
         trace_macro_diagnostics: bool,
-    ) -> Option<lsp_types::PublishDiagnosticsParams> {
-        let file = db.file_for_url(&self.url)?;
-
+    ) -> lsp_types::PublishDiagnosticsParams {
         let mut diagnostics = Vec::new();
         map_cairo_diagnostics_to_lsp(
             (*db).upcast(),
             &mut diagnostics,
             &self.parser,
-            file,
+            file_id,
             trace_macro_diagnostics,
         );
         map_cairo_diagnostics_to_lsp(
             (*db).upcast(),
             &mut diagnostics,
             &self.semantic,
-            file,
+            file_id,
             trace_macro_diagnostics,
         );
         map_cairo_diagnostics_to_lsp(
             (*db).upcast(),
             &mut diagnostics,
             &self.lowering,
-            file,
+            file_id,
             trace_macro_diagnostics,
         );
 
-        Some(lsp_types::PublishDiagnosticsParams {
-            uri: self.url.clone(),
-            diagnostics,
-            version: None,
-        })
+        lsp_types::PublishDiagnosticsParams { uri: self.url.clone(), diagnostics, version: None }
     }
 }
diff --git a/src/lang/diagnostics/mod.rs b/src/lang/diagnostics/mod.rs
index 7e783d25..0828b64c 100644
--- a/src/lang/diagnostics/mod.rs
+++ b/src/lang/diagnostics/mod.rs
@@ -124,7 +124,7 @@ impl DiagnosticsControllerThread {
             .collect();
 
         self.spawn_worker(move |project_diagnostics, notifier| {
-            clear_old_diagnostics(&state.db, files_to_preserve, project_diagnostics, notifier);
+            clear_old_diagnostics(files_to_preserve, project_diagnostics, notifier);
         });
     }
 
diff --git a/src/lang/diagnostics/refresh.rs b/src/lang/diagnostics/refresh.rs
index ba68a770..5e732950 100644
--- a/src/lang/diagnostics/refresh.rs
+++ b/src/lang/diagnostics/refresh.rs
@@ -2,8 +2,8 @@ use std::collections::HashSet;
 
 use cairo_lang_defs::ids::ModuleId;
 use cairo_lang_filesystem::ids::FileId;
-use lsp_types::Url;
 use lsp_types::notification::PublishDiagnostics;
+use lsp_types::{PublishDiagnosticsParams, Url};
 
 use crate::lang::db::AnalysisDatabase;
 use crate::lang::diagnostics::file_diagnostics::FileDiagnostics;
@@ -35,6 +35,10 @@ pub fn refresh_diagnostics(
 }
 
 /// Refresh diagnostics for a single file.
+///
+/// IMPORTANT: keep updating diagnostics state between server and client ATOMIC!
+/// I.e, if diagnostics are updated on the server side they MUST be sent successfully to the
+/// client (and vice-versa).
 #[tracing::instrument(skip_all, fields(url = tracing_file_url(db, file)))]
 fn refresh_file_diagnostics(
     db: &AnalysisDatabase,
@@ -48,17 +52,21 @@ fn refresh_file_diagnostics(
         return;
     };
 
+    // IMPORTANT: DO NOT change the order of operations here. `to_lsp` may panic, so it has to come
+    // before `insert`. It is to make sure that if `insert` succeeds, `notify` executes as well.
+    let params = new_file_diagnostics.to_lsp(db, file, trace_macro_diagnostics);
     if project_diagnostics.insert(new_file_diagnostics.clone()) {
-        if let Some(params) = new_file_diagnostics.to_lsp(db, trace_macro_diagnostics) {
-            notifier.notify::<PublishDiagnostics>(params);
-        }
+        notifier.notify::<PublishDiagnostics>(params);
     }
 }
 
 /// Wipes diagnostics for any files not present in the preserve set.
+///
+/// IMPORTANT: keep updating diagnostics state between server and client ATOMIC!
+/// I.e, if diagnostics are updated on the server side they MUST be sent successfully to the
+/// client (and vice-versa).
 #[tracing::instrument(skip_all)]
 pub fn clear_old_diagnostics(
-    db: &AnalysisDatabase,
     files_to_preserve: HashSet<Url>,
     project_diagnostics: ProjectDiagnostics,
     notifier: Notifier,
@@ -71,11 +79,12 @@ pub fn clear_old_diagnostics(
         // to preserve any extra state it might contain.
         file_diagnostics.clear();
 
-        // We can safely assume `trace_macro_diagnostics` = false here, as we are explicitly
-        // sending a "no diagnostics" message.
-        if let Some(params) = file_diagnostics.to_lsp(db, false) {
-            notifier.notify::<PublishDiagnostics>(params);
-        }
+        let params = PublishDiagnosticsParams {
+            uri: file_diagnostics.url,
+            diagnostics: vec![],
+            version: None,
+        };
+        notifier.notify::<PublishDiagnostics>(params);
     }
 }