From 08b0795e51e6f02650aa0a5aec1e6c4569e15301 Mon Sep 17 00:00:00 2001
From: Piotr Figiela <77412592+Draggu@users.noreply.github.com>
Date: Thu, 28 Nov 2024 17:45:18 +0100
Subject: [PATCH] Restart proc macro server on `Scarb.toml` change

commit-id:cf55e524
---
 src/lang/db/mod.rs                 |  7 +++++++
 src/lang/proc_macros/controller.rs | 21 +++++++++++++++++++++
 src/server/routing/traits.rs       |  2 ++
 3 files changed, 30 insertions(+)

diff --git a/src/lang/db/mod.rs b/src/lang/db/mod.rs
index c81b6995..7edfc65e 100644
--- a/src/lang/db/mod.rs
+++ b/src/lang/db/mod.rs
@@ -126,6 +126,13 @@ impl AnalysisDatabase {
         })
     }
 
+    /// Removes plugin suit from database.
+    pub fn remove_plugin_suite(&mut self, plugins: PluginSuite) {
+        self.with_plugins_mut(move |macro_plugins, analyzer_plugins, inline_macro_plugins| {
+            remove_plugin_suite(plugins, macro_plugins, analyzer_plugins, inline_macro_plugins)
+        })
+    }
+
     /// Adds plugin suit to database.
     pub fn add_plugin_suite(&mut self, plugins: PluginSuite) {
         self.with_plugins_mut(move |macro_plugins, analyzer_plugins, inline_macro_plugins| {
diff --git a/src/lang/proc_macros/controller.rs b/src/lang/proc_macros/controller.rs
index 535f37a4..b5851fe3 100644
--- a/src/lang/proc_macros/controller.rs
+++ b/src/lang/proc_macros/controller.rs
@@ -88,6 +88,21 @@ impl ProcMacroClientController {
         }
     }
 
+    pub fn force_restart(&mut self, db: &mut AnalysisDatabase, config: &Config) {
+        // We have to make sure that snapshots will not report errors from previous client after we
+        // create new one.
+        db.cancel_all();
+
+        // Otherwise we can get messages from old client after initialization of new one.
+        self.channels.clear_all();
+
+        if let Some(plugin_suite) = self.plugin_suite.take() {
+            db.remove_plugin_suite(plugin_suite);
+        }
+
+        self.try_initialize(db, config);
+    }
+
     /// Check if an error was reported. If so, try to restart.
     pub fn handle_error(&mut self, db: &mut AnalysisDatabase, config: &Config) {
         if !self.try_initialize(db, config) {
@@ -272,4 +287,10 @@ impl ProcMacroChannels {
 
         Self { response_sender, response_receiver, error_sender, error_receiver }
     }
+
+    /// Make all channels empty in a non-blocking manner.
+    fn clear_all(&self) {
+        self.error_receiver.try_iter().for_each(|_| {});
+        self.response_receiver.try_iter().for_each(|_| {});
+    }
 }
diff --git a/src/server/routing/traits.rs b/src/server/routing/traits.rs
index bfe05f51..92362f10 100644
--- a/src/server/routing/traits.rs
+++ b/src/server/routing/traits.rs
@@ -198,6 +198,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
             if ["Scarb.toml", "cairo_project.toml"].map(Some).contains(&changed_file_name.to_str())
             {
                 Backend::reload(state, requester)?;
+
+                state.proc_macro_controller.force_restart(&mut state.db, &state.config);
             }
         }