[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

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) {

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.

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();

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
 //===----------------------------------------------------------------------===//

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

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;

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"}