Moved over to TS Query API for checking node type

src/lsp.rs

@@ -54,6 +54,7 @@ pub fn get_word_from_file_params(
     }
 }
 
+// TODO: Replace with tree-sitter Query API
 // BUGBUG: Add some sort of depth limit here?
 pub fn at_node_type(
     root: &tree_sitter::Node,
@@ -161,50 +162,168 @@ pub fn get_comp_resp(
     instr_comps: &[CompletionItem],
     reg_comps: &[CompletionItem],
 ) -> Option<CompletionList> {
-    let mut is_register = false;
-    let mut is_instruction = false;
+    let cursor_line = params.text_document_position.position.line as usize;
+    let cursor_char = params.text_document_position.position.character as usize;
+    let mut comp_items = None;
 
     // prepend register names with "%" in GAS
     if let Some(ctx) = params.context.as_ref() {
         if ctx.trigger_kind == CompletionTriggerKind::TRIGGER_CHARACTER {
-            is_register = true;
+            return Some(CompletionList {
+                is_incomplete: true,
+                items: reg_comps.to_owned(),
+            });
         }
     }
 
-    if !is_register {
-        // would like to use incremental parsing, but can't figure out
-        // translating lsp_server's edit info to the data format tree sitter is
-        // expecting...
-        let new_tree = parser.parse(curr_doc, None);
-        let line = params.text_document_position.position.line as usize;
-        let col = params.text_document_position.position.character as usize;
-
-        if let Some(ref tree) = new_tree {
-            is_register = at_node_type(&tree.root_node(), TSASMNodeType::sym_reg, line, col);
-
-            if !is_register {
-                // not sure why this matches instructions but sym_instruction doesn't
-                is_instruction =
-                    at_node_type(&tree.root_node(), TSASMNodeType::sym_string, line, col);
+    // TODO: filter by width allowed by corresponding instruction
+    // Would like to do incremental parsing but don't see a straightforward
+    // way to map the LSP's edit info the the format/ info tree-sitter is
+    // expecting...
+    let curr_tree = parser.parse(curr_doc, None);
+    if let Some(tree) = curr_tree {
+        let mut cursor = tree_sitter::QueryCursor::new();
+        cursor.set_point_range(std::ops::Range {
+            start: tree_sitter::Point {
+                row: cursor_line,
+                column: 0,
+            },
+            end: tree_sitter::Point {
+                row: cursor_line,
+                column: usize::MAX,
+            },
+        });
+
+        // Instruction and two register arguments
+        let query = match tree_sitter::Query::new(
+            tree_sitter_asm::language(),
+            "(instruction kind: (word) @instr_name (ident (reg) @r1) (ident (reg) @r2))",
+        ) {
+            Ok(query_) => query_,
+            Err(e) => {
+                error!(
+                    "Error creating tree-sitter instruction register register query - Error: {e}"
+                );
+                return None;
+            }
+        };
+
+        let matches: Vec<tree_sitter::QueryMatch<'_, '_>> = cursor
+            .matches(&query, tree.root_node(), curr_doc.as_bytes())
+            .collect();
+        if !matches.is_empty() && matches[0].captures.len() == 3 {
+            let instr_start = matches[0].captures[0].node.range().start_point;
+            let instr_end = matches[0].captures[0].node.range().end_point;
+            let reg_1_start = matches[0].captures[1].node.range().start_point;
+            let reg_1_end = matches[0].captures[1].node.range().end_point;
+            let reg_2_start = matches[0].captures[2].node.range().start_point;
+            let reg_2_end = matches[0].captures[2].node.range().end_point;
+            if instr_start.row == cursor_line
+                && instr_end.row == cursor_line
+                && instr_start.column <= cursor_char
+                && instr_end.column >= cursor_char
+            {
+                comp_items = Some(instr_comps.to_owned());
+            } else if (reg_1_start.row == cursor_line
+                // match with first register
+                && reg_1_end.row == cursor_line
+                && reg_1_start.column <= cursor_char
+                && reg_1_end.column >= cursor_char)
+                // match with second register
+                || (reg_2_start.row == cursor_line
+                    && reg_2_end.row == cursor_line
+                    && reg_2_start.column <= cursor_char
+                    && reg_2_end.column >= cursor_char)
+            {
+                comp_items = Some(reg_comps.to_owned());
+            }
+            if let Some(items) = comp_items {
+                return Some(CompletionList {
+                    is_incomplete: true,
+                    items,
+                });
             }
         }
-    }
 
-    let items: Vec<CompletionItem> = match (is_register, is_instruction) {
-        // We know we're at a register node
-        (true, false) => reg_comps.to_owned(),
-        // We know we're at an instruction node
-        (false, true) => instr_comps.to_owned(),
-        // We don't know where we are, offer nothing
-        _ => {
-            return None;
+        // Instruction and one register argument
+        let query = match tree_sitter::Query::new(
+            tree_sitter_asm::language(),
+            "(instruction kind: (word) @instr_name (ident (reg) @r1))",
+        ) {
+            Ok(query_) => query_,
+
+            Err(e) => {
+                error!("Error creating tree-sitter instruction register query - Error: {e}");
+                return None;
+            }
+        };
+
+        let matches: Vec<tree_sitter::QueryMatch<'_, '_>> = cursor
+            .matches(&query, tree.root_node(), curr_doc.as_bytes())
+            .collect();
+        if !matches.is_empty() && matches[0].captures.len() == 2 {
+            let instr_start = matches[0].captures[0].node.range().start_point;
+            let instr_end = matches[0].captures[0].node.range().end_point;
+            let reg_start = matches[0].captures[1].node.range().start_point;
+            let reg_end = matches[0].captures[1].node.range().end_point;
+            if instr_start.row == cursor_line
+                && instr_end.row == cursor_line
+                && instr_start.column <= cursor_char
+                && instr_end.column >= cursor_char
+            {
+                comp_items = Some(instr_comps.to_owned());
+            } else if reg_start.row == cursor_line
+                && reg_end.row == cursor_line
+                && reg_start.column <= cursor_char
+                && reg_end.column >= cursor_char
+            {
+                comp_items = Some(reg_comps.to_owned());
+            }
+
+            if let Some(items) = comp_items {
+                return Some(CompletionList {
+                    is_incomplete: true,
+                    items,
+                });
+            }
         }
-    };
 
-    Some(CompletionList {
-        is_incomplete: true,
-        items,
-    })
+        // Just an instruction
+        let query = match tree_sitter::Query::new(
+            tree_sitter_asm::language(),
+            "(instruction kind: (word) @instr_name)",
+        ) {
+            Ok(query_) => query_,
+            Err(e) => {
+                error!("Error creating tree-sitter instruction query - Error: {e}");
+                return None;
+            }
+        };
+
+        let matches: Vec<tree_sitter::QueryMatch<'_, '_>> = cursor
+            .matches(&query, tree.root_node(), curr_doc.as_bytes())
+            .collect();
+        if !matches.is_empty() && matches[0].captures.len() == 1 {
+            let instr_start = matches[0].captures[0].node.range().start_point;
+            let instr_end = matches[0].captures[0].node.range().end_point;
+            if instr_start.row == cursor_line
+                && instr_end.row == cursor_line
+                && instr_start.column <= cursor_char
+                && instr_end.column >= cursor_char
+            {
+                comp_items = Some(instr_comps.to_owned());
+            }
+
+            if let Some(items) = comp_items {
+                return Some(CompletionList {
+                    is_incomplete: true,
+                    items,
+                });
+            }
+        }
+    }
+
+    None
 }
 
 // Note: Some issues here regarding entangled lifetimes