Tree-sitter Performance

lsp/Cargo.toml

@@ -20,3 +20,4 @@ tree-sitter.workspace = true
 tree-sitter-html.workspace = true
 maplit = "1.0.2"
 phf = { version = "0.11.2", features = ["macros"] }
+lsp-textdocument = "0.3.2"

lsp/src/handle.rs

@@ -1,42 +1,15 @@
 use crate::{
     htmx::{hx_completion, hx_hover, HxCompletion},
-    text_store::TEXT_STORE,
+    text_store::{DocInfo, DOCUMENT_STORE},
+    tree_sitter::text_doc_change_to_ts_edit,
 };
 use log::{debug, error, warn};
 use lsp_server::{Message, Notification, Request, RequestId};
-use lsp_types::{CompletionContext, CompletionParams, CompletionTriggerKind, HoverParams};
-
-#[derive(serde::Deserialize, Debug)]
-struct Text {
-    text: String,
-}
-
-#[derive(serde::Deserialize, Debug)]
-struct TextDocumentLocation {
-    uri: String,
-}
-
-#[derive(serde::Deserialize, Debug)]
-struct TextDocumentChanges {
-    #[serde(rename = "textDocument")]
-    text_document: TextDocumentLocation,
-
-    #[serde(rename = "contentChanges")]
-    content_changes: Vec<Text>,
-}
-
-#[derive(serde::Deserialize, Debug)]
-struct TextDocumentOpened {
-    uri: String,
-
-    text: String,
-}
-
-#[derive(serde::Deserialize, Debug)]
-struct TextDocumentOpen {
-    #[serde(rename = "textDocument")]
-    text_document: TextDocumentOpened,
-}
+use lsp_textdocument::FullTextDocument;
+use lsp_types::{
+    notification::{DidChangeTextDocument, DidOpenTextDocument},
+    CompletionContext, CompletionParams, CompletionTriggerKind, HoverParams,
+};
 
 #[derive(Debug)]
 pub struct HtmxAttributeCompletion {
@@ -61,41 +34,84 @@ pub enum HtmxResult {
 // ignore snakeCase
 #[allow(non_snake_case)]
 fn handle_didChange(noti: Notification) -> Option<HtmxResult> {
-    let text_document_changes: TextDocumentChanges = serde_json::from_value(noti.params).ok()?;
-    let uri = text_document_changes.text_document.uri;
-    let text = text_document_changes.content_changes[0].text.to_string();
-
-    if text_document_changes.content_changes.len() > 1 {
-        error!("more than one content change, please be wary");
+    match cast_notif::<DidChangeTextDocument>(noti) {
+        Ok(params) => {
+            match DOCUMENT_STORE
+                .get()
+                .expect("text store not initialized")
+                .lock()
+                .expect("text store mutex poisoned")
+                .get_mut(params.text_document.uri.as_str())
+            {
+                Some(entry) => {
+                    entry
+                        .doc
+                        .update(&params.content_changes, params.text_document.version);
+
+                    if let Some(ref mut curr_tree) = entry.tree {
+                        for edit in params.content_changes.iter() {
+                            match text_doc_change_to_ts_edit(edit, &entry.doc) {
+                                Ok(edit) => {
+                                    curr_tree.edit(&edit);
+                                }
+                                Err(e) => {
+                                    error!("handle_didChange Bad edit info, failed to edit tree -- Error: {e}");
+                                }
+                            }
+                        }
+                    } else {
+                        error!(
+                            "handle_didChange tree for {} is None",
+                            params.text_document.uri.as_str()
+                        );
+                    }
+                }
+                None => {
+                    error!(
+                        "handle_didChange No corresponding doc for supplied edits -- {}",
+                        params.text_document.uri.as_str()
+                    );
+                }
+            }
+        }
+        Err(e) => {
+            error!("Failed the deserialize DidChangeTextDocument params -- Error {e}");
+        }
     }
 
-    TEXT_STORE
-        .get()
-        .expect("text store not initialized")
-        .lock()
-        .expect("text store mutex poisoned")
-        .insert(uri, text);
-
     None
 }
 
 #[allow(non_snake_case)]
 fn handle_didOpen(noti: Notification) -> Option<HtmxResult> {
     debug!("handle_didOpen params {:?}", noti.params);
-    let text_document_changes = match serde_json::from_value::<TextDocumentOpen>(noti.params) {
-        Ok(p) => p.text_document,
+    let text_doc_open = match cast_notif::<DidOpenTextDocument>(noti) {
+        Ok(params) => params,
         Err(err) => {
             error!("handle_didOpen parsing params error : {:?}", err);
             return None;
         }
     };
 
-    TEXT_STORE
+    let doc = FullTextDocument::new(
+        text_doc_open.text_document.language_id,
+        text_doc_open.text_document.version,
+        text_doc_open.text_document.text,
+    );
+    let mut parser = ::tree_sitter::Parser::new();
+    parser
+        .set_language(tree_sitter_html::language())
+        .expect("Failed to load HTML grammar");
+    let tree = parser.parse(doc.get_content(None), None);
+
+    let doc = DocInfo { doc, parser, tree };
+
+    DOCUMENT_STORE
         .get()
         .expect("text store not initialized")
         .lock()
         .expect("text store mutex poisoned")
-        .insert(text_document_changes.uri, text_document_changes.text);
+        .insert(text_doc_open.text_document.uri.to_string(), doc);
 
     None
 }
@@ -186,10 +202,23 @@ pub fn handle_other(msg: Message) -> Option<HtmxResult> {
     None
 }
 
+fn cast_notif<R>(notif: Notification) -> anyhow::Result<R::Params>
+where
+    R: lsp_types::notification::Notification,
+    R::Params: serde::de::DeserializeOwned,
+{
+    match notif.extract(R::METHOD) {
+        Ok(value) => Ok(value),
+        Err(e) => Err(anyhow::anyhow!(
+            "cast_notif Failed to extract params -- Error: {e}"
+        )),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{handle_request, HtmxResult, Request};
-    use crate::text_store::{init_text_store, TEXT_STORE};
+    use crate::text_store::{init_text_store, DocInfo, DOCUMENT_STORE};
     use std::sync::Once;
 
     static SETUP: Once = Once::new();
@@ -198,12 +227,21 @@ mod tests {
             init_text_store();
         });
 
-        TEXT_STORE
+        let doc =
+            lsp_textdocument::FullTextDocument::new("html".to_string(), 0, content.to_string());
+        let mut parser = ::tree_sitter::Parser::new();
+        parser
+            .set_language(tree_sitter_html::language())
+            .expect("Failed to load HTML grammar");
+        let tree = parser.parse(doc.get_content(None), None);
+        let doc_info = DocInfo { doc, parser, tree };
+
+        DOCUMENT_STORE
             .get()
             .expect("text store not initialized")
             .lock()
             .expect("text store mutex poisoned")
-            .insert(file.to_string(), content.to_string());
+            .insert(file.to_string(), doc_info);
     }
 
     #[test]

lsp/src/lib.rs

@@ -9,8 +9,8 @@ use htmx::HxCompletion;
 use log::{debug, error, info, warn};
 use lsp_types::{
     CompletionItem, CompletionItemKind, CompletionList, HoverContents, InitializeParams,
-    MarkupContent, ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
-    WorkDoneProgressOptions,
+    MarkupContent, PositionEncodingKind, ServerCapabilities, TextDocumentSyncCapability,
+    TextDocumentSyncKind, WorkDoneProgressOptions,
 };
 
 use lsp_server::{Connection, Message, Response};
@@ -123,7 +123,10 @@ pub fn start_lsp() -> Result<()> {
 
     // Run the server and wait for the two threads to end (typically by trigger LSP Exit event).
     let server_capabilities = serde_json::to_value(ServerCapabilities {
-        text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::FULL)),
+        position_encoding: Some(PositionEncodingKind::UTF16), // compatability with lsp_textdocument crate
+        text_document_sync: Some(TextDocumentSyncCapability::Kind(
+            TextDocumentSyncKind::INCREMENTAL,
+        )),
         completion_provider: Some(lsp_types::CompletionOptions {
             resolve_provider: Some(false),
             trigger_characters: Some(vec!["-".to_string(), "\"".to_string(), " ".to_string()]),

lsp/src/text_store.rs

@@ -4,39 +4,48 @@ use std::{
     sync::{Arc, Mutex, OnceLock},
 };
 
-use lsp_types::{TextDocumentPositionParams, Url};
+use lsp_textdocument::FullTextDocument;
+use lsp_types::{Position, Range, TextDocumentPositionParams, Url};
+use tree_sitter::{Parser, Tree};
 
-type TxtStore = HashMap<String, String>;
+pub struct DocInfo {
+    pub doc: FullTextDocument,
+    pub parser: Parser,
+    pub tree: Option<Tree>,
+}
+
+type DocStore = HashMap<String, DocInfo>;
 
-pub struct TextStore(TxtStore);
+#[derive(Default)]
+pub struct DocumentStore(DocStore);
 
-impl Deref for TextStore {
-    type Target = TxtStore;
+impl Deref for DocumentStore {
+    type Target = DocStore;
     fn deref(&self) -> &Self::Target {
         &self.0
     }
 }
 
-impl DerefMut for TextStore {
+impl DerefMut for DocumentStore {
     fn deref_mut(&mut self) -> &mut Self::Target {
         &mut self.0
     }
 }
 
-pub static TEXT_STORE: OnceLock<Arc<Mutex<TextStore>>> = OnceLock::new();
+pub static DOCUMENT_STORE: OnceLock<Arc<Mutex<DocumentStore>>> = OnceLock::new();
 
 pub fn init_text_store() {
-    _ = TEXT_STORE.set(Arc::new(Mutex::new(TextStore(HashMap::new()))));
+    _ = DOCUMENT_STORE.set(Arc::new(Mutex::new(DocumentStore::default())));
 }
 
-pub fn get_text_document(uri: &Url) -> Option<String> {
-    return TEXT_STORE
+pub fn get_text_document(uri: &Url, range: Option<Range>) -> Option<String> {
+    return DOCUMENT_STORE
         .get()
         .expect("text store not initialized")
         .lock()
         .expect("text store mutex poisoned")
         .get(&uri.to_string())
-        .cloned();
+        .map(|doc| doc.doc.get_content(range).to_string());
 }
 
 /// Find the start and end indices of a word inside the given line
@@ -67,20 +76,20 @@ fn find_word_at_pos(line: &str, col: usize) -> (usize, usize) {
 
 pub fn get_word_from_pos_params(pos_params: &TextDocumentPositionParams) -> anyhow::Result<String> {
     let uri = &pos_params.text_document.uri;
-    let line = pos_params.position.line as usize;
+    let line = pos_params.position.line;
     let col = pos_params.position.character as usize;
 
-    match get_text_document(uri) {
-        Some(text) => {
-            let line_conts = match text.lines().nth(line) {
-                Some(conts) => conts,
-                None => {
-                    return Err(anyhow::anyhow!(
-                        "get_word_from_pos_params Failed to get word under cursor"
-                    ));
-                }
-            };
-            let (start, end) = find_word_at_pos(line_conts, col);
+    let range = Range {
+        start: Position { line, character: 0 },
+        end: Position {
+            line,
+            character: u32::MAX,
+        },
+    };
+
+    match get_text_document(uri, Some(range)) {
+        Some(line_conts) => {
+            let (start, end) = find_word_at_pos(&line_conts, col);
             Ok(String::from(&line_conts[start..end]))
         }
         None => Err(anyhow::anyhow!(