[mlir:LSP] Add support for MLIR code completions

[rriddle]

Summary

This commit adds code completion results to the MLIR LSP using
a new code completion context in the MLIR parser. This commit
adds initial completion for dialect, operation, SSA value, and
block names.

[jpienaar]

mlir/lib/Parser/Parser.h

314

This feels a bit weird. Is ParseResult correct type here?

[rriddle]

mlir/lib/Parser/Parser.h

314

Either ParseResult or LogicalResult is what we want. We really just want a silent failure that stops the parsing process, given that reaching a code completion point is generally a sign that we've done all that we wanted to do during parsing.

[jpienaar]

Looks good in general thanks!

mlir/lib/Parser/Parser.h

314

What does failure signify here? That it did successfully code complete? Why not have it something like "bool hasCodeCompletedDialectName()" and then at the call site decide what to do with it? (this reminds me too much of Status situations where folks want to use RETURN_IF_ERROR macro for convenience and so leaf functions end up with Status return types while they aren't really exceptional cases, and so this may irritate me more than most readers :)). Returning failure on success feels weird, it makes sense at the caller, but for the function itself it feels odd as failure is successfully performed task, and success means nothing happened. [not blocker]

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

643

What does this do? Just give ordering preference when returning completions?

707

Why is this needed? Is completion just heavier and want to avoid it unless requested? Or is it because we abort on complete above that we couldn't do it always?

717

kwargs style comment for nullptr?

[rriddle]

mlir/lib/Parser/Parser.h

314

It's less about the success of completion, completion is essentially just enqueueing something in the LSP, and more about convenience and keeping the meaning of the result/failure specific to the completion methods. I agree that we are using "failure" here, but what we really want to do is interrupt the parser process (given that these methods always "succeed"). One alternative would be sprinkle return failure() in each place this is called, but that one doesn't seem as nice to me given that here we can document and define what the failure is supposed to mean (plus it's not as clean). Another would be to have a special/local result type used here that is self documenting. I have a slight preference to keep as-is, given that this aligns with how we ended up doing it in PDLL, but happy to change if there is a strong preference otherwise.

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

643

Yeah, sortText is used to determine the ordering of the results (if not present, the label field gets used for this).

707

The latter for a couple reasons. "Code completion" is a "special" mode and we want to be able to inject completions that can only come from being in the immediate context (e.g. keywords, dialect names, etc.). Aside from that, completion often runs on incomplete IR (given that it runs while you type), so we want to be able to build tokens in special ways depending on where the completion location is.

717

Thanks for the catch.

mlir/include/mlir/Parser/CodeComplete.h

+//===- CodeComplete.h - MLIR Asm CodeComplete Context -----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_PARSER_CODECOMPLETE_H
+#define MLIR_PARSER_CODECOMPLETE_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/Support/SourceMgr.h"
+
+namespace mlir {
+/// This class provides an abstract interface into the parser for hooking in
+/// code completion events. This class is only really useful for providing
+/// language tooling for MLIR, general clients should not need to use this
+/// class.
+class AsmParserCodeCompleteContext {
+public:
+  virtual ~AsmParserCodeCompleteContext();
+
+  /// Return the source location used to provide code completion.
+  SMLoc getCodeCompleteLoc() const { return codeCompleteLoc; }
+
+  //===--------------------------------------------------------------------===//
+  // Completion Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Signal code completion for a dialect name.
+  virtual void completeDialectName() = 0;
+
+  /// Signal code completion for an operation name within the given dialect.
+  virtual void completeOperationName(StringRef dialectName) = 0;
+
+  /// Append the given SSA value as a code completion result for SSA value
+  /// 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;
+
+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/include/mlir/Parser/Parser.h

 namespace llvm {
 class SourceMgr;
 class SMDiagnostic;
 class StringRef;
 } // namespace llvm
 
 namespace mlir {
 class AsmParserState;
+class AsmParserCodeCompleteContext;
 
 namespace detail {
 
 /// Given a block containing operations that have just been parsed, if the block
 /// contains a single operation of `ContainerOpT` type then remove it from the
 /// block and return it. If the block does not contain just that operation,
 /// create a new operation instance of `ContainerOpT` and move all of the
 /// operations within `parsedBlock` into the first block of the first region.

 /// operations to the given block. If the block is non-empty, the operations are
 /// placed before the current terminator. If parsing is successful, success is
 /// returned. Otherwise, an error message is emitted through the error handler
 /// registered in the context, and failure is returned. If `sourceFileLoc` is
 /// non-null, it is populated with a file location representing the start of the
 /// source file that is being parsed. If `asmState` is non-null, it is populated
 /// with detailed information about the parsed IR (including exact locations for
 /// SSA uses and definitions). `asmState` should only be provided if this
-/// detailed information is desired.
-LogicalResult parseSourceFile(const llvm::SourceMgr &sourceMgr, Block *block,
-                              const ParserConfig &config,
-                              LocationAttr *sourceFileLoc = nullptr,
-                              AsmParserState *asmState = nullptr);
+/// detailed information is desired. If `codeCompleteContext` is non-null, it is
+/// used to signal tracking of a code completion event (generally only ever
+/// useful for LSP or other high level language tooling).
+LogicalResult parseSourceFile(
+    const llvm::SourceMgr &sourceMgr, Block *block, const ParserConfig &config,
+    LocationAttr *sourceFileLoc = nullptr, AsmParserState *asmState = nullptr,
+    AsmParserCodeCompleteContext *codeCompleteContext = nullptr);
 
 /// This parses the file specified by the indicated filename and appends parsed
 /// operations to the given block. If the block is non-empty, the operations are
 /// placed before the current terminator. If parsing is successful, success is
 /// returned. Otherwise, an error message is emitted through the error handler
 /// registered in the context, and failure is returned. If `sourceFileLoc` is
 /// non-null, it is populated with a file location representing the start of the
 /// source file that is being parsed.

mlir/lib/Parser/AffineParser.cpp

IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context,
                                  bool printDiagnosticInfo) {
   llvm::SourceMgr sourceMgr;
   auto memBuffer = llvm::MemoryBuffer::getMemBuffer(
       inputStr, /*BufferName=*/"<mlir_parser_buffer>",
       /*RequiresNullTerminator=*/false);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
   SymbolState symbolState;
   ParserConfig config(context);
-  ParserState state(sourceMgr, config, symbolState, /*asmState=*/nullptr);
+  ParserState state(sourceMgr, config, symbolState, /*asmState=*/nullptr,
+                    /*codeCompleteContext=*/nullptr);
   Parser parser(state);
 
   raw_ostream &os = printDiagnosticInfo ? llvm::errs() : llvm::nulls();
   SourceMgrDiagnosticHandler handler(sourceMgr, context, os);
   IntegerSet set;
   if (parser.parseIntegerSetReference(set))
     return IntegerSet();
 
   Token endTok = parser.getToken();
   if (endTok.isNot(Token::eof)) {
     parser.emitError(endTok.getLoc(), "encountered unexpected token");
     return IntegerSet();
   }
 
   return set;
 }

mlir/lib/Parser/DialectSymbolParser.cpp

static T parseSymbol(StringRef inputStr, MLIRContext *context, size_t &numRead,
                      ParserFn &&parserFn) {
   SourceMgr sourceMgr;
   auto memBuffer = MemoryBuffer::getMemBuffer(
       inputStr, /*BufferName=*/"<mlir_parser_buffer>",
       /*RequiresNullTerminator=*/false);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
   SymbolState aliasState;
   ParserConfig config(context);
-  ParserState state(sourceMgr, config, aliasState, /*asmState=*/nullptr);
+  ParserState state(sourceMgr, config, aliasState, /*asmState=*/nullptr,
+                    /*codeCompleteContext=*/nullptr);
   Parser parser(state);
 
   SourceMgrDiagnosticHandler handler(
       const_cast<llvm::SourceMgr &>(parser.getSourceMgr()),
       parser.getContext());
   Token startTok = parser.getToken();
   T symbol = parserFn(parser);
   if (!symbol)

mlir/lib/Parser/Lexer.h

 #include "mlir/Parser/Parser.h"
 
 namespace mlir {
 class Location;
 
 /// This class breaks up the current file into a token stream.
 class Lexer {
 public:
-  explicit Lexer(const llvm::SourceMgr &sourceMgr, MLIRContext *context);
+  explicit Lexer(const llvm::SourceMgr &sourceMgr, MLIRContext *context,
+                 AsmParserCodeCompleteContext *codeCompleteContext);
 
   const llvm::SourceMgr &getSourceMgr() { return sourceMgr; }
 
   Token lexToken();
 
   /// Encode the specified source location information into a Location object
   /// for attachment to the IR or error reporting.
   Location getEncodedSourceLocation(SMLoc loc);
private:
   void skipComment();
 
   const llvm::SourceMgr &sourceMgr;
   MLIRContext *context;
 
   StringRef curBuffer;
   const char *curPtr;
 
+  /// An optional code completion point within the input file, used to indicate
+  /// the position of a code completion token.
+  const char *codeCompleteLoc;
+
   Lexer(const Lexer &) = delete;
   void operator=(const Lexer &) = delete;
 };
 
 } // namespace mlir
 
 #endif // MLIR_LIB_PARSER_LEXER_H

mlir/lib/Parser/Lexer.cpp

 //===- Lexer.cpp - MLIR Lexer Implementation ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the lexer for the MLIR textual form.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Lexer.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/Parser/CodeComplete.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
 // Returns true if 'c' is an allowable punctuation character: [$._-]
 // Returns false otherwise.
 static bool isPunct(char c) {
   return c == '$' || c == '.' || c == '_' || c == '-';
 }
 
-Lexer::Lexer(const llvm::SourceMgr &sourceMgr, MLIRContext *context)
-    : sourceMgr(sourceMgr), context(context) {
+Lexer::Lexer(const llvm::SourceMgr &sourceMgr, MLIRContext *context,
+             AsmParserCodeCompleteContext *codeCompleteContext)
+    : sourceMgr(sourceMgr), context(context), codeCompleteLoc(nullptr) {
   auto bufferID = sourceMgr.getMainFileID();
   curBuffer = sourceMgr.getMemoryBuffer(bufferID)->getBuffer();
   curPtr = curBuffer.begin();
+
+  // Set the code completion location if it was provided.
+  if (codeCompleteContext)
+    codeCompleteLoc = codeCompleteContext->getCodeCompleteLoc().getPointer();
 }
 
 /// Encode the specified source location information into an attribute for
 /// attachment to the IR.
 Location Lexer::getEncodedSourceLocation(SMLoc loc) {
   auto &sourceMgr = getSourceMgr();
   unsigned mainFileID = sourceMgr.getMainFileID();
 
Token Lexer::emitError(const char *loc, const Twine &message) {
   mlir::emitError(getEncodedSourceLocation(SMLoc::getFromPointer(loc)),
                   message);
   return formToken(Token::error, loc);
 }
 
 Token Lexer::lexToken() {
   while (true) {
     const char *tokStart = curPtr;
+
+    // Check to see if the current token is at the code completion location.
+    if (tokStart == codeCompleteLoc)
+      return formToken(Token::code_complete, tokStart);
+
+    // Lex the next token.
     switch (*curPtr++) {
     default:
       // Handle bare identifiers.
       if (isalpha(curPtr[-1]))
         return lexBareIdentifierOrKeyword(tokStart);
 
       // Unknown character, emit an error.
       return emitError(tokStart, "unexpected character");
case '!':
     break;
   default:
     llvm_unreachable("invalid caller");
   }
 
   // Parse suffix-id.
   if (isdigit(*curPtr)) {
     // If suffix-id starts with a digit, the rest must be digits.
-    while (isdigit(*curPtr)) {
+    while (isdigit(*curPtr))
       ++curPtr;
-    }
   } else if (isalpha(*curPtr) || isPunct(*curPtr)) {
     do {
       ++curPtr;
     } while (isalpha(*curPtr) || isdigit(*curPtr) || isPunct(*curPtr));
+  } else if (curPtr == codeCompleteLoc) {
+    return formToken(Token::code_complete, tokStart);
   } else {
     return emitError(curPtr - 1, errorKind);
   }
 
+  // Check for a code completion within the identifier.
+  if (codeCompleteLoc && codeCompleteLoc >= tokStart &&
+      codeCompleteLoc <= curPtr) {
+    return Token(Token::code_complete,
+                 StringRef(tokStart, codeCompleteLoc - tokStart));
+  }
+
   return formToken(kind, tokStart);
 }
 
 /// Lex a string literal.
 ///
 ///   string-literal ::= '"' [^"\n\f\v\r]* '"'
 ///
 /// TODO: define escaping rules.
 Token Lexer::lexString(const char *tokStart) {
   assert(curPtr[-1] == '"');
 
   while (true) {
+    // Check to see if there is a code completion location within the string. In
+    // these cases we generate a completion location and place the currently
+    // lexed string within the token. This allows for the parser to use the
+    // partially lexed string when computing the completion results.
+    if (curPtr == codeCompleteLoc)
+      return formToken(Token::code_complete, tokStart);
+
     switch (*curPtr++) {
     case '"':
       return formToken(Token::string, tokStart);
     case 0:
       // If this is a random nul character in the middle of a string, just
       // include it.  If it is the end of file, then it is an error.
       if (curPtr - 1 != curBuffer.end())
         continue;

mlir/lib/Parser/Parser.h

parseAffineMapOfSSAIds(AffineMap &map,
                          function_ref<ParseResult(bool)> parseElement,
                          Delimiter delimiter);
 
   /// Parse an AffineExpr where dim and symbol identifiers are SSA ids.
   ParseResult
   parseAffineExprOfSSAIds(AffineExpr &expr,
                           function_ref<ParseResult(bool)> parseElement);
 
+  //===--------------------------------------------------------------------===//
+  // Code Completion
+  //===--------------------------------------------------------------------===//
+
+  /// The set of various code completion methods. Every completion method
+  /// returns `failure` to stop the parsing process after providing completion
+  /// results.
+
+  ParseResult codeCompleteDialectName();
+  ParseResult codeCompleteOperationName(StringRef dialectName);
+  ParseResult codeCompleteDialectOrElidedOpName(SMLoc loc);
+  ParseResult codeCompleteStringDialectOrOperationName(StringRef name);
+
 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

 #include "Parser.h"
 #include "AsmParserImpl.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Parser/AsmParserState.h"
+#include "mlir/Parser/CodeComplete.h"
 #include "mlir/Parser/Parser.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/bit.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/SourceMgr.h"
 #include <algorithm>
 
 using namespace mlir;
 using namespace mlir::detail;
 using llvm::MemoryBuffer;
 using llvm::SourceMgr;
 
 //===----------------------------------------------------------------------===//
+// CodeComplete
+//===----------------------------------------------------------------------===//
+
+AsmParserCodeCompleteContext::~AsmParserCodeCompleteContext() = default;
+
+//===----------------------------------------------------------------------===//
 // Parser
 //===----------------------------------------------------------------------===//
 
 /// Parse a list of comma-separated items with an optional delimiter.  If a
 /// delimiter is provided, then an empty list is allowed.  If not, then at
 /// least one element will be parsed.
 ParseResult
 Parser::parseCommaSeparatedList(Delimiter delimiter,
if (entry.first.empty()) {
     entry.second = *result;
   }
 
   name = entry.first;
   return entry.second;
 }
 
 //===----------------------------------------------------------------------===//
+// Code Completion
+
+ParseResult Parser::codeCompleteDialectName() {
+  state.codeCompleteContext->completeDialectName();
+  return failure();
+}
+
+ParseResult Parser::codeCompleteOperationName(StringRef dialectName) {
+  // Perform some simple validation on the dialect name. This doesn't need to be
+  // extensive, it's more of an optimization (to avoid checking completion
+  // results when we know they will fail).
+  if (dialectName.empty() || dialectName.contains('.'))
+    return failure();
+  state.codeCompleteContext->completeOperationName(dialectName);
+  return failure();
+}
+
+ParseResult Parser::codeCompleteDialectOrElidedOpName(SMLoc loc) {
+  // Check to see if there is anything else on the current line. This check
+  // isn't strictly necessary, but it does avoid unnecessarily triggering
+  // completions for operations and dialects in situations where we don't want
+  // them (e.g. at the end of an operation).
+  auto shouldIgnoreOpCompletion = [&]() {
+    const char *bufBegin = state.lex.getBufferBegin();
+    const char *it = loc.getPointer() - 1;
+    for (; it > bufBegin && *it != '\n'; --it)
+      if (!llvm::is_contained(StringRef(" \t\r"), *it))
+        return true;
+    return false;
+  };
+  if (shouldIgnoreOpCompletion())
+    return failure();
+
+  // The completion here is either for a dialect name, or an operation name
+  // whose dialect prefix was elided. For this we simply invoke both of the
+  // individual completion methods.
+  (void)codeCompleteDialectName();
+  return codeCompleteOperationName(state.defaultDialectStack.back());
+}
+
+ParseResult Parser::codeCompleteStringDialectOrOperationName(StringRef name) {
+  // If the name is empty, this is the start of the string and contains the
+  // dialect.
+  if (name.empty())
+    return codeCompleteDialectName();
+
+  // 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();
+}
+
+//===----------------------------------------------------------------------===//
 // OperationParser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for parsing operations and regions of
 /// operations.
 class OperationParser : public Parser {
 public:
public:
   /// Parse a (possibly empty) list of block arguments.
   ParseResult parseOptionalBlockArgList(Block *owner);
 
   /// Get the block with the specified name, creating it if it doesn't
   /// already exist.  The location specified is the point of use, which allows
   /// us to diagnose references to blocks that are not defined precisely.
   Block *getBlockNamed(StringRef name, SMLoc loc);
 
+  //===--------------------------------------------------------------------===//
+  // Code Completion
+  //===--------------------------------------------------------------------===//
+
+  /// The set of various code completion methods. Every completion method
+  /// returns `failure` to stop the parsing process after providing completion
+  /// results.
+
+  ParseResult codeCompleteSSAUse();
+  ParseResult codeCompleteBlock();
+
 private:
   /// This class represents a definition of a Block.
   struct BlockDefinition {
     /// A pointer to the defined Block.
     Block *block;
     /// The location that the Block was defined at.
     SMLoc loc;
   };

 
 /// Parse a (possibly empty) list of SSA operands.
 ///
 ///   ssa-use-list ::= ssa-use (`,` ssa-use)*
 ///   ssa-use-list-opt ::= ssa-use-list?
 ///
 ParseResult OperationParser::parseOptionalSSAUseList(
     SmallVectorImpl<UnresolvedOperand> &results) {
-  if (getToken().isNot(Token::percent_identifier))
+  if (!getToken().isOrIsCodeCompletionFor(Token::percent_identifier))
     return success();
   return parseCommaSeparatedList([&]() -> ParseResult {
     UnresolvedOperand result;
     if (parseSSAUse(result))
       return failure();
     results.push_back(result);
     return success();
   });
 }
 
 /// Parse a SSA operand for an operation.
 ///
 ///   ssa-use ::= ssa-id
 ///
 ParseResult OperationParser::parseSSAUse(UnresolvedOperand &result,
                                          bool allowResultNumber) {
+  if (getToken().isCodeCompletion())
+    return codeCompleteSSAUse();
+
   result.name = getTokenSpelling();
   result.number = 0;
   result.location = getToken().getLoc();
   if (parseToken(Token::percent_identifier, "expected SSA operand"))
     return failure();
 
   // If we have an attribute ID, it is a result number.
   if (getToken().is(Token::hash_identifier)) {
ParseResult OperationParser::parseOperation() {
   }
 
   Operation *op;
   Token nameTok = getToken();
   if (nameTok.is(Token::bare_identifier) || nameTok.isKeyword())
     op = parseCustomOperation(resultIDs);
   else if (nameTok.is(Token::string))
     op = parseGenericOperation();
+  else if (nameTok.isCodeCompletionFor(Token::string))
+    return codeCompleteStringDialectOrOperationName(nameTok.getStringValue());
+  else if (nameTok.isCodeCompletion())
+    return codeCompleteDialectOrElidedOpName(loc);
   else
     return emitWrongTokenError("expected operation name in quotes");
 
   // If parsing of the basic operation failed, then this whole thing fails.
   if (!op)
     return failure();
 
   // If the operation had a name, register it.
ParseResult OperationParser::parseOperation() {
   return success();
 }
 
 /// Parse a single operation successor.
 ///
 ///   successor ::= block-id
 ///
 ParseResult OperationParser::parseSuccessor(Block *&dest) {
+  if (getToken().isCodeCompletion())
+    return codeCompleteBlock();
+
   // Verify branch is identifier and get the matching block.
   if (!getToken().is(Token::caret_identifier))
     return emitWrongTokenError("expected block name");
   dest = getBlockNamed(getTokenSpelling(), getToken().getLoc());
   consumeToken();
   return success();
 }
 
ParseResult parseOperand(UnresolvedOperand &result,
     result = {useInfo.location, useInfo.name, useInfo.number};
     return success();
   }
 
   /// Parse a single operand if present.
   OptionalParseResult
   parseOptionalOperand(UnresolvedOperand &result,
                        bool allowResultNumber = true) override {
-    if (parser.getToken().is(Token::percent_identifier))
+    if (parser.getToken().isOrIsCodeCompletionFor(Token::percent_identifier))
       return parseOperand(result, allowResultNumber);
     return llvm::None;
   }
 
   /// Parse zero or more SSA comma-separated operand references with a specified
   /// surrounding delimiter, and an optional required operand count.
   ParseResult parseOperandList(SmallVectorImpl<UnresolvedOperand> &result,
                                Delimiter delimiter = Delimiter::None,
                                bool allowResultNumber = true,
                                int requiredOperandCount = -1) override {
     // The no-delimiter case has some special handling for better diagnostics.
     if (delimiter == Delimiter::None) {
       // parseCommaSeparatedList doesn't handle the missing case for "none",
       // so we handle it custom here.
-      if (parser.getToken().isNot(Token::percent_identifier)) {
+      Token tok = parser.getToken();
+      if (!tok.isOrIsCodeCompletionFor(Token::percent_identifier)) {
         // If we didn't require any operands or required exactly zero (weird)
         // then this is success.
         if (requiredOperandCount == -1 || requiredOperandCount == 0)
           return success();
 
         // Otherwise, try to produce a nice error message.
-        if (parser.getToken().isAny(Token::l_paren, Token::l_square))
+        if (tok.isAny(Token::l_paren, Token::l_square))
           return parser.emitError("unexpected delimiter");
         return parser.emitWrongTokenError("expected operand");
       }
     }
 
     auto parseOneOperand = [&]() -> ParseResult {
       return parseOperand(result.emplace_back(), allowResultNumber);
     };
public:
 
   /// Parse a single operation successor.
   ParseResult parseSuccessor(Block *&dest) override {
     return parser.parseSuccessor(dest);
   }
 
   /// Parse an optional operation successor and its operand list.
   OptionalParseResult parseOptionalSuccessor(Block *&dest) override {
-    if (parser.getToken().isNot(Token::caret_identifier))
+    if (!parser.getToken().isOrIsCodeCompletionFor(Token::caret_identifier))
       return llvm::None;
     return parseSuccessor(dest);
   }
 
   /// Parse a single operation successor and its operand list.
   ParseResult
   parseSuccessorAndUseList(Block *&dest,
                            SmallVectorImpl<Value> &operands) override {
private:
   StringRef opName;
 
   /// The backing operation parser.
   OperationParser &parser;
 };
 } // namespace
 
 FailureOr<OperationName> OperationParser::parseCustomOperationName() {
-  std::string opName = getTokenSpelling().str();
+  Token nameTok = getToken();
+  StringRef opName = nameTok.getSpelling();
   if (opName.empty())
     return (emitError("empty operation name is invalid"), failure());
   consumeToken();
 
   // Check to see if this operation name is already registered.
   Optional<RegisteredOperationName> opInfo =
       RegisteredOperationName::lookup(opName, getContext());
   if (opInfo)
     return *opInfo;
 
   // If the operation doesn't have a dialect prefix try using the default
   // dialect.
-  auto opNameSplit = StringRef(opName).split('.');
+  auto opNameSplit = opName.split('.');
   StringRef dialectName = opNameSplit.first;
+  std::string opNameStorage;
   if (opNameSplit.second.empty()) {
+    // If the name didn't have a prefix, check for a code completion request.
+    if (getToken().isCodeCompletion() && opName.back() == '.')
+      return codeCompleteOperationName(dialectName);
+
     dialectName = getState().defaultDialectStack.back();
-    opName = (dialectName + "." + opName).str();
+    opNameStorage = (dialectName + "." + opName).str();
+    opName = opNameStorage;
   }
 
   // Try to load the dialect before returning the operation name to make sure
   // the operation has a chance to be registered.
   getContext()->getOrLoadDialect(dialectName);
   return OperationName(opName, getContext());
 }
 
return parseSSADefOrUseAndType(
             state.asmState->addDefinition(arg, useInfo.location);
 
           return addDefinition(useInfo, arg);
         });
   });
 }
 
 //===----------------------------------------------------------------------===//
+// Code Completion
+//===----------------------------------------------------------------------===//
+
+ParseResult OperationParser::codeCompleteSSAUse() {
+  std::string detailData;
+  llvm::raw_string_ostream detailOS(detailData);
+  for (IsolatedSSANameScope &scope : isolatedNameScopes) {
+    for (auto &it : scope.values) {
+      if (it.second.empty())
+        continue;
+      Value frontValue = it.second.front().value;
+
+      // If the value isn't a forward reference, we also add the name of the op
+      // to the detail.
+      if (auto result = frontValue.dyn_cast<OpResult>()) {
+        if (!forwardRefPlaceholders.count(result))
+          detailOS << result.getOwner()->getName() << ": ";
+      } else {
+        detailOS << "arg #" << frontValue.cast<BlockArgument>().getArgNumber()
+                 << ": ";
+      }
+
+      // Emit the type of the values to aid with completion selection.
+      detailOS << frontValue.getType();
+
+      // FIXME: We should define a policy for packed values, e.g. with a limit
+      // on the detail size, but it isn't clear what would be useful right now.
+      // For now we just only emit the first type.
+      if (it.second.size() > 1)
+        detailOS << ", ...";
+
+      state.codeCompleteContext->appendSSAValueCompletion(
+          it.getKey(), std::move(detailOS.str()));
+    }
+  }
+
+  return failure();
+}
+
+ParseResult OperationParser::codeCompleteBlock() {
+  // Don't provide completions if the token isn't empty, e.g. this avoids
+  // weirdness when we encounter a `.` within the identifier.
+  StringRef spelling = getTokenSpelling();
+  if (!(spelling.empty() || spelling == "^"))
+    return failure();
+
+  for (const auto &it : blocksByName.back())
+    state.codeCompleteContext->appendBlockCompletion(it.getFirst());
+  return failure();
+}
+
+//===----------------------------------------------------------------------===//
 // Top-level entity parsing.
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This parser handles entities that are only valid at the top level of the
 /// file.
 class TopLevelOperationParser : public Parser {
 public:
case Token::file_metadata_begin:
         return failure();
       break;
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 
-LogicalResult mlir::parseSourceFile(const llvm::SourceMgr &sourceMgr,
-                                    Block *block, const ParserConfig &config,
-                                    LocationAttr *sourceFileLoc,
-                                    AsmParserState *asmState) {
+LogicalResult
+mlir::parseSourceFile(const llvm::SourceMgr &sourceMgr, Block *block,
+                      const ParserConfig &config, LocationAttr *sourceFileLoc,
+                      AsmParserState *asmState,
+                      AsmParserCodeCompleteContext *codeCompleteContext) {
   const auto *sourceBuf = sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
 
   Location parserLoc =
       FileLineColLoc::get(config.getContext(), sourceBuf->getBufferIdentifier(),
                           /*line=*/0, /*column=*/0);
   if (sourceFileLoc)
     *sourceFileLoc = parserLoc;
 
   SymbolState aliasState;
-  ParserState state(sourceMgr, config, aliasState, asmState);
+  ParserState state(sourceMgr, config, aliasState, asmState,
+                    codeCompleteContext);
   return TopLevelOperationParser(state).parse(block, parserLoc);
 }
 
 LogicalResult mlir::parseSourceFile(llvm::StringRef filename, Block *block,
                                     const ParserConfig &config,
                                     LocationAttr *sourceFileLoc) {
   llvm::SourceMgr sourceMgr;
   return parseSourceFile(filename, sourceMgr, block, config, sourceFileLoc);

mlir/lib/Parser/ParserState.h

 //===----------------------------------------------------------------------===//
 // ParserState
 //===----------------------------------------------------------------------===//
 
 /// This class refers to all of the state maintained globally by the parser,
 /// such as the current lexer position etc.
 struct ParserState {
   ParserState(const llvm::SourceMgr &sourceMgr, const ParserConfig &config,
-              SymbolState &symbols, AsmParserState *asmState)
-      : config(config), lex(sourceMgr, config.getContext()),
-        curToken(lex.lexToken()), symbols(symbols), asmState(asmState) {}
+              SymbolState &symbols, AsmParserState *asmState,
+              AsmParserCodeCompleteContext *codeCompleteContext)
+      : config(config),
+        lex(sourceMgr, config.getContext(), codeCompleteContext),
+        curToken(lex.lexToken()), symbols(symbols), asmState(asmState),
+        codeCompleteContext(codeCompleteContext) {}
   ParserState(const ParserState &) = delete;
   void operator=(const ParserState &) = delete;
 
   /// The configuration used to setup the parser.
   const ParserConfig &config;
 
   /// The lexer for the source file we're parsing.
   Lexer lex;
 
   /// This is the next token that hasn't been consumed yet.
   Token curToken;
 
   /// The current state for symbol parsing.
   SymbolState &symbols;
 
   /// An optional pointer to a struct containing high level parser state to be
   /// populated during parsing.
   AsmParserState *asmState;
 
+  /// An optional code completion context.
+  AsmParserCodeCompleteContext *codeCompleteContext;
+
   // Contains the stack of default dialect to use when parsing regions.
   // A new dialect get pushed to the stack before parsing regions nested
   // under an operation implementing `OpAsmOpInterface`, and
   // popped when done. At the top-level we start with "builtin" as the
   // default, so that the top-level `module` operation parses as-is.
   SmallVector<StringRef> defaultDialectStack{"builtin"};
 };
 
 } // namespace detail
 } // namespace mlir
 
 #endif // MLIR_LIB_PARSER_PARSERSTATE_H

mlir/lib/Parser/Token.h

#include "TokenKinds.def"
   template <typename... T>
   bool isNot(Kind k1, Kind k2, T... others) const {
     return !isAny(k1, k2, others...);
   }
 
   /// Return true if this is one of the keyword token kinds (e.g. kw_if).
   bool isKeyword() const;
 
+  /// Returns true if the current token represents a code completion.
+  bool isCodeCompletion() const { return is(code_complete); }
+
+  /// Returns true if the current token represents a code completion for the
+  /// "normal" token type.
+  bool isCodeCompletionFor(Kind kind) const;
+
+  /// Returns true if the current token is the given type, or represents a code
+  /// completion for that type.
+  bool isOrIsCodeCompletionFor(Kind kind) const {
+    return is(kind) || isCodeCompletionFor(kind);
+  }
+
   // Helpers to decode specific sorts of tokens.
 
   /// For an integer token, return its value as an unsigned.  If it doesn't fit,
   /// return None.
   Optional<unsigned> getUnsignedIntegerValue() const;
 
   /// For an integer token, return its value as an uint64_t.  If it doesn't fit,
   /// return None.

mlir/lib/Parser/Token.cpp

Optional<bool> Token::getIntTypeSignedness() const {
   assert(spelling[0] == 'u');
   return false;
 }
 
 /// Given a token containing a string literal, return its value, including
 /// removing the quote characters and unescaping the contents of the string. The
 /// lexer has already verified that this token is valid.
 std::string Token::getStringValue() const {
-  assert(getKind() == string ||
+  assert(getKind() == string || getKind() == code_complete ||
          (getKind() == at_identifier && getSpelling()[1] == '"'));
   // Start by dropping the quotes.
-  StringRef bytes = getSpelling().drop_front().drop_back();
+  StringRef bytes = getSpelling().drop_front();
+  if (getKind() != Token::code_complete) {
+    bytes = bytes.drop_back();
     if (getKind() == at_identifier)
       bytes = bytes.drop_front();
+  }
 
   std::string result;
   result.reserve(bytes.size());
   for (unsigned i = 0, e = bytes.size(); i != e;) {
     auto c = bytes[i++];
     if (c != '\\') {
       result.push_back(c);
       continue;
bool Token::isKeyword() const {
   default:
     return false;
 #define TOK_KEYWORD(SPELLING)                                                  \
   case kw_##SPELLING:                                                          \
     return true;
 #include "TokenKinds.def"
   }
 }
+
+bool Token::isCodeCompletionFor(Kind kind) const {
+  if (!isCodeCompletion() || spelling.empty())
+    return false;
+  switch (kind) {
+  case Kind::string:
+    return spelling[0] == '"';
+  case Kind::hash_identifier:
+    return spelling[0] == '#';
+  case Kind::percent_identifier:
+    return spelling[0] == '%';
+  case Kind::caret_identifier:
+    return spelling[0] == '^';
+  case Kind::exclamation_identifier:
+    return spelling[0] == '!';
+  default:
+    return false;
+  }
+}

mlir/lib/Parser/TokenKinds.def

 #endif
 #ifndef TOK_KEYWORD
 #define TOK_KEYWORD(SPELLING)
 #endif
 
 // Markers
 TOK_MARKER(eof)
 TOK_MARKER(error)
+TOK_MARKER(code_complete)
 
 // Identifiers.
 TOK_IDENTIFIER(bare_identifier)        // foo
 TOK_IDENTIFIER(at_identifier)          // @foo
 TOK_IDENTIFIER(hash_identifier)        // #foo
 TOK_IDENTIFIER(percent_identifier)     // %foo
 TOK_IDENTIFIER(caret_identifier)       // ^foo
 TOK_IDENTIFIER(exclamation_identifier) // !foo

mlir/lib/Tools/lsp-server-support/Protocol.h

enum class InsertTextFormat {
   /// typing in one will update others too.
   ///
   /// See also:
   /// https//github.com/Microsoft/vscode/blob/master/src/vs/editor/contrib/snippet/common/snippet.md
   Snippet = 2,
 };
 
 struct CompletionItem {
+  CompletionItem() = default;
+  CompletionItem(StringRef label, CompletionItemKind kind)
+      : label(label.str()), kind(kind),
+        insertTextFormat(InsertTextFormat::PlainText) {}
+
   /// The label of this completion item. By default also the text that is
   /// inserted when selecting this completion.
   std::string label;
 
   /// The kind of this completion item. Based of the kind an icon is chosen by
   /// the editor.
   CompletionItemKind kind = CompletionItemKind::Missing;
 

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

struct LSPServer::Impl {
 
   //===--------------------------------------------------------------------===//
   // Document Symbols
 
   void onDocumentSymbol(const DocumentSymbolParams &params,
                         Callback<std::vector<DocumentSymbol>> reply);
 
   //===--------------------------------------------------------------------===//
+  // Code Completion
+
+  void onCompletion(const CompletionParams &params,
+                    Callback<CompletionList> reply);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   MLIRServer &server;
   JSONTransport &transport;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.
void LSPServer::Impl::onInitialize(const InitializeParams &params,
   // Send a response with the capabilities of this server.
   llvm::json::Object serverCaps{
       {"textDocumentSync",
        llvm::json::Object{
            {"openClose", true},
            {"change", (int)TextDocumentSyncKind::Full},
            {"save", true},
        }},
+      {"completionProvider",
+       llvm::json::Object{
+           {"allCommitCharacters",
+            {"\t", "(", ")", "[", "]", "<", ">", ";", ",", "+", "-", "/", "*",
+             "&", "?", ".", "=", "|"}},
+           {"resolveProvider", false},
+           {"triggerCharacters",
+            {".", "%", "^", "!", "#", "(", ",", "<", ":", "[", " ", "\"", "/"}},
+       }},
       {"definitionProvider", true},
       {"referencesProvider", true},
       {"hoverProvider", true},
 
       // For now we only support documenting symbols when the client supports
       // hierarchical symbols.
       {"documentSymbolProvider",
        params.capabilities.hierarchicalDocumentSymbol},
void LSPServer::Impl::onDocumentSymbol(
     const DocumentSymbolParams &params,
     Callback<std::vector<DocumentSymbol>> reply) {
   std::vector<DocumentSymbol> symbols;
   server.findDocumentSymbols(params.textDocument.uri, symbols);
   reply(std::move(symbols));
 }
 
 //===----------------------------------------------------------------------===//
+// Code Completion
+
+void LSPServer::Impl::onCompletion(const CompletionParams &params,
+                                   Callback<CompletionList> reply) {
+  reply(server.getCodeCompletion(params.textDocument.uri, params.position));
+}
+
+//===----------------------------------------------------------------------===//
 // LSPServer
 //===----------------------------------------------------------------------===//
 
 LSPServer::LSPServer(MLIRServer &server, JSONTransport &transport)
     : impl(std::make_unique<Impl>(server, transport)) {}
 LSPServer::~LSPServer() = default;
 
 LogicalResult LSPServer::run() {
LogicalResult LSPServer::run() {
 
   // Hover
   messageHandler.method("textDocument/hover", impl.get(), &Impl::onHover);
 
   // Document Symbols
   messageHandler.method("textDocument/documentSymbol", impl.get(),
                         &Impl::onDocumentSymbol);
 
+  // Code Completion
+  messageHandler.method("textDocument/completion", impl.get(),
+                        &Impl::onCompletion);
+
   // Diagnostics
   impl->publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
           "textDocument/publishDiagnostics");
 
   // Run the main loop of the transport.
   LogicalResult result = success();
   if (llvm::Error error = impl->transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     result = failure();
   } else {
     result = success(impl->shutdownRequestReceived);
   }
   return result;
 }

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

 
 #include "mlir/Support/LLVM.h"
 #include <memory>
 
 namespace mlir {
 class DialectRegistry;
 
 namespace lsp {
+struct CompletionList;
 struct Diagnostic;
 struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
 class URIForFile;
 
 /// This class implements all of the MLIR related functionality necessary for a
public:
   /// Find a hover description for the given hover position, or None if one
   /// couldn't be found.
   Optional<Hover> findHover(const URIForFile &uri, const Position &hoverPos);
 
   /// Find all of the document symbols within the given file.
   void findDocumentSymbols(const URIForFile &uri,
                            std::vector<DocumentSymbol> &symbols);
 
+  /// Get the code completion list for the position within the given file.
+  CompletionList getCodeCompletion(const URIForFile &uri,
+                                   const Position &completePos);
+
 private:
   struct Impl;
 
   std::unique_ptr<Impl> impl;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_

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

 //===- MLIRServer.cpp - MLIR Generic Language Server ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "MLIRServer.h"
 #include "../lsp-server-support/Logging.h"
 #include "../lsp-server-support/Protocol.h"
 #include "../lsp-server-support/SourceMgrUtils.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Parser/AsmParserState.h"
+#include "mlir/Parser/CodeComplete.h"
 #include "mlir/Parser/Parser.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
 /// Returns a language server location from the given MLIR file location.
 static Optional<lsp::Location> getLocationFromLoc(FileLineColLoc loc) {
   llvm::Expected<lsp::URIForFile> sourceURI =
struct MLIRDocument {
   // Document Symbols
   //===--------------------------------------------------------------------===//
 
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
   void findDocumentSymbols(Operation *op,
                            std::vector<lsp::DocumentSymbol> &symbols);
 
   //===--------------------------------------------------------------------===//
+  // Code Completion
+  //===--------------------------------------------------------------------===//
+
+  lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
+                                        const lsp::Position &completePos,
+                                        const DialectRegistry &registry);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// The high level parser state used to find definitions and references within
   /// the source file.
   AsmParserState asmState;
 
   /// The container for the IR parsed from the input file.
void MLIRDocument::findDocumentSymbols(
   if (!op->getNumRegions())
     return;
   for (Region &region : op->getRegions())
     for (Operation &childOp : region.getOps())
       findDocumentSymbols(&childOp, *childSymbols);
 }
 
 //===----------------------------------------------------------------------===//
+// MLIRDocument: Code Completion
+//===----------------------------------------------------------------------===//
+
+namespace {
+class LSPCodeCompleteContext : public AsmParserCodeCompleteContext {
+public:
+  LSPCodeCompleteContext(SMLoc completeLoc, lsp::CompletionList &completionList,
+                         MLIRContext *ctx)
+      : AsmParserCodeCompleteContext(completeLoc),
+        completionList(completionList), ctx(ctx) {}
+
+  /// Signal code completion for a dialect name.
+  void completeDialectName() final {
+    for (StringRef dialect : ctx->getAvailableDialects()) {
+      lsp::CompletionItem item(dialect, lsp::CompletionItemKind::Module);
+      item.sortText = "2";
+      item.detail = "dialect";
+      completionList.items.emplace_back(item);
+    }
+  }
+
+  /// Signal code completion for an operation name within the given dialect.
+  void completeOperationName(StringRef dialectName) final {
+    Dialect *dialect = ctx->getOrLoadDialect(dialectName);
+    if (!dialect)
+      return;
+
+    for (const auto &op : ctx->getRegisteredOperations()) {
+      if (&op.getDialect() != dialect)
+        continue;
+
+      lsp::CompletionItem item(
+          op.getStringRef().drop_front(dialectName.size() + 1),
+          lsp::CompletionItemKind::Field);
+      item.sortText = "1";
+      item.detail = "operation";
+      completionList.items.emplace_back(item);
+    }
+  }
+
+  /// Append the given SSA value as a code completion result for SSA value
+  /// completions.
+  void appendSSAValueCompletion(StringRef name, std::string typeData) final {
+    // Check if we need to insert the `%` or not.
+    bool stripPrefix = getCodeCompleteLoc().getPointer()[-1] == '%';
+
+    lsp::CompletionItem item(name, lsp::CompletionItemKind::Variable);
+    if (stripPrefix)
+      item.insertText = name.drop_front(1).str();
+    item.detail = std::move(typeData);
+    completionList.items.emplace_back(item);
+  }
+
+  /// Append the given block as a code completion result for block name
+  /// completions.
+  void appendBlockCompletion(StringRef name) final {
+    // Check if we need to insert the `^` or not.
+    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);
+  }
+
+private:
+  lsp::CompletionList &completionList;
+  MLIRContext *ctx;
+};
+} // namespace
+
+lsp::CompletionList
+MLIRDocument::getCodeCompletion(const lsp::URIForFile &uri,
+                                const lsp::Position &completePos,
+                                const DialectRegistry &registry) {
+  SMLoc posLoc = completePos.getAsSMLoc(sourceMgr);
+  if (!posLoc.isValid())
+    return lsp::CompletionList();
+
+  // To perform code completion, we run another parse of the module with the
+  // code completion context provided.
+  MLIRContext tmpContext(registry, MLIRContext::Threading::DISABLED);
+  tmpContext.allowUnregisteredDialects();
+  lsp::CompletionList completionList;
+  LSPCodeCompleteContext lspCompleteContext(posLoc, completionList,
+                                            &tmpContext);
+
+  Block tmpIR;
+  AsmParserState tmpState;
+  (void)parseSourceFile(sourceMgr, &tmpIR, &tmpContext, nullptr, &tmpState,
+                        &lspCompleteContext);
+  return completionList;
+}
+
+//===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an MLIR text file.
 struct MLIRTextFileChunk {
   MLIRTextFileChunk(MLIRContext &context, uint64_t lineOffset,
                     const lsp::URIForFile &uri, StringRef contents,
public:
 
   void getLocationsOf(const lsp::URIForFile &uri, lsp::Position defPos,
                       std::vector<lsp::Location> &locations);
   void findReferencesOf(const lsp::URIForFile &uri, lsp::Position pos,
                         std::vector<lsp::Location> &references);
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  lsp::Position hoverPos);
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
+  lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
+                                        lsp::Position completePos);
 
 private:
   /// Find the MLIR document that contains the given position, and update the
   /// position to be anchored at the start of the found chunk instead of the
   /// beginning of the file.
   MLIRTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The context used to hold the state contained by the parsed document.
if (i != 0) {
       }
     }
 
     // Push the symbol for this chunk.
     symbols.emplace_back(std::move(symbol));
   }
 }
 
+lsp::CompletionList MLIRTextFile::getCodeCompletion(const lsp::URIForFile &uri,
+                                                    lsp::Position completePos) {
+  MLIRTextFileChunk &chunk = getChunkFor(completePos);
+  lsp::CompletionList completionList = chunk.document.getCodeCompletion(
+      uri, completePos, context.getDialectRegistry());
+
+  // Adjust any completion locations.
+  for (lsp::CompletionItem &item : completionList.items) {
+    if (item.textEdit)
+      chunk.adjustLocForChunkOffset(item.textEdit->range);
+    for (lsp::TextEdit &edit : item.additionalTextEdits)
+      chunk.adjustLocForChunkOffset(edit.range);
+  }
+  return completionList;
+}
+
 MLIRTextFileChunk &MLIRTextFile::getChunkFor(lsp::Position &pos) {
   if (chunks.size() == 1)
     return *chunks.front();
 
   // Search for the first chunk with a greater line offset, the previous chunk
   // is the one that contains `pos`.
   auto it = llvm::upper_bound(
       chunks, pos, [](const lsp::Position &pos, const auto &chunk) {

 }
 
 void lsp::MLIRServer::findDocumentSymbols(
     const URIForFile &uri, std::vector<DocumentSymbol> &symbols) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     fileIt->second->findDocumentSymbols(symbols);
 }
+
+lsp::CompletionList
+lsp::MLIRServer::getCodeCompletion(const URIForFile &uri,
+                                   const Position &completePos) {
+  auto fileIt = impl->files.find(uri.file());
+  if (fileIt != impl->files.end())
+    return fileIt->second->getCodeCompletion(uri, completePos);
+  return CompletionList();
+}

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

+// RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.mlir",
+  "languageId":"mlir",
+  "version":1,
+  "text":"func.func private @foo(%arg: i32) -> i32 {\n%cast = \"builtin.unrealized_conversion_cast\"() : () -> (i32)\nreturn %"
+}}}
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":0,"character":0}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK:           {
+// CHECK:             "detail": "dialect",
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 9,
+// CHECK:             "label": "builtin",
+// CHECK:             "sortText": "2"
+// CHECK:           },
+// CHECK:           {
+// CHECK:             "detail": "operation",
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "module",
+// CHECK:             "sortText": "1"
+// CHECK:           }
+// CHECK:         ]
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":1,"character":9}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK:           {
+// CHECK:             "detail": "dialect",
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 9,
+// CHECK:             "label": "builtin",
+// CHECK:             "sortText": "2"
+// CHECK:           },
+// CHECK-NOT:       "detail": "operation",
+// CHECK:         ]
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":1,"character":17}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK-NOT:       "detail": "dialect",
+// CHECK:           {
+// CHECK:             "detail": "operation",
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "module",
+// CHECK:             "sortText": "1"
+// CHECK:           }
+// CHECK:         ]
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":2,"character":8}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "detail": "builtin.unrealized_conversion_cast: i32",
+// CHECK-NEXT:        "insertText": "cast",
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 6,
+// CHECK-NEXT:        "label": "%cast"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "detail": "arg #0: i32",
+// 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":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-lsp-server/initialize-params.test

 // RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
 // Test initialize request parameters with rootUri
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{},"trace":"off"}}
 //      CHECK:  "id": 0,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "capabilities": {
+// CHECK-NEXT:      "completionProvider": {
+// CHECK-NEXT:        "allCommitCharacters": [
+// CHECK:             ],
+// CHECK-NEXT:        "resolveProvider": false,
+// CHECK-NEXT:        "triggerCharacters": [ 
+// CHECK:             ]
+// CHECK-NEXT:      },
 // CHECK-NEXT:      "definitionProvider": true,
 // CHECK-NEXT:      "documentSymbolProvider": false,
 // CHECK-NEXT:      "hoverProvider": true,
 // CHECK-NEXT:      "referencesProvider": true,
 // CHECK-NEXT:      "textDocumentSync": {
 // CHECK-NEXT:        "change": 1,
 // CHECK-NEXT:        "openClose": true,
 // CHECK-NEXT:        "save": true