improve: parse workspace in worker thread with progress (#24)

Co-authored-by: mohammadkhan <mohammadkhan@digitalocean.com>

Author: coder3101

src/lsp.rs

@@ -1,5 +1,7 @@
 use std::fs::read_to_string;
 use std::ops::ControlFlow;
+use std::sync::mpsc;
+use std::thread;
 use tracing::{error, info};
 
 use async_lsp::lsp_types::{
@@ -9,9 +11,10 @@ use async_lsp::lsp_types::{
     DocumentSymbolResponse, FileOperationFilter, FileOperationPattern, FileOperationPatternKind,
     FileOperationRegistrationOptions, GotoDefinitionParams, GotoDefinitionResponse, Hover,
     HoverContents, HoverParams, HoverProviderCapability, InitializeParams, InitializeResult, OneOf,
-    PrepareRenameResponse, RenameFilesParams, RenameOptions, RenameParams, ServerCapabilities,
-    ServerInfo, TextDocumentPositionParams, TextDocumentSyncCapability, TextDocumentSyncKind, Url,
-    WorkspaceEdit, WorkspaceFileOperationsServerCapabilities, WorkspaceFoldersServerCapabilities,
+    PrepareRenameResponse, ProgressParams, RenameFilesParams, RenameOptions,
+    RenameParams, ServerCapabilities, ServerInfo, TextDocumentPositionParams,
+    TextDocumentSyncCapability, TextDocumentSyncKind, Url, WorkspaceEdit,
+    WorkspaceFileOperationsServerCapabilities, WorkspaceFoldersServerCapabilities,
     WorkspaceServerCapabilities,
 };
 use async_lsp::{LanguageClient, LanguageServer, ResponseError};
@@ -46,6 +49,25 @@ impl LanguageServer for ProtoLanguageServer {
             },
         }];
 
+        let worktoken = params.work_done_progress_params.work_done_token;
+        let (tx, rx) = mpsc::channel();
+        let mut socket = self.client.clone();
+
+        thread::spawn(move || {
+            let Some(token) = worktoken else {
+                return;
+            };
+
+            while let Ok(value) = rx.recv() {
+                if let Err(e) = socket.progress(ProgressParams {
+                    token: token.clone(),
+                    value,
+                }) {
+                    error!(error=%e, "failed to report parse progress");
+                }
+            }
+        });
+
         let file_registration_option = FileOperationRegistrationOptions {
             filters: file_operation_filers.clone(),
         };
@@ -54,7 +76,7 @@ impl LanguageServer for ProtoLanguageServer {
         if let Some(folders) = params.workspace_folders {
             for workspace in folders {
                 info!("Workspace folder: {workspace:?}");
-                self.state.add_workspace_folder(workspace)
+                self.state.add_workspace_folder_async(workspace, tx.clone())
             }
             workspace_capabilities = Some(WorkspaceServerCapabilities {
                 workspace_folders: Some(WorkspaceFoldersServerCapabilities {

src/parser/definition.rs

@@ -22,7 +22,7 @@ impl ParsedTree {
             return;
         }
 
-        if !identifier.contains(".") {
+        if !identifier.contains('.') {
             let locations: Vec<Location> = self
                 .filter_nodes_from(n, NodeKind::is_userdefined)
                 .into_iter()
@@ -41,7 +41,7 @@ impl ParsedTree {
         }
 
         // Safety: identifier contains a .
-        let (parent_identifier, remaining) = identifier.split_once(".").unwrap();
+        let (parent_identifier, remaining) = identifier.split_once('.').unwrap();
         let child_node = self
             .filter_nodes_from(n, NodeKind::is_userdefined)
             .into_iter()

src/parser/hover.rs

@@ -64,7 +64,7 @@ impl ParsedTree {
             return;
         }
 
-        if !identifier.contains(".") {
+        if !identifier.contains('.') {
             let comments: Vec<MarkedString> = self
                 .filter_nodes_from(n, NodeKind::is_userdefined)
                 .into_iter()
@@ -78,7 +78,7 @@ impl ParsedTree {
         }
 
         // Safety: identifier contains a .
-        let (parent_identifier, remaining) = identifier.split_once(".").unwrap();
+        let (parent_identifier, remaining) = identifier.split_once('.').unwrap();
         let child_node = self
             .filter_nodes_from(n, NodeKind::is_userdefined)
             .into_iter()

src/parser/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use async_lsp::lsp_types::Url;
 use tree_sitter::Tree;
 
@@ -12,9 +14,10 @@ pub struct ProtoParser {
     parser: tree_sitter::Parser,
 }
 
+#[derive(Clone)]
 pub struct ParsedTree {
     pub uri: Url,
-    tree: Tree,
+    tree: Arc<Tree>,
 }
 
 impl ProtoParser {
@@ -27,8 +30,9 @@ impl ProtoParser {
     }
 
     pub fn parse(&mut self, uri: Url, contents: impl AsRef<[u8]>) -> Option<ParsedTree> {
-        self.parser
-            .parse(contents, None)
-            .map(|t| ParsedTree { tree: t, uri })
+        self.parser.parse(contents, None).map(|t| ParsedTree {
+            tree: Arc::new(t),
+            uri,
+        })
     }
 }

src/state.rs

@@ -1,8 +1,15 @@
-use std::{collections::HashMap, fs::read_to_string};
+use std::{
+    collections::HashMap,
+    fs::read_to_string,
+    sync::{mpsc::Sender, Arc, Mutex, MutexGuard, RwLock, RwLockWriteGuard},
+    thread,
+};
 use tracing::{error, info};
 
 use async_lsp::lsp_types::{
-    CompletionItem, CompletionItemKind, PublishDiagnosticsParams, Url, WorkspaceFolder,
+    CompletionItem, CompletionItemKind, ProgressParamsValue, PublishDiagnosticsParams, Url,
+    WorkDoneProgress, WorkDoneProgressBegin, WorkDoneProgressEnd, WorkDoneProgressReport,
+    WorkspaceFolder,
 };
 use tree_sitter::Node;
 use walkdir::WalkDir;
@@ -13,79 +20,146 @@ use crate::{
 };
 
 pub struct ProtoLanguageState {
-    documents: HashMap<Url, String>,
-    pub trees: HashMap<Url, ParsedTree>,
-    parser: ProtoParser,
+    documents: Arc<RwLock<HashMap<Url, String>>>,
+    trees: Arc<RwLock<HashMap<Url, ParsedTree>>>,
+    parser: Arc<Mutex<ProtoParser>>,
 }
 
 impl ProtoLanguageState {
     pub fn new() -> Self {
         ProtoLanguageState {
             documents: Default::default(),
             trees: Default::default(),
-            parser: ProtoParser::new(),
+            parser: Arc::new(Mutex::new(ProtoParser::new())),
         }
     }
 
-    pub fn get_content(&self, uri: &Url) -> &str {
+    pub fn get_content(&self, uri: &Url) -> String {
         self.documents
+            .read()
+            .expect("poison")
             .get(uri)
-            .map(|s| s.as_str())
+            .map(|s| s.to_string())
             .unwrap_or_default()
     }
 
-    pub fn get_tree(&self, uri: &Url) -> Option<&ParsedTree> {
-        self.trees.get(uri)
+    pub fn get_tree(&self, uri: &Url) -> Option<ParsedTree> {
+        self.trees.read().expect("poison").get(uri).cloned()
     }
 
-    pub fn get_trees_for_package(&self, package: &str) -> Vec<&ParsedTree> {
+    pub fn get_trees_for_package(&self, package: &str) -> Vec<ParsedTree> {
         self.trees
+            .read()
+            .expect("poison")
             .values()
             .filter(|tree| {
                 let content = self.get_content(&tree.uri);
                 tree.get_package_name(content.as_bytes())
                     .unwrap_or_default()
                     == package
             })
+            .map(ToOwned::to_owned)
             .collect()
     }
 
-    pub fn upsert_content(&mut self, uri: &Url, content: String) -> bool {
-        if let Some(parsed) = self.parser.parse(uri.clone(), content.as_bytes()) {
-            self.trees.insert(uri.clone(), parsed);
-            self.documents.insert(uri.clone(), content);
+    fn upsert_content_impl(
+        mut parser: MutexGuard<ProtoParser>,
+        uri: &Url,
+        content: String,
+        mut docs: RwLockWriteGuard<HashMap<Url, String>>,
+        mut trees: RwLockWriteGuard<HashMap<Url, ParsedTree>>,
+    ) -> bool {
+        if let Some(parsed) = parser.parse(uri.clone(), content.as_bytes()) {
+            trees.insert(uri.clone(), parsed);
+            docs.insert(uri.clone(), content);
             true
         } else {
-            error!(uri=%uri, "failed to parse content");
             false
         }
     }
 
-    pub fn add_workspace_folder(&mut self, workspace: WorkspaceFolder) {
-        for entry in WalkDir::new(workspace.uri.path())
-            .into_iter()
-            .filter_map(|e| e.ok())
-        {
-            let path = entry.path();
-            if path.is_absolute() && path.is_file() {
-                let Some(ext) = path.extension() else {
-                    continue;
-                };
-
-                let Ok(content) = read_to_string(path) else {
-                    continue;
-                };
-
-                let Ok(uri) = Url::from_file_path(path) else {
-                    continue;
-                };
-
-                if ext == "proto" {
-                    let r = self.upsert_content(&uri, content);
-                    info!("workspace parse file: {}, result: {}", path.display(), r);
+    pub fn upsert_content(&mut self, uri: &Url, content: String) -> bool {
+        let parser = self.parser.lock().expect("poison");
+        let tree = self.trees.write().expect("poison");
+        let docs = self.documents.write().expect("poison");
+        Self::upsert_content_impl(parser, uri, content, docs, tree)
+    }
+
+    pub fn add_workspace_folder_async(
+        &mut self,
+        workspace: WorkspaceFolder,
+        tx: Sender<ProgressParamsValue>,
+    ) {
+        let parser = self.parser.clone();
+        let tree = self.trees.clone();
+        let docs = self.documents.clone();
+
+        let begin = ProgressParamsValue::WorkDone(WorkDoneProgress::Begin(WorkDoneProgressBegin {
+            title: String::from("indexing"),
+            cancellable: Some(false),
+            percentage: Some(0),
+            ..Default::default()
+        }));
+
+        if let Err(e) = tx.send(begin) {
+            error!(error=%e, "failed to send work begin progress");
+        }
+
+        thread::spawn(move || {
+            let files: Vec<_> = WalkDir::new(workspace.uri.path())
+                .into_iter()
+                .filter_map(|e| e.ok())
+                .filter(|e| e.path().extension().is_some())
+                .filter(|e| e.path().extension().unwrap() == "proto")
+                .collect();
+
+            let total_files = files.len();
+            let mut current = 0;
+
+            for file in files.into_iter() {
+                let path = file.path();
+                if path.is_absolute() && path.is_file() {
+                    let Ok(content) = read_to_string(path) else {
+                        continue;
+                    };
+
+                    let Ok(uri) = Url::from_file_path(path) else {
+                        continue;
+                    };
+
+                    Self::upsert_content_impl(
+                        parser.lock().expect("poison"),
+                        &uri,
+                        content,
+                        docs.write().expect("poison"),
+                        tree.write().expect("poison"),
+                    );
+
+                    current += 1;
+
+                    let report = ProgressParamsValue::WorkDone(WorkDoneProgress::Report(
+                        WorkDoneProgressReport {
+                            cancellable: Some(false),
+                            message: Some(path.display().to_string()),
+                            percentage: Some((current * 100 / total_files) as u32),
+                        },
+                    ));
+
+                    if let Err(e) = tx.send(report) {
+                        error!(error=%e, "failed to send work report progress");
+                    }
                 }
             }
-        }
+            let report =
+                ProgressParamsValue::WorkDone(WorkDoneProgress::End(WorkDoneProgressEnd {
+                    message: Some(String::from("completed")),
+                }));
+
+            info!(len = total_files, "workspace file parsing completed");
+            if let Err(e) = tx.send(report) {
+                error!(error=%e, "failed to send work completed result");
+            }
+        });
     }
 
     pub fn upsert_file(&mut self, uri: &Url, content: String) -> Option<PublishDiagnosticsParams> {
@@ -96,20 +170,26 @@ impl ProtoLanguageState {
 
     pub fn delete_file(&mut self, uri: &Url) {
         info!(uri=%uri, "deleting file");
-        self.documents.remove(uri);
-        self.trees.remove(uri);
+        self.documents.write().expect("poison").remove(uri);
+        self.trees.write().expect("poison").remove(uri);
     }
 
     pub fn rename_file(&mut self, new_uri: &Url, old_uri: &Url) {
         info!(new_uri=%new_uri, old_uri=%new_uri, "renaming file");
 
-        if let Some(v) = self.documents.remove(old_uri) {
-            self.documents.insert(new_uri.clone(), v);
+        if let Some(v) = self.documents.write().expect("poison").remove(old_uri) {
+            self.documents
+                .write()
+                .expect("poison")
+                .insert(new_uri.clone(), v);
         }
 
-        if let Some(mut v) = self.trees.remove(old_uri) {
+        if let Some(mut v) = self.trees.write().expect("poison").remove(old_uri) {
             v.uri = new_uri.clone();
-            self.trees.insert(new_uri.clone(), v);
+            self.trees
+                .write()
+                .expect("poison")
+                .insert(new_uri.clone(), v);
         }
     }
 

src/utils.rs

@@ -30,19 +30,19 @@ fn is_first_lower_case(s: &&str) -> bool {
 }
 
 pub fn is_inner_identifier(s: &str) -> bool {
-    if !s.contains(".") {
+    if !s.contains('.') {
         return false;
     }
-    s.split(".").all(is_title_case)
+    s.split('.').all(is_title_case)
 }
 
 pub fn split_identifier_package(s: &str) -> (&str, &str) {
-    if is_inner_identifier(s) || !s.contains(".") {
+    if is_inner_identifier(s) || !s.contains('.') {
         return ("", s);
     }
 
     let i = s
-        .split(".")
+        .split('.')
         .take_while(is_first_lower_case)
         .fold(0, |mut c, s| {
             if c != 0 {