From 6e02ee734e2fb623f97746df88650f0976865775 Mon Sep 17 00:00:00 2001
From: He1pa <56333845+He1pa@users.noreply.github.com>
Date: Tue, 17 Oct 2023 16:00:25 +0800
Subject: [PATCH] Fix: lsp crash caused by vfs loader (#778)

* bugfix: fix lsp crash caused by vfs loader

Signed-off-by: He1pa <18012015693@163.com>

* test: add open-close file mock test for vfs

Signed-off-by: He1pa <18012015693@163.com>

---------

Signed-off-by: He1pa <18012015693@163.com>
---
 kclvm/tools/src/LSP/src/notification.rs |  4 +--
 kclvm/tools/src/LSP/src/state.rs        | 21 ++++++++----
 kclvm/tools/src/LSP/src/tests.rs        | 45 +++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/kclvm/tools/src/LSP/src/notification.rs b/kclvm/tools/src/LSP/src/notification.rs
index 3830968ad..605f8497e 100644
--- a/kclvm/tools/src/LSP/src/notification.rs
+++ b/kclvm/tools/src/LSP/src/notification.rs
@@ -115,7 +115,7 @@ impl LanguageServerState {
         if let Some(id) = self.vfs.read().file_id(&path.clone().into()) {
             self.opened_files.remove(&id);
         }
-        self.vfs_handle.invalidate(path);
+        self.loader.handle.invalidate(path);
 
         Ok(())
     }
@@ -127,7 +127,7 @@ impl LanguageServerState {
     ) -> anyhow::Result<()> {
         for change in params.changes {
             let path = from_lsp::abs_path(&change.uri)?;
-            self.vfs_handle.invalidate(path);
+            self.loader.handle.invalidate(path);
         }
         Ok(())
     }
diff --git a/kclvm/tools/src/LSP/src/state.rs b/kclvm/tools/src/LSP/src/state.rs
index 5dd8e0b62..a3f51afcc 100644
--- a/kclvm/tools/src/LSP/src/state.rs
+++ b/kclvm/tools/src/LSP/src/state.rs
@@ -13,7 +13,6 @@ use lsp_types::{
 };
 use parking_lot::RwLock;
 use ra_ap_vfs::{FileId, Vfs};
-use ra_ap_vfs_notify::NotifyHandle;
 use std::collections::HashMap;
 use std::{sync::Arc, time::Instant};
 
@@ -34,6 +33,11 @@ pub(crate) enum Event {
     Lsp(lsp_server::Message),
 }
 
+pub(crate) struct Handle<H, C> {
+    pub(crate) handle: H,
+    pub(crate) _receiver: C,
+}
+
 /// State for the language server
 pub(crate) struct LanguageServerState {
     /// Channel to send language server messages to the client
@@ -67,7 +71,7 @@ pub(crate) struct LanguageServerState {
     pub opened_files: IndexSet<FileId>,
 
     /// The VFS loader
-    pub vfs_handle: Box<dyn ra_ap_vfs::loader::Handle>,
+    pub loader: Handle<Box<dyn ra_ap_vfs::loader::Handle>, Receiver<ra_ap_vfs::loader::Message>>,
 
     /// The word index map
     pub word_index_map: HashMap<Url, HashMap<String, Vec<Location>>>,
@@ -95,10 +99,13 @@ impl LanguageServerState {
     ) -> Self {
         let (task_sender, task_receiver) = unbounded::<Task>();
 
-        let (vfs_sender, receiver) = unbounded::<ra_ap_vfs::loader::Message>();
-        let handle: NotifyHandle =
-            ra_ap_vfs::loader::Handle::spawn(Box::new(move |msg| vfs_sender.send(msg).unwrap()));
-        let handle = Box::new(handle) as Box<dyn ra_ap_vfs::loader::Handle>;
+        let loader = {
+            let (sender, _receiver) = unbounded::<ra_ap_vfs::loader::Message>();
+            let handle: ra_ap_vfs_notify::NotifyHandle =
+                ra_ap_vfs::loader::Handle::spawn(Box::new(move |msg| sender.send(msg).unwrap()));
+            let handle = Box::new(handle) as Box<dyn ra_ap_vfs::loader::Handle>;
+            Handle { handle, _receiver }
+        };
 
         // build word index for all the workspace folders
         // todo: async
@@ -127,8 +134,8 @@ impl LanguageServerState {
             shutdown_requested: false,
             analysis: Analysis::default(),
             opened_files: IndexSet::new(),
-            vfs_handle: handle,
             word_index_map,
+            loader,
         }
     }
 
diff --git a/kclvm/tools/src/LSP/src/tests.rs b/kclvm/tools/src/LSP/src/tests.rs
index c2979ac11..e4ff5e745 100644
--- a/kclvm/tools/src/LSP/src/tests.rs
+++ b/kclvm/tools/src/LSP/src/tests.rs
@@ -582,6 +582,51 @@ fn notification_test() {
     }
 }
 
+#[test]
+fn close_file_test() {
+    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    let mut path = root.clone();
+
+    path.push("src/test_data/diagnostics.k");
+
+    let path = path.to_str().unwrap();
+    let src = std::fs::read_to_string(path.clone()).unwrap();
+    let server = Project {}.server(InitializeParams::default());
+
+    // Mock open file
+    server.notification::<lsp_types::notification::DidOpenTextDocument>(
+        lsp_types::DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: Url::from_file_path(path).unwrap(),
+                language_id: "KCL".to_string(),
+                version: 0,
+                text: src.clone(),
+            },
+        },
+    );
+
+    // Mock close file
+    server.notification::<lsp_types::notification::DidCloseTextDocument>(
+        lsp_types::DidCloseTextDocumentParams {
+            text_document: TextDocumentIdentifier {
+                uri: Url::from_file_path(path).unwrap(),
+            },
+        },
+    );
+
+    // Mock reopen file
+    server.notification::<lsp_types::notification::DidOpenTextDocument>(
+        lsp_types::DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: Url::from_file_path(path).unwrap(),
+                language_id: "KCL".to_string(),
+                version: 0,
+                text: src,
+            },
+        },
+    );
+}
+
 #[test]
 fn goto_def_test() {
     let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));