refactor: modularize server handlers

README.md

@@ -85,16 +85,15 @@ vim.api.nvim_create_autocmd('FileType', {
     should be exposed to gather pattern information more efficiently
 - [x] Recognize parsers built for `WASM`
 - [x] Document formatting compatible with the `nvim-treesitter` formatter
-- [ ] Code cleanup
+- [x] Code cleanup
 - [ ] Add tests for all functionality
 
 ### Packaging
 
 - [ ] [`homebrew`](https://github.com/Homebrew/homebrew-core)
       ([in progress](https://github.com/Homebrew/homebrew-core/pull/197587),
       requires repo to reach 75 GitHub stars)
-- [ ] [`nixpkgs`](https://github.com/NixOS/nixpkgs)
-      ([in progress](https://github.com/NixOS/nixpkgs/pull/350834))
+- [x] [`nixpkgs`](https://github.com/NixOS/nixpkgs)
 - [ ] [`mason.nvim`](https://github.com/mason-org/mason-registry)
       ([in progress](https://github.com/mason-org/mason-registry/pull/7849))
 

src/handlers/completion.rs

@@ -0,0 +1,119 @@
+use std::collections::HashSet;
+
+use log::warn;
+use streaming_iterator::StreamingIterator;
+use tower_lsp::jsonrpc::Result;
+use tower_lsp::lsp_types::{
+    CompletionItem, CompletionItemKind, CompletionParams, CompletionResponse,
+};
+use tree_sitter::{Query, QueryCursor};
+
+use crate::util::{
+    get_node_text, lsp_position_to_byte_offset, lsp_position_to_ts_point, node_is_or_has_ancestor,
+    TextProviderRope,
+};
+use crate::{Backend, QUERY_LANGUAGE};
+
+pub async fn completion(
+    backend: &Backend,
+    params: CompletionParams,
+) -> Result<Option<CompletionResponse>> {
+    let uri = &params.text_document_position.text_document.uri;
+
+    let Some(tree) = backend.cst_map.get(uri) else {
+        warn!("No CST built for URI: {uri:?}");
+        return Ok(None);
+    };
+    let Some(rope) = backend.document_map.get(uri) else {
+        warn!("No document built for URI: {uri:?}");
+        return Ok(None);
+    };
+
+    let mut position = params.text_document_position.position;
+    if position.character > 0 {
+        position.character -= 1;
+    }
+    let point = lsp_position_to_ts_point(position, &rope);
+    let query = Query::new(&QUERY_LANGUAGE, "(capture) @cap").unwrap();
+    let mut cursor = QueryCursor::new();
+    let current_node = tree
+        .root_node()
+        .named_descendant_for_point_range(point, point)
+        .unwrap();
+
+    // Don't offer completions when in a comment
+    if current_node.kind() == "comment" {
+        return Ok(None);
+    }
+
+    let mut completion_items = vec![];
+
+    // Node and field name completions
+    let cursor_after_at_sign = lsp_position_to_byte_offset(position, &rope)
+        .and_then(|b| rope.try_byte_to_char(b))
+        .is_ok_and(|c| rope.char(c) == '@');
+    let in_capture =
+        cursor_after_at_sign || node_is_or_has_ancestor(tree.root_node(), current_node, "capture");
+    if !in_capture && !node_is_or_has_ancestor(tree.root_node(), current_node, "predicate") {
+        let in_anon = node_is_or_has_ancestor(tree.root_node(), current_node, "anonymous_node");
+        if let Some(symbols) = backend.symbols_vec_map.get(uri) {
+            for symbol in symbols.iter() {
+                if (in_anon && !symbol.named) || (!in_anon && symbol.named) {
+                    completion_items.push(CompletionItem {
+                        label: symbol.label.clone(),
+                        kind: if symbol.named {
+                            Some(CompletionItemKind::CLASS)
+                        } else {
+                            Some(CompletionItemKind::CONSTANT)
+                        },
+                        ..Default::default()
+                    });
+                }
+            }
+        }
+        if !in_anon {
+            if let Some(fields) = backend.fields_vec_map.get(uri) {
+                for field in fields.iter() {
+                    completion_items.push(CompletionItem {
+                        label: format!("{field}: "),
+                        kind: Some(CompletionItemKind::FIELD),
+                        ..Default::default()
+                    });
+                }
+            }
+        }
+    }
+
+    // Capture completions
+    if !node_is_or_has_ancestor(tree.root_node(), current_node, "predicate")
+        || node_is_or_has_ancestor(tree.root_node(), current_node, "string")
+    {
+        return Ok(Some(CompletionResponse::Array(completion_items)));
+    }
+    let node = match tree.root_node().child_with_descendant(current_node) {
+        None => return Ok(Some(CompletionResponse::Array(completion_items))),
+        Some(value) => value,
+    };
+    let provider = TextProviderRope(&rope);
+    let mut iter = cursor.matches(&query, node, &provider);
+    let mut seen = HashSet::new();
+    while let Some(match_) = iter.next() {
+        for capture in match_.captures {
+            let node_text = get_node_text(&capture.node, &rope);
+            let parent_params = capture
+                .node
+                .parent()
+                .is_none_or(|p| p.kind() != "parameters");
+            if parent_params && !seen.contains(&node_text) {
+                seen.insert(node_text.clone());
+                completion_items.push(CompletionItem {
+                    label: node_text.clone(),
+                    kind: Some(CompletionItemKind::VARIABLE),
+                    ..Default::default()
+                });
+            }
+        }
+    }
+
+    Ok(Some(CompletionResponse::Array(completion_items)))
+}

src/handlers/did_change.rs

@@ -0,0 +1,83 @@
+use tower_lsp::lsp_types::{DidChangeTextDocumentParams, Position, Range};
+use tree_sitter::Parser;
+
+use crate::{
+    util::{get_diagnostics, lsp_textdocchange_to_ts_inputedit, TextProviderRope},
+    Backend, QUERY_LANGUAGE,
+};
+
+pub async fn did_change(backend: &Backend, params: DidChangeTextDocumentParams) {
+    let uri = &params.text_document.uri;
+    let mut rope = backend.document_map.get_mut(uri).unwrap();
+    let mut parser = Parser::new();
+    parser
+        .set_language(&QUERY_LANGUAGE)
+        .expect("Error loading Query grammar");
+
+    let mut edits = vec![];
+    for change in &params.content_changes {
+        let text = change.text.as_str();
+        let text_bytes = text.as_bytes();
+        let text_end_byte_idx = text_bytes.len();
+
+        let range = if let Some(range) = change.range {
+            range
+        } else {
+            let start_line_idx = rope.byte_to_line(0);
+            let end_line_idx = rope.byte_to_line(text_end_byte_idx);
+
+            let start = Position::new(start_line_idx as u32, 0);
+            let end = Position::new(end_line_idx as u32, 0);
+            Range { start, end }
+        };
+
+        edits.push(lsp_textdocchange_to_ts_inputedit(&rope, change).unwrap());
+
+        let start_row_char_idx = rope.line_to_char(range.start.line as usize);
+        let start_row_cu = rope.char_to_utf16_cu(start_row_char_idx);
+        let start_col_char_idx = rope
+            .utf16_cu_to_char(start_row_cu + range.start.character as usize)
+            - start_row_char_idx;
+        let end_row_char_idx = rope.line_to_char(range.end.line as usize);
+        let end_row_cu = rope.char_to_utf16_cu(end_row_char_idx);
+        let end_col_char_idx =
+            rope.utf16_cu_to_char(end_row_cu + range.end.character as usize) - end_row_char_idx;
+
+        let start_char_idx = start_row_char_idx + start_col_char_idx;
+        let end_char_idx = end_row_char_idx + end_col_char_idx;
+        rope.remove(start_char_idx..end_char_idx);
+
+        if !change.text.is_empty() {
+            rope.insert(start_char_idx, text);
+        }
+    }
+    let contents = rope.to_string();
+    let result = {
+        let mut old_tree = backend.cst_map.get_mut(uri).unwrap();
+
+        for edit in edits {
+            old_tree.edit(&edit);
+        }
+
+        parser.parse(&contents, Some(&old_tree))
+    };
+
+    if let Some(tree) = result {
+        *backend.cst_map.get_mut(uri).unwrap() = tree.clone();
+        // Update diagnostics
+        if let (Some(symbols), Some(fields)) = (
+            backend.symbols_set_map.get(uri),
+            backend.fields_set_map.get(uri),
+        ) {
+            let provider = TextProviderRope(&rope);
+            backend
+                .client
+                .publish_diagnostics(
+                    uri.clone(),
+                    get_diagnostics(&tree, &rope, &provider, &symbols, &fields),
+                    None,
+                )
+                .await;
+        }
+    }
+}

src/handlers/did_change_configuration.rs

@@ -0,0 +1,15 @@
+use log::warn;
+use tower_lsp::lsp_types::DidChangeConfigurationParams;
+
+use crate::{Backend, Options};
+
+pub async fn did_change_configuration(backend: &Backend, params: DidChangeConfigurationParams) {
+    let Ok(changed_options) = serde_json::from_value::<Options>(params.settings) else {
+        warn!("Unable to parse configuration settings!",);
+        return;
+    };
+    let mut options = backend.options.write().unwrap();
+    options.parser_install_directories = changed_options.parser_install_directories;
+    options.parser_aliases = changed_options.parser_aliases;
+    options.language_retrieval_patterns = changed_options.language_retrieval_patterns;
+}

src/handlers/did_open.rs

@@ -0,0 +1,125 @@
+use std::collections::HashSet;
+
+use log::info;
+use regex::Regex;
+use ropey::Rope;
+use tower_lsp::lsp_types::DidOpenTextDocumentParams;
+use tree_sitter::Parser;
+
+use crate::{
+    util::{get_diagnostics, get_language, TextProviderRope},
+    Backend, SymbolInfo, ENGINE, QUERY_LANGUAGE,
+};
+
+pub async fn did_open(backend: &Backend, params: DidOpenTextDocumentParams) {
+    let uri = &params.text_document.uri;
+    info!("ts_query_ls did_ops: {params:?}");
+    let contents = params.text_document.text;
+    let rope = Rope::from_str(&contents);
+    let mut parser = Parser::new();
+    parser
+        .set_language(&QUERY_LANGUAGE)
+        .expect("Error loading Query grammar");
+    backend.document_map.insert(uri.clone(), rope.clone());
+    backend
+        .cst_map
+        .insert(uri.clone(), parser.parse(&contents, None).unwrap());
+
+    // Get language, if it exists
+    let mut lang = None;
+    if let Ok(options) = backend.options.read() {
+        let mut language_retrieval_regexes: Vec<Regex> = options
+            .language_retrieval_patterns
+            .clone()
+            .unwrap_or(vec![])
+            .iter()
+            .map(|r| Regex::new(r).unwrap())
+            .collect();
+        language_retrieval_regexes.push(Regex::new(r"queries/([^/]+)/[^/]+\.scm$").unwrap());
+        language_retrieval_regexes
+            .push(Regex::new(r"tree-sitter-([^/]+)/queries/[^/]+\.scm$").unwrap());
+        let mut captures = None;
+        for re in language_retrieval_regexes {
+            if let Some(caps) = re.captures(uri.as_str()) {
+                captures = Some(caps);
+                break;
+            }
+        }
+        lang = captures
+            .and_then(|captures| captures.get(1))
+            .and_then(|cap| {
+                let cap_str = cap.as_str();
+                get_language(
+                    options
+                        .parser_aliases
+                        .as_ref()
+                        .and_then(|map| map.get(cap_str))
+                        .unwrap_or(&cap_str.to_owned())
+                        .as_str(),
+                    &options.parser_install_directories,
+                    &ENGINE,
+                )
+            });
+    }
+
+    // Initialize language info
+    let mut symbols_vec: Vec<SymbolInfo> = vec![];
+    let mut symbols_set: HashSet<SymbolInfo> = HashSet::new();
+    let mut fields_vec: Vec<String> = vec![];
+    let mut fields_set: HashSet<String> = HashSet::new();
+    if let Some(lang) = lang {
+        let error_symbol = SymbolInfo {
+            label: "ERROR".to_owned(),
+            named: true,
+        };
+        symbols_set.insert(error_symbol.clone());
+        symbols_vec.push(error_symbol);
+        for i in 0..lang.node_kind_count() as u16 {
+            let named = lang.node_kind_is_named(i);
+            let label = if named {
+                lang.node_kind_for_id(i).unwrap().to_owned()
+            } else {
+                lang.node_kind_for_id(i)
+                    .unwrap()
+                    .replace('\\', r"\\")
+                    .replace('"', r#"\""#)
+                    .replace('\n', r"\n")
+            };
+            let symbol_info = SymbolInfo { label, named };
+            if symbols_set.contains(&symbol_info) || !lang.node_kind_is_visible(i) {
+                continue;
+            }
+            symbols_set.insert(symbol_info.clone());
+            symbols_vec.push(symbol_info);
+        }
+        // Field IDs go from 1 to nfields inclusive (extra index 0 maps to NULL)
+        for i in 1..=lang.field_count() as u16 {
+            let field_name = lang.field_name_for_id(i).unwrap().to_owned();
+            if !fields_set.contains(&field_name) {
+                fields_set.insert(field_name.clone());
+                fields_vec.push(field_name);
+            }
+        }
+    }
+    backend.symbols_vec_map.insert(uri.to_owned(), symbols_vec);
+    backend.symbols_set_map.insert(uri.to_owned(), symbols_set);
+    backend.fields_vec_map.insert(uri.to_owned(), fields_vec);
+    backend.fields_set_map.insert(uri.to_owned(), fields_set);
+
+    // Publish diagnostics
+    if let (Some(tree), Some(symbols), Some(fields)) = (
+        backend.cst_map.get(uri),
+        backend.symbols_set_map.get(uri),
+        backend.fields_set_map.get(uri),
+    ) {
+        let provider = TextProviderRope(&rope);
+        backend
+            .client
+            .publish_diagnostics(
+                uri.clone(),
+                get_diagnostics(&tree, &rope, &provider, &symbols, &fields),
+                None,
+            )
+            .await;
+    }
+}