[mlir:LSP] Add support for keyword code completions

This commit adds code completion results to the MLIR LSP when
parsing keywords. Keyword support is currently limited to the
case where the expected keyword is provided, but a followup will
work on expanding the set of keyword cases we handle (e.g. to
allow capturing attribute/type mnemonics).

Differential Revision: https://reviews.llvm.org/D129184

diff --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index 84e71c4..33967b0 100644 (file)
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/mlir/include/mlir/IR/OpImplementation.h
@@ -572,12 +572,10 @@ public:
   virtual ParseResult parseOptionalString(std::string *string) = 0;
 
   /// Parse a given keyword.
-  ParseResult parseKeyword(StringRef keyword, const Twine &msg = "") {
-    auto loc = getCurrentLocation();
-    if (parseOptionalKeyword(keyword))
-      return emitError(loc, "expected '") << keyword << "'" << msg;
-    return success();
+  ParseResult parseKeyword(StringRef keyword) {
+    return parseKeyword(keyword, "");
   }
+  virtual ParseResult parseKeyword(StringRef keyword, const Twine &msg) = 0;
 
   /// Parse a keyword into 'keyword'.
   ParseResult parseKeyword(StringRef *keyword) {

diff --git a/mlir/include/mlir/Parser/CodeComplete.h b/mlir/include/mlir/Parser/CodeComplete.h
index 24d0181..7dd7b61 100644 (file)
--- a/mlir/include/mlir/Parser/CodeComplete.h
+++ b/mlir/include/mlir/Parser/CodeComplete.h
@@ -43,6 +43,11 @@ public:
   /// completions.
   virtual void appendBlockCompletion(StringRef name) = 0;
 
+  /// Signal a completion for the given expected tokens, which are optional if
+  /// `optional` is set.
+  virtual void completeExpectedTokens(ArrayRef<StringRef> tokens,
+                                      bool optional) = 0;
+
 protected:
   /// Create a new code completion context with the given code complete
   /// location.

diff --git a/mlir/lib/Parser/AsmParserImpl.h b/mlir/lib/Parser/AsmParserImpl.h
index 98fc824..e15738e 100644 (file)
--- a/mlir/lib/Parser/AsmParserImpl.h
+++ b/mlir/lib/Parser/AsmParserImpl.h
@@ -242,8 +242,21 @@ public:
     return success();
   }
 
+  ParseResult parseKeyword(StringRef keyword, const Twine &msg) override {
+    if (parser.getToken().isCodeCompletion())
+      return parser.codeCompleteExpectedTokens(keyword);
+
+    auto loc = getCurrentLocation();
+    if (parseOptionalKeyword(keyword))
+      return emitError(loc, "expected '") << keyword << "'" << msg;
+    return success();
+  }
+
   /// Parse the given keyword if present.
   ParseResult parseOptionalKeyword(StringRef keyword) override {
+    if (parser.getToken().isCodeCompletion())
+      return parser.codeCompleteOptionalTokens(keyword);
+
     // Check that the current token has the same spelling.
     if (!parser.isCurrentTokenAKeyword() ||
         parser.getTokenSpelling() != keyword)
@@ -267,6 +280,9 @@ public:
   ParseResult
   parseOptionalKeyword(StringRef *keyword,
                        ArrayRef<StringRef> allowedKeywords) override {
+    if (parser.getToken().isCodeCompletion())
+      return parser.codeCompleteOptionalTokens(allowedKeywords);
+
     // Check that the current token is a keyword.
     if (!parser.isCurrentTokenAKeyword())
       return failure();

diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 51a65a9..be4f886 100644 (file)
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -395,6 +395,15 @@ ParseResult Parser::codeCompleteStringDialectOrOperationName(StringRef name) {
   return failure();
 }
 
+ParseResult Parser::codeCompleteExpectedTokens(ArrayRef<StringRef> tokens) {
+  state.codeCompleteContext->completeExpectedTokens(tokens, /*optional=*/false);
+  return failure();
+}
+ParseResult Parser::codeCompleteOptionalTokens(ArrayRef<StringRef> tokens) {
+  state.codeCompleteContext->completeExpectedTokens(tokens, /*optional=*/true);
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // OperationParser
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Parser/Parser.h b/mlir/lib/Parser/Parser.h
index f1c87ca..34a95a5 100644 (file)
--- a/mlir/lib/Parser/Parser.h
+++ b/mlir/lib/Parser/Parser.h
@@ -319,6 +319,8 @@ public:
   ParseResult codeCompleteOperationName(StringRef dialectName);
   ParseResult codeCompleteDialectOrElidedOpName(SMLoc loc);
   ParseResult codeCompleteStringDialectOrOperationName(StringRef name);
+  ParseResult codeCompleteExpectedTokens(ArrayRef<StringRef> tokens);
+  ParseResult codeCompleteOptionalTokens(ArrayRef<StringRef> tokens);
 
 protected:
   /// The Parser is subclassed and reinstantiated.  Do not add additional

diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
index c209774..9d3d51b 100644 (file)
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
@@ -690,6 +690,16 @@ public:
     completionList.items.emplace_back(item);
   }
 
+  /// Signal a completion for the given expected token.
+  void completeExpectedTokens(ArrayRef<StringRef> tokens, bool optional) final {
+    for (StringRef token : tokens) {
+      lsp::CompletionItem item(token, lsp::CompletionItemKind::Keyword);
+      item.sortText = "0";
+      item.detail = optional ? "optional" : "";
+      completionList.items.emplace_back(item);
+    }
+  }
+
 private:
   lsp::CompletionList &completionList;
   MLIRContext *ctx;

diff --git a/mlir/test/mlir-lsp-server/completion.test b/mlir/test/mlir-lsp-server/completion.test
index b13dbe1..4264fe6 100644 (file)
--- a/mlir/test/mlir-lsp-server/completion.test
+++ b/mlir/test/mlir-lsp-server/completion.test
@@ -100,6 +100,39 @@
 // CHECK:         ]
 // CHECK-NEXT:  }
 // -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":0,"character":10}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "detail": "optional",
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "public",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "detail": "optional",
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "private",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "detail": "optional",
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "nested",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      }
+// CHECK-NEXT:    ]
+// CHECK-NEXT:  }
+// -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}