[mlir:LSP] Add support for keyword code completions

[rriddle]

Summary

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

[jpienaar]

Thanks (sorry missed this one yesterday)

mlir/include/mlir/IR/OpImplementation.h

if (parseOptionalString(string))
       return emitError(loc, "expected string");
     return success();
   }
 
   /// Parse a quoted string token if present.
   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) {
     auto loc = getCurrentLocation();
     if (parseOptionalKeyword(keyword))
       return emitError(loc, "expected valid keyword");
     return success();
   }

mlir/include/mlir/Parser/CodeComplete.h

public:
   /// completions.
   virtual void appendSSAValueCompletion(StringRef name,
                                         std::string typeData) = 0;
 
   /// Append the given block as a code completion result for block name
   /// 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.
   explicit AsmParserCodeCompleteContext(SMLoc codeCompleteLoc)
       : codeCompleteLoc(codeCompleteLoc) {}
 
 private:
   /// The location used to code complete.
   SMLoc codeCompleteLoc;
 };
 } // namespace mlir
 
 #endif // MLIR_PARSER_CODECOMPLETE_H

mlir/lib/Parser/AsmParserImpl.h

if (!parser.getToken().is(Token::string))
       return failure();
 
     if (string)
       *string = parser.getToken().getStringValue();
     parser.consumeToken();
     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)
       return failure();
     parser.consumeToken();
     return success();
   }
 
   /// Parse a keyword, if present, into 'keyword'.
   ParseResult parseOptionalKeyword(StringRef *keyword) override {
     // Check that the current token is a keyword.
     if (!parser.isCurrentTokenAKeyword())
       return failure();
 
     *keyword = parser.getTokenSpelling();
     parser.consumeToken();
     return success();
   }
 
   /// Parse a keyword if it is one of the 'allowedKeywords'.
   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();
 
     StringRef currentKeyword = parser.getTokenSpelling();
     if (llvm::is_contained(allowedKeywords, currentKeyword)) {
       *keyword = currentKeyword;
       parser.consumeToken();

mlir/lib/Parser/Parser.h

public:
   /// returns `failure` to signal that parsing should abort after any desired
   /// completions have been enqueued. Note that `failure` is does not mean
   /// completion failed, it's just a signal to the parser to stop.
 
   ParseResult codeCompleteDialectName();
   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
   /// non-trivial state here, add it to the ParserState class.
   ParserState &state;
 };
 } // namespace detail
 } // namespace mlir
 
 #endif // MLIR_LIB_PARSER_PARSER_H

mlir/lib/Parser/Parser.cpp

ParseResult Parser::codeCompleteStringDialectOrOperationName(StringRef name) {
 
   // Otherwise, we treat this as completing an operation name. The current name
   // is used as the dialect namespace.
   if (name.consume_back("."))
     return codeCompleteOperationName(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
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for parsing operations and regions of
 /// operations.
 class OperationParser : public Parser {

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

void appendBlockCompletion(StringRef name) final {
     bool stripPrefix = getCodeCompleteLoc().getPointer()[-1] == '^';
 
     lsp::CompletionItem item(name, lsp::CompletionItemKind::Field);
     if (stripPrefix)
       item.insertText = name.drop_front(1).str();
     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;
 };
 } // namespace
 
 lsp::CompletionList
 MLIRDocument::getCodeCompletion(const lsp::URIForFile &uri,

mlir/test/mlir-lsp-server/completion.test

 // CHECK-NEXT:        "insertText": "arg",
 // CHECK-NEXT:        "insertTextFormat": 1,
 // CHECK-NEXT:        "kind": 6,
 // CHECK-NEXT:        "label": "%arg"
 // CHECK-NEXT:      }
 // 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"}