fix: fix sema token request failed when open file (#1610)

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

kclvm/tools/src/LSP/src/analysis.rs

@@ -29,6 +29,7 @@ pub enum DBState {
     // The previous version of db
     Compiling(Arc<AnalysisDatabase>),
     Init,
+    Failed(String),
 }
 
 /// AnalysisDatabase holds the result of the compile

kclvm/tools/src/LSP/src/request.rs

@@ -99,6 +99,7 @@ impl LanguageServerSnapshot {
                 Some(db) => match db {
                     DBState::Ready(db) => Ok(Some(db.clone())),
                     DBState::Compiling(_) | DBState::Init => Ok(None),
+                    DBState::Failed(e) => Err(anyhow::anyhow!(e)),
                 },
                 None => Ok(None),
             },
@@ -312,6 +313,7 @@ pub(crate) fn handle_completion(
     if matches!(completion_trigger_character, Some('\n')) {
         match db_state {
             DBState::Compiling(_) | DBState::Init => return Err(anyhow!(LSPError::Retry)),
+            DBState::Failed(_) => return Ok(None),
             _ => {}
         }
     }
@@ -320,6 +322,7 @@ pub(crate) fn handle_completion(
         DBState::Ready(db) => db,
         DBState::Compiling(db) => db,
         DBState::Init => return Err(anyhow!(LSPError::Retry)),
+        DBState::Failed(_) => return Ok(None),
     };
 
     let kcl_pos = kcl_pos(&file, params.text_document_position.position);

kclvm/tools/src/LSP/src/state.rs

@@ -200,6 +200,7 @@ impl LanguageServerState {
         };
 
         // 2. Process changes
+
         self.process_vfs_changes();
         Ok(())
     }
@@ -256,6 +257,7 @@ impl LanguageServerState {
                                         .send(Task::ChangedFile(file.file_id, file.change_kind))
                                         .unwrap();
                                 }
+                                DBState::Failed(_) => continue,
                             }
                         }
 
@@ -264,17 +266,17 @@ impl LanguageServerState {
                             let tool = Arc::clone(&self.tool);
                             let (workspaces, failed) =
                                 lookup_compile_workspaces(&*tool.read(), &filename, true);
-                            for (workspace, opts) in workspaces {
-                                self.async_compile(workspace, opts, Some(file.file_id), true);
-                            }
-                            if self
-                                .temporary_workspace
-                                .read()
-                                .get(&file.file_id)
-                                .unwrap_or(&None)
-                                .is_none()
-                            {
+
+                            if workspaces.is_empty() {
                                 self.temporary_workspace.write().remove(&file.file_id);
+                                self.log_message(format!(
+                                    "Not found any workspace for {:?}",
+                                    filename
+                                ));
+                            } else {
+                                for (workspace, opts) in workspaces {
+                                    self.async_compile(workspace, opts, Some(file.file_id), true);
+                                }
                             }
                         } else {
                             self.temporary_workspace.write().remove(&file.file_id);
@@ -500,7 +502,7 @@ impl LanguageServerState {
                         Some(option_db) => match option_db {
                             DBState::Ready(db) => db.diags.clone(),
                             DBState::Compiling(db) => db.diags.clone(),
-                            DBState::Init => IndexSet::new(),
+                            DBState::Init | DBState::Failed(_) => IndexSet::new(),
                         },
                         None => IndexSet::new(),
                     }
@@ -512,15 +514,17 @@ impl LanguageServerState {
                         Some(state) => match state {
                             DBState::Ready(db) => DBState::Compiling(db.clone()),
                             DBState::Compiling(db) => DBState::Compiling(db.clone()),
-                            DBState::Init => DBState::Init,
+                            DBState::Init | DBState::Failed(_) => DBState::Init,
                         },
                         None => DBState::Init,
                     };
                     workspaces.insert(workspace.clone(), state);
                 }
+                let start = Instant::now();
+
                 let (diags, compile_res) = compile(
                     Params {
-                        file: filename,
+                        file: filename.clone(),
                         module_cache: Some(module_cache),
                         scope_cache: Some(scope_cache),
                         vfs: Some(snapshot.vfs),
@@ -530,6 +534,16 @@ impl LanguageServerState {
                     opts.1,
                 );
 
+                log_message(
+                    format!(
+                        "Compile workspace: {:?}, changed file: {:?}, use {:?} micros",
+                        workspace,
+                        filename,
+                        start.elapsed().as_micros()
+                    ),
+                    &sender,
+                );
+
                 let mut old_diags_maps = HashMap::new();
                 for diag in &old_diags {
                     let lsp_diag = kcl_diag_to_lsp_diags(diag);
@@ -538,7 +552,7 @@ impl LanguageServerState {
                     }
                 }
 
-                // publich diags
+                // publish diags
                 let mut new_diags_maps = HashMap::new();
 
                 for diag in &diags {
@@ -595,7 +609,7 @@ impl LanguageServerState {
                     }
                     Err(e) => {
                         let mut workspaces = snapshot.workspaces.write();
-                        workspaces.remove(&workspace);
+                        workspaces.insert(workspace, DBState::Failed(e.to_string()));
                         if temp && changed_file_id.is_some() {
                             let mut temporary_workspace = snapshot.temporary_workspace.write();
                             temporary_workspace.remove(&changed_file_id.unwrap());

kclvm/tools/src/LSP/src/test_data/workspace/init/folder/a.k

@@ -0,0 +1 @@
+a = 1

kclvm/tools/src/LSP/src/test_data/workspace/init/folder/sub/c.k

@@ -0,0 +1 @@
+a = 1

kclvm/tools/src/LSP/src/tests.rs

@@ -28,15 +28,18 @@ use lsp_types::HoverContents;
 use lsp_types::HoverParams;
 use lsp_types::InitializeParams;
 use lsp_types::MarkedString;
+use lsp_types::PartialResultParams;
 use lsp_types::PublishDiagnosticsParams;
 use lsp_types::ReferenceContext;
 use lsp_types::ReferenceParams;
 use lsp_types::RenameParams;
+use lsp_types::SemanticTokensParams;
 use lsp_types::TextDocumentIdentifier;
 use lsp_types::TextDocumentItem;
 use lsp_types::TextDocumentPositionParams;
 use lsp_types::TextEdit;
 use lsp_types::Url;
+use lsp_types::WorkDoneProgressParams;
 use lsp_types::WorkspaceEdit;
 use lsp_types::WorkspaceFolder;
 
@@ -2248,3 +2251,100 @@ fn kcl_workspace_init_folder_test() {
 
     assert!(failed.is_none());
 }
+
+#[test]
+fn init_workspace_sema_token_test() {
+    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("src")
+        .join("test_data")
+        .join("workspace")
+        .join("init")
+        .join("folder");
+
+    let mut a_path = root.clone();
+    a_path.push("a.k");
+
+    let mut c_path = root.clone();
+    c_path.push("sub");
+    c_path.push("c.k");
+
+    let a_path = a_path.to_str().unwrap();
+    let c_path = c_path.to_str().unwrap();
+    let a_src = std::fs::read_to_string(a_path).unwrap();
+    let c_src = std::fs::read_to_string(c_path).unwrap();
+    let initialize_params = InitializeParams {
+        workspace_folders: Some(vec![WorkspaceFolder {
+            uri: Url::from_file_path(root.clone()).unwrap(),
+            name: "test".to_string(),
+        }]),
+        ..Default::default()
+    };
+    let server = Project {}.server(initialize_params);
+
+    let a_url = Url::from_file_path(a_path).unwrap();
+    let c_url = Url::from_file_path(c_path).unwrap();
+
+    // Mock open file in init workspace
+    server.notification::<lsp_types::notification::DidOpenTextDocument>(
+        lsp_types::DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: a_url.clone(),
+                language_id: "KCL".to_string(),
+                version: 0,
+                text: a_src,
+            },
+        },
+    );
+
+    let id = server.next_request_id.get();
+    server.next_request_id.set(id.wrapping_add(1));
+
+    let r: Request = Request::new(
+        id.into(),
+        "textDocument/semanticTokens/full".to_string(),
+        SemanticTokensParams {
+            text_document: TextDocumentIdentifier { uri: a_url },
+            work_done_progress_params: WorkDoneProgressParams {
+                work_done_token: None,
+            },
+            partial_result_params: PartialResultParams {
+                partial_result_token: None,
+            },
+        },
+    );
+
+    let res = server.send_and_receive(r);
+    assert!(res.result.is_some());
+
+    // Mock open file not in init workspace
+    server.notification::<lsp_types::notification::DidOpenTextDocument>(
+        lsp_types::DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: c_url.clone(),
+                language_id: "KCL".to_string(),
+                version: 0,
+                text: c_src,
+            },
+        },
+    );
+
+    let id = server.next_request_id.get();
+    server.next_request_id.set(id.wrapping_add(1));
+
+    let r: Request = Request::new(
+        id.into(),
+        "textDocument/semanticTokens/full".to_string(),
+        SemanticTokensParams {
+            text_document: TextDocumentIdentifier { uri: c_url },
+            work_done_progress_params: WorkDoneProgressParams {
+                work_done_token: None,
+            },
+            partial_result_params: PartialResultParams {
+                partial_result_token: None,
+            },
+        },
+    );
+
+    let res = server.send_and_receive(r);
+    assert!(res.result.is_some());
+}