fix(ls): code completion suggestions that don't make too much sense (VirusTotal#549)

Fix issues with code completion making suggestions that doesn't make sense, like keywords after a `.` character.

Author: prosperritty

ls/src/features/completion.rs

@@ -1,6 +1,7 @@
 use async_lsp::lsp_types::{
-    CompletionItem, CompletionItemKind, CompletionItemLabelDetails,
-    InsertTextFormat, InsertTextMode, Position,
+    CompletionContext, CompletionItem, CompletionItemKind,
+    CompletionItemLabelDetails, CompletionTriggerKind, InsertTextFormat,
+    InsertTextMode, Position,
 };
 use itertools::Itertools;
 
@@ -76,6 +77,7 @@ const CONDITION_SUGGESTIONS: [(&str, Option<&str>); 16] = [
 pub fn completion(
     document: &Document,
     pos: Position,
+    context: Option<CompletionContext>,
 ) -> Option<Vec<CompletionItem>> {
     let cst = &document.cst;
     // Get the token before cursor. There might be no token at cursor when the
@@ -84,13 +86,29 @@ pub fn completion(
         .and_then(|token| token.prev_token())
         .or_else(|| cst.root().last_token())?;
 
+    // Trigger characters are: `.`, `!`, `$`, `@`, `#`.
+    let is_trigger_character = context.is_some_and(|ctx| {
+        ctx.trigger_kind == CompletionTriggerKind::TRIGGER_CHARACTER
+    });
+
     // If the token is a direct child of `SOURCE_FILE`, return top-level suggestions.
-    if non_error_parent(&token)?.kind() == SyntaxKind::SOURCE_FILE {
+    if !is_trigger_character
+        && non_error_parent(&token)?.kind() == SyntaxKind::SOURCE_FILE
+    {
         return Some(source_file_suggestions());
     }
 
     let prev_token = prev_non_trivia_token(&token)?;
 
+    if prev_token.ancestors().any(|n| n.kind() == SyntaxKind::CONDITION_BLK) {
+        return condition_suggestions(cst, token);
+    }
+
+    // Trigger characters are recognized in the condition block only.
+    if is_trigger_character {
+        return Some(vec![]);
+    }
+
     if prev_token.kind() == SyntaxKind::IMPORT_KW {
         #[cfg(feature = "full-compiler")]
         return Some(import_suggestions());
@@ -104,10 +122,6 @@ pub fn completion(
         return Some(pattern_modifier_suggestions(pattern_def));
     }
 
-    if prev_token.ancestors().any(|n| n.kind() == SyntaxKind::CONDITION_BLK) {
-        return condition_suggestions(cst, token);
-    }
-
     if prev_token.ancestors().any(|n| n.kind() == SyntaxKind::RULE_DECL) {
         return Some(rule_suggestions());
     }
@@ -192,6 +206,8 @@ fn condition_suggestions(
                 }
             }
         }
+        // Do not propose keywords for condition block after a dot.
+        SyntaxKind::DOT => {}
         _ => {
             CONDITION_SUGGESTIONS.iter().for_each(|(kw, insert)| {
                 result.push(CompletionItem {

ls/src/server.rs

@@ -290,13 +290,14 @@ impl LanguageServer for YARALanguageServer {
     {
         let uri = params.text_document_position.text_document.uri;
         let position = params.text_document_position.position;
+        let context = params.context;
         let document = match self.documents.get(&uri) {
             Some(entry) => entry,
             None => return Box::pin(async { Ok(None) }),
         };
 
-        let completions =
-            completion(document, position).map(CompletionResponse::Array);
+        let completions = completion(document, position, context)
+            .map(CompletionResponse::Array);
 
         Box::pin(async move { Ok(completions) })
     }

ls/src/tests/mod.rs

@@ -267,6 +267,12 @@ async fn completion() {
 
     #[cfg(all(feature = "full-compiler", not(feature = "magic-module")))]
     test_lsp_request::<_, Completion>("completion11.yar").await;
+
+    #[cfg(all(feature = "full-compiler", not(feature = "magic-module")))]
+    test_lsp_request::<_, Completion>("completion12.yar").await;
+
+    #[cfg(all(feature = "full-compiler", not(feature = "magic-module")))]
+    test_lsp_request::<_, Completion>("completion13.yar").await;
 }
 
 #[tokio::test]

ls/src/tests/testdata/completion12.request.json

@@ -0,0 +1,13 @@
+{
+  "textDocument": {
+    "uri": "file:///completion12.yar"
+  },
+  "position": {
+    "line": 2,
+    "character": 15
+  },
+  "context": {
+    "triggerKind": 2,
+    "triggerCharacter": "."
+  }
+}

ls/src/tests/testdata/completion12.response.json

@@ -0,0 +1 @@
+[]

ls/src/tests/testdata/completion12.yar

@@ -0,0 +1,6 @@
+rule test {
+  strings:
+    $a = "foo".
+  condition:
+    $a
+}

ls/src/tests/testdata/completion13.request.json

@@ -0,0 +1,13 @@
+{
+  "textDocument": {
+    "uri": "file:///completion13.yar"
+  },
+  "position": {
+    "line": 3,
+    "character": 24
+  },
+  "context": {
+    "triggerKind": 2,
+    "triggerCharacter": "."
+  }
+}

ls/src/tests/testdata/completion13.response.json

@@ -0,0 +1 @@
+[]

ls/src/tests/testdata/completion13.yar

@@ -0,0 +1,5 @@
+import "pe"
+rule test {
+  condition:
+    pe.data_directories.
+}