[mlir:LSP] Add support for code completing attributes and types

[rriddle]

Summary

This required changing a bit of how attributes/types are parsed. A new
KeywordSwitch class was added to AsmParser that provides a StringSwitch
like API for parsing keywords with a set of potential matches. It intends to
both provide a cleaner API, and enable injection for code completion. This
required changing the API of generated(Attr|Type)Parser to handle the
parsing of the keyword, instead of having the user do it. Most upstream
dialects use the autogenerated handling and didn't require a direct update.

Depends on D129184

[jpienaar]

Thanks!

mlir/include/mlir/IR/OpImplementation.h

679

I would have made this a separate commit as this is a nice feature not LSP related.

mlir/lib/Parser/DialectSymbolParser.cpp

55

Should this be part of Parser's state instead?

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

705

Could these be queried from registered attributes instead? (could be in future too and TODO )

[rriddle]

mlir/lib/Parser/DialectSymbolParser.cpp

55

For now no, given that the isCodeCompletion is only local/used when parsing the immediate attribute/type. I think in the future as we build more complex completions, it might be useful to have it there.

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

705

Thinking about it, let me add completions for dialects with attributes and types. We can't query the ones here because they are not encoded in the registry, given that builtin attributes/types have special syntax/parsing.

flang/lib/Optimizer/Dialect/FIRType.cpp

 
 } // namespace
 
 // Implementation of the thin interface from dialect to type parser
 
 mlir::Type fir::parseFirType(FIROpsDialect *dialect,
                              mlir::DialectAsmParser &parser) {
   mlir::StringRef typeTag;
-  if (parser.parseKeyword(&typeTag))
-    return {};
   mlir::Type genType;
-  auto parseResult = generatedTypeParser(parser, typeTag, genType);
+  auto parseResult = generatedTypeParser(parser, &typeTag, genType);
   if (parseResult.hasValue())
     return genType;
   parser.emitError(parser.getNameLoc(), "unknown fir type: ") << typeTag;
   return {};
 }
 
 namespace fir {
 namespace detail {

mlir/docs/AttributesAndTypes.md

 and Types with no parameters need not use either of these fields, in which case
 the syntax for the Attribute or Type is simply the mnemonic.
 
 For each dialect, two "dispatch" functions will be created: one for parsing and
 one for printing. These static functions placed alongside the class definitions
 and have the following function signatures:
 
 ```c++
-static ParseResult generatedAttributeParser(DialectAsmParser& parser, StringRef mnemonic, Type attrType, Attribute &result);
+static ParseResult generatedAttributeParser(DialectAsmParser& parser, StringRef *mnemonic, Type attrType, Attribute &result);
 static LogicalResult generatedAttributePrinter(Attribute attr, DialectAsmPrinter& printer);
 
-static ParseResult generatedTypeParser(DialectAsmParser& parser, StringRef mnemonic, Type &result);
+static ParseResult generatedTypeParser(DialectAsmParser& parser, StringRef *mnemonic, Type &result);
 static LogicalResult generatedTypePrinter(Type type, DialectAsmPrinter& printer);
 ```
 
 The above functions should be added to the respective in your
 `Dialect::printType` and `Dialect::parseType` methods, or consider using the
 `useDefaultAttributePrinterParser` and `useDefaultTypePrinterParser` ODS Dialect
 options if all attributes or types define a mnemonic.
 

mlir/include/mlir/IR/OpImplementation.h

ParseResult parseString(std::string *string) {
     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) {
-    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();
-  }
-
-  /// Parse the given keyword if present.
-  virtual ParseResult parseOptionalKeyword(StringRef keyword) = 0;
-
-  /// Parse a keyword, if present, into 'keyword'.
-  virtual ParseResult parseOptionalKeyword(StringRef *keyword) = 0;
-
-  /// Parse a keyword, if present, and if one of the 'allowedValues',
-  /// into 'keyword'
-  virtual ParseResult
-  parseOptionalKeyword(StringRef *keyword,
-                       ArrayRef<StringRef> allowedValues) = 0;
-
-  /// Parse a keyword or a quoted string.
-  ParseResult parseKeywordOrString(std::string *result) {
-    if (failed(parseOptionalKeywordOrString(result)))
-      return emitError(getCurrentLocation())
-             << "expected valid keyword or string";
-    return success();
-  }
-
-  /// Parse an optional keyword or string.
-  virtual ParseResult parseOptionalKeywordOrString(std::string *result) = 0;
-
   /// Parse a `(` token.
   virtual ParseResult parseLParen() = 0;
 
   /// Parse a `(` token if present.
   virtual ParseResult parseOptionalLParen() = 0;
 
   /// Parse a `)` token.
   virtual ParseResult parseRParen() = 0;
public:
   /// Parse a comma separated list of elements that must have at least one entry
   /// in it.
   ParseResult
   parseCommaSeparatedList(function_ref<ParseResult()> parseElementFn) {
     return parseCommaSeparatedList(Delimiter::None, parseElementFn);
   }
 
   //===--------------------------------------------------------------------===//
+  // Keyword Parsing
+  //===--------------------------------------------------------------------===//
+
+  /// This class represents a StringSwitch like class that is useful for parsing
+  /// expected keywords. On construction, it invokes `parseKeyword` and
+  /// processes each of the provided cases statements until a match is hit. The
+  /// provided `ResultT` must be assignable from `failure()`.
+  template <typename ResultT = ParseResult>
+  class KeywordSwitch {
+  public:
+    KeywordSwitch(AsmParser &parser)
+        : parser(parser), loc(parser.getCurrentLocation()) {
+      if (failed(parser.parseKeywordOrCompletion(&keyword)))
+        result = failure();
+    }
+
+    /// Case that uses the provided value when true.
+    KeywordSwitch &Case(StringLiteral str, ResultT value) {
+      return Case(str, [&](StringRef, SMLoc) { return std::move(value); });
+    }
+    KeywordSwitch &Default(ResultT value) {
+      return Default([&](StringRef, SMLoc) { return std::move(value); });
+    }
+    /// Case that invokes the provided functor when true. The parameters passed
+    /// to the functor are the keyword, and the location of the keyword (in case
+    /// any errors need to be emitted).
+    template <typename FnT>
+    std::enable_if_t<!std::is_convertible<FnT, ResultT>::value, KeywordSwitch &>
+    Case(StringLiteral str, FnT &&fn) {
+      if (result)
+        return *this;
+
+      // If the word was empty, record this as a completion.
+      if (keyword.empty())
+        parser.codeCompleteExpectedTokens(str);
+      else if (keyword == str)
+        result.emplace(std::move(fn(keyword, loc)));
+      return *this;
+    }
+    template <typename FnT>
+    std::enable_if_t<!std::is_convertible<FnT, ResultT>::value, KeywordSwitch &>
+    Default(FnT &&fn) {
+      if (!result)
+        result.emplace(fn(keyword, loc));
+      return *this;
+    }
+
+    /// Returns true if this switch has a value yet.
+    bool hasValue() const { return result.hasValue(); }
+
+    /// Return the result of the switch.
+    LLVM_NODISCARD operator ResultT() {
+      if (!result)
+        return parser.emitError(loc, "unexpected keyword: ") << keyword;
+      return std::move(*result);
+    }
+
+  private:
+    /// The parser used to construct this switch.
+    AsmParser &parser;
+
+    /// The location of the keyword, used to emit errors as necessary.
+    SMLoc loc;
+
+    /// The parsed keyword itself.
+    StringRef keyword;
+
+    /// The result of the switch statement or none if currently unknown.
+    Optional<ResultT> result;
+  };
+
+  /// Parse a given keyword.
+  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();
+  }
+
+  /// Parse the given keyword if present.
+  virtual ParseResult parseOptionalKeyword(StringRef keyword) = 0;
+
+  /// Parse a keyword, if present, into 'keyword'.
+  virtual ParseResult parseOptionalKeyword(StringRef *keyword) = 0;
+
+  /// Parse a keyword, if present, and if one of the 'allowedValues',
+  /// into 'keyword'
+  virtual ParseResult
+  parseOptionalKeyword(StringRef *keyword,
+                       ArrayRef<StringRef> allowedValues) = 0;
+
+  /// Parse a keyword or a quoted string.
+  ParseResult parseKeywordOrString(std::string *result) {
+    if (failed(parseOptionalKeywordOrString(result)))
+      return emitError(getCurrentLocation())
+             << "expected valid keyword or string";
+    return success();
+  }
+
+  /// Parse an optional keyword or string.
+  virtual ParseResult parseOptionalKeywordOrString(std::string *result) = 0;
+
+  //===--------------------------------------------------------------------===//
   // Attribute/Type Parsing
   //===--------------------------------------------------------------------===//
 
   /// Invoke the `getChecked` method of the given Attribute or Type class, using
   /// the provided location to emit errors in the case of failure. Note that
   /// unlike `OpBuilder::getType`, this method does not implicitly insert a
   /// context parameter.
   template <typename T, typename... ParamsT>
public:
   virtual ParseResult parseXInDimensionList() = 0;
 
 protected:
   /// Parse a handle to a resource within the assembly format for the given
   /// dialect.
   virtual FailureOr<AsmDialectResourceHandle>
   parseResourceHandle(Dialect *dialect) = 0;
 
+  //===--------------------------------------------------------------------===//
+  // Code Completion
+  //===--------------------------------------------------------------------===//
+
+  /// Parse a keyword, or an empty string if the current location signals a code
+  /// completion.
+  virtual ParseResult parseKeywordOrCompletion(StringRef *keyword) = 0;
+
+  /// Signal the code completion of a set of expected tokens.
+  virtual void codeCompleteExpectedTokens(ArrayRef<StringRef> tokens) = 0;
+
 private:
   AsmParser(const AsmParser &) = delete;
   void operator=(const AsmParser &) = delete;
 };
 
 //===----------------------------------------------------------------------===//
 // OpAsmParser
 //===----------------------------------------------------------------------===//

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/ADT/StringMap.h"
 #include "llvm/Support/SourceMgr.h"
 
 namespace mlir {
+class Attribute;
+class Type;
+
 /// 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();
 
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;
 
+  /// Signal a completion for an attribute.
+  virtual void completeAttribute() = 0;
+  virtual void completeDialectAttributeOrAlias(
+      const llvm::StringMap<Attribute> &aliases) = 0;
+
+  /// Signal a completion for a type.
+  virtual void completeType() = 0;
+  virtual void
+  completeDialectTypeOrAlias(const llvm::StringMap<Type> &aliases) = 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/Dialect/PDL/IR/PDLTypes.cpp

void PDLDialect::registerTypes() {
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/PDL/IR/PDLOpsTypes.cpp.inc"
       >();
 }
 
 static Type parsePDLType(AsmParser &parser) {
   StringRef typeTag;
-  if (parser.parseKeyword(&typeTag))
-    return Type();
   {
     Type genType;
-    auto parseResult = generatedTypeParser(parser, typeTag, genType);
+    auto parseResult = generatedTypeParser(parser, &typeTag, genType);
     if (parseResult.hasValue())
       return genType;
   }
 
   // FIXME: This ends up with a double error being emitted if `RangeType` also
   // emits an error. We should rework the `generatedTypeParser` to better
   // support when the keyword is valid but the individual type parser itself
   // emits an error.

mlir/lib/Dialect/SPIRV/IR/SPIRVAttributes.cpp

Attribute SPIRVDialect::parseAttribute(DialectAsmParser &parser,
   // SPIR-V attributes are dictionaries so they do not have type.
   if (type) {
     parser.emitError(parser.getNameLoc(), "unexpected type");
     return {};
   }
 
   // Parse the kind keyword first.
   StringRef attrKind;
-  if (parser.parseKeyword(&attrKind))
-    return {};
-
   Attribute attr;
   OptionalParseResult result =
-      generatedAttributeParser(parser, attrKind, type, attr);
-  if (result.hasValue()) {
-    if (failed(result.getValue()))
-      return {};
+      generatedAttributeParser(parser, &attrKind, type, attr);
+  if (result.hasValue())
     return attr;
-  }
 
   if (attrKind == spirv::TargetEnvAttr::getKindName())
     return parseTargetEnvAttr(parser);
   if (attrKind == spirv::VerCapExtAttr::getKindName())
     return parseVerCapExtAttr(parser);
   if (attrKind == spirv::InterfaceVarABIAttr::getKindName())
     return parseInterfaceVarABIAttr(parser);
 

mlir/lib/Parser/AsmParserImpl.h

if (!parser.getToken().is(Token::string))
       return failure();
 
     if (string)
       *string = parser.getToken().getStringValue();
     parser.consumeToken();
     return success();
   }
 
+  /// Parse a floating point value from the stream.
+  ParseResult parseFloat(double &result) override {
+    bool isNegative = parser.consumeIf(Token::minus);
+    Token curTok = parser.getToken();
+    SMLoc loc = curTok.getLoc();
+
+    // Check for a floating point value.
+    if (curTok.is(Token::floatliteral)) {
+      auto val = curTok.getFloatingPointValue();
+      if (!val)
+        return emitError(loc, "floating point value too large");
+      parser.consumeToken(Token::floatliteral);
+      result = isNegative ? -*val : *val;
+      return success();
+    }
+
+    // Check for a hexadecimal float value.
+    if (curTok.is(Token::integer)) {
+      Optional<APFloat> apResult;
+      if (failed(parser.parseFloatFromIntegerLiteral(
+              apResult, curTok, isNegative, APFloat::IEEEdouble(),
+              /*typeSizeInBits=*/64)))
+        return failure();
+
+      parser.consumeToken(Token::integer);
+      result = apResult->convertToDouble();
+      return success();
+    }
+
+    return emitError(loc, "expected floating point literal");
+  }
+
+  /// Parse an optional integer value from the stream.
+  OptionalParseResult parseOptionalInteger(APInt &result) override {
+    return parser.parseOptionalInteger(result);
+  }
+
+  /// 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 parseCommaSeparatedList(Delimiter delimiter,
+                                      function_ref<ParseResult()> parseElt,
+                                      StringRef contextMessage) override {
+    return parser.parseCommaSeparatedList(delimiter, parseElt, contextMessage);
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Keyword Parsing
+  //===--------------------------------------------------------------------===//
+
   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();
   }
+  using AsmParser::parseKeyword;
 
   /// 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() ||
ParseResult parseOptionalKeywordOrString(std::string *result) override {
     if (succeeded(parseOptionalKeyword(&keyword))) {
       *result = keyword.str();
       return success();
     }
 
     return parseOptionalString(result);
   }
 
-  /// Parse a floating point value from the stream.
-  ParseResult parseFloat(double &result) override {
-    bool isNegative = parser.consumeIf(Token::minus);
-    Token curTok = parser.getToken();
-    SMLoc loc = curTok.getLoc();
-
-    // Check for a floating point value.
-    if (curTok.is(Token::floatliteral)) {
-      auto val = curTok.getFloatingPointValue();
-      if (!val)
-        return emitError(loc, "floating point value too large");
-      parser.consumeToken(Token::floatliteral);
-      result = isNegative ? -*val : *val;
-      return success();
-    }
-
-    // Check for a hexadecimal float value.
-    if (curTok.is(Token::integer)) {
-      Optional<APFloat> apResult;
-      if (failed(parser.parseFloatFromIntegerLiteral(
-              apResult, curTok, isNegative, APFloat::IEEEdouble(),
-              /*typeSizeInBits=*/64)))
-        return failure();
-
-      parser.consumeToken(Token::integer);
-      result = apResult->convertToDouble();
-      return success();
-    }
-
-    return emitError(loc, "expected floating point literal");
-  }
-
-  /// Parse an optional integer value from the stream.
-  OptionalParseResult parseOptionalInteger(APInt &result) override {
-    return parser.parseOptionalInteger(result);
-  }
-
-  /// 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 parseCommaSeparatedList(Delimiter delimiter,
-                                      function_ref<ParseResult()> parseElt,
-                                      StringRef contextMessage) override {
-    return parser.parseCommaSeparatedList(delimiter, parseElt, contextMessage);
-  }
-
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse an arbitrary attribute and return it in result.
   ParseResult parseAttribute(Attribute &result, Type type) override {
     result = parser.parseAttribute(type);
     return success(static_cast<bool>(result));
ParseResult parseDimensionList(SmallVectorImpl<int64_t> &dimensions,
     return parser.parseDimensionListRanked(dimensions, allowDynamic,
                                            withTrailingX);
   }
 
   ParseResult parseXInDimensionList() override {
     return parser.parseXInDimensionList();
   }
 
+  //===--------------------------------------------------------------------===//
+  // Code Completion
+  //===--------------------------------------------------------------------===//
+
+  /// Parse a keyword, or an empty string if the current location signals a code
+  /// completion.
+  ParseResult parseKeywordOrCompletion(StringRef *keyword) override {
+    Token tok = parser.getToken();
+    if (tok.isCodeCompletion() && tok.getSpelling().empty()) {
+      *keyword = "";
+      return success();
+    }
+    return parseKeyword(keyword);
+  }
+
+  /// Signal the code completion of a set of expected tokens.
+  void codeCompleteExpectedTokens(ArrayRef<StringRef> tokens) override {
+    Token tok = parser.getToken();
+    if (tok.isCodeCompletion() && tok.getSpelling().empty())
+      (void)parser.codeCompleteExpectedTokens(tokens);
+  }
+
 protected:
   /// The source location of the dialect symbol.
   SMLoc nameLoc;
 
   /// The main parser.
   Parser &parser;
 
   /// A flag that indicates if any errors were emitted during parsing.
   bool emittedError = false;
 };
 } // namespace detail
 } // namespace mlir
 
 #endif // MLIR_LIB_PARSER_ASMPARSERIMPL_H

mlir/lib/Parser/AttributeParser.cpp

case Token::at_identifier: {
     return symbolRefAttr;
   }
 
   // Parse a 'unit' attribute.
   case Token::kw_unit:
     consumeToken(Token::kw_unit);
     return builder.getUnitAttr();
 
+    // Handle completion of an attribute.
+  case Token::code_complete:
+    if (getToken().isCodeCompletionFor(Token::hash_identifier))
+      return parseExtendedAttr(type);
+    return codeCompleteAttribute();
+
   default:
     // Parse a type attribute. We parse `Optional` here to allow for providing a
     // better error message.
     Type type;
     OptionalParseResult result = parseOptionalType(type);
     if (!result.hasValue())
       return emitWrongTokenError("expected attribute value"), Attribute();
     return failed(*result) ? Attribute() : TypeAttr::get(type);

mlir/lib/Parser/DialectSymbolParser.cpp

public:
   StringRef getFullSymbolSpec() const override { return fullSpec; }
 
 private:
   /// The full symbol specification.
   StringRef fullSpec;
 };
 } // namespace
 
-/// Parse the body of a dialect symbol, which starts and ends with <>'s, and may
-/// be recursive. Return with the 'body' StringRef encompassing the entire
-/// body.
 ///
 ///   pretty-dialect-sym-body ::= '<' pretty-dialect-sym-contents+ '>'
 ///   pretty-dialect-sym-contents ::= pretty-dialect-sym-body
 ///                                  | '(' pretty-dialect-sym-contents+ ')'
 ///                                  | '[' pretty-dialect-sym-contents+ ']'
 ///                                  | '{' pretty-dialect-sym-contents+ '}'
 ///                                  | '[^[<({>\])}\0]+'
 ///
-ParseResult Parser::parseDialectSymbolBody(StringRef &body) {
+ParseResult Parser::parseDialectSymbolBody(StringRef &body,
+                                           bool &isCodeCompletion) {
   // Symbol bodies are a relatively unstructured format that contains a series
   // of properly nested punctuation, with anything else in the middle. Scan
   // ahead to find it and consume it if successful, otherwise emit an error.
   const char *curPtr = getTokenSpelling().data();
 
   // Scan over the nested punctuation, bailing out on error and consuming until
   // we find the end. We know that we're currently looking at the '<', so we can
   // go until we find the matching '>' character.
   assert(*curPtr == '<');
   SmallVector<char, 8> nestedPunctuation;
+  const char *codeCompleteLoc = state.lex.getCodeCompleteLoc();
   do {
+    // Handle code completions, which may appear in the middle of the symbol
+    // body.
+    if (curPtr == codeCompleteLoc) {
+      isCodeCompletion = true;
+      nestedPunctuation.clear();
+      break;
+    }
+
     char c = *curPtr++;
     switch (c) {
     case '\0':
       // This also handles the EOF case.
       if (!nestedPunctuation.empty()) {
         return emitError() << "unbalanced '" << nestedPunctuation.back()
                            << "' character in pretty dialect name";
       }
case ')':
       break;
     case '}':
       if (nestedPunctuation.pop_back_val() != '{')
         return emitError("unbalanced '}' character in pretty dialect name");
       break;
     case '"': {
       // Dispatch to the lexer to lex past strings.
       resetToken(curPtr - 1);
+      curPtr = state.curToken.getEndLoc().getPointer();
+
+      // Handle code completions, which may appear in the middle of the symbol
+      // body.
+      if (state.curToken.isCodeCompletion()) {
+        isCodeCompletion = true;
+        nestedPunctuation.clear();
+        break;
+      }
+
+      // Otherwise, ensure this token was actually a string.
       if (state.curToken.isNot(Token::string))
         return failure();
-      curPtr = state.curToken.getEndLoc().getPointer();
       break;
     }
 
     default:
       continue;
     }
   } while (!nestedPunctuation.empty());
 
   // Ok, we succeeded, remember where we stopped, reset the lexer to know it is
   // consuming all this stuff, and return.
   resetToken(curPtr);
 
   unsigned length = curPtr - body.begin();
   body = StringRef(body.data(), length);
   return success();
 }
 
 /// Parse an extended dialect symbol.
 template <typename Symbol, typename SymbolAliasMap, typename CreateFn>
-static Symbol parseExtendedSymbol(Parser &p, Token::Kind identifierTok,
-                                  SymbolAliasMap &aliases,
+static Symbol parseExtendedSymbol(Parser &p, SymbolAliasMap &aliases,
                                   CreateFn &&createSymbol) {
+  Token tok = p.getToken();
+
+  // Handle code completion of the extended symbol.
+  StringRef identifier = tok.getSpelling().drop_front();
+  if (tok.isCodeCompletion() && identifier.empty())
+    return p.codeCompleteDialectSymbol(aliases);
+
   // Parse the dialect namespace.
-  StringRef identifier = p.getTokenSpelling().drop_front();
   SMLoc loc = p.getToken().getLoc();
-  p.consumeToken(identifierTok);
+  p.consumeToken();
 
   // Check to see if this is a pretty name.
   StringRef dialectName;
   StringRef symbolData;
   std::tie(dialectName, symbolData) = identifier.split('.');
-  bool isPrettyName = !symbolData.empty();
+  bool isPrettyName = !symbolData.empty() || identifier.back() == '.';
 
   // Check to see if the symbol has trailing data, i.e. has an immediately
   // following '<'.
   bool hasTrailingData =
       p.getToken().is(Token::less) &&
       identifier.bytes_end() == p.getTokenSpelling().bytes_begin();
 
   // If there is no '<' token following this, and if the typename contains no
static Symbol parseExtendedSymbol(Parser &p, SymbolAliasMap &aliases,
   }
 
   // If this isn't an alias, we are parsing a dialect-specific symbol. If the
   // name contains a dot, then this is the "pretty" form. If not, it is the
   // verbose form that looks like <...>.
   if (!isPrettyName) {
     // Point the symbol data to the end of the dialect name to start.
     symbolData = StringRef(dialectName.end(), 0);
-    if (p.parseDialectSymbolBody(symbolData))
+
+    // Parse the body of the symbol.
+    bool isCodeCompletion = false;
+    if (p.parseDialectSymbolBody(symbolData, isCodeCompletion))
       return nullptr;
-    symbolData = symbolData.drop_front().drop_back();
+    symbolData = symbolData.drop_front();
+
+    // If the body contained a code completion it won't have the trailing `>`
+    // token, so don't drop it.
+    if (!isCodeCompletion)
+      symbolData = symbolData.drop_back();
   } else {
     loc = SMLoc::getFromPointer(symbolData.data());
 
     // If the dialect's symbol is followed immediately by a <, then lex the body
     // of it into prettyName.
     if (hasTrailingData && p.parseDialectSymbolBody(symbolData))
       return nullptr;
   }
 
   return createSymbol(dialectName, symbolData, loc);
 }
 
 /// Parse an extended attribute.
 ///
 ///   extended-attribute ::= (dialect-attribute | attribute-alias)
 ///   dialect-attribute  ::= `#` dialect-namespace `<` `"` attr-data `"` `>`
 ///   dialect-attribute  ::= `#` alias-name pretty-dialect-sym-body?
 ///   attribute-alias    ::= `#` alias-name
 ///
 Attribute Parser::parseExtendedAttr(Type type) {
   MLIRContext *ctx = getContext();
   Attribute attr = parseExtendedSymbol<Attribute>(
-      *this, Token::hash_identifier, state.symbols.attributeAliasDefinitions,
+      *this, state.symbols.attributeAliasDefinitions,
       [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Attribute {
         // Parse an optional trailing colon type.
         Type attrType = type;
         if (consumeIf(Token::colon) && !(attrType = parseType()))
           return Attribute();
 
         // If we found a registered dialect, then ask it to parse the attribute.
         if (Dialect *dialect =

 ///   extended-type ::= (dialect-type | type-alias)
 ///   dialect-type  ::= `!` dialect-namespace `<` `"` type-data `"` `>`
 ///   dialect-type  ::= `!` alias-name pretty-dialect-attribute-body?
 ///   type-alias    ::= `!` alias-name
 ///
 Type Parser::parseExtendedType() {
   MLIRContext *ctx = getContext();
   return parseExtendedSymbol<Type>(
-      *this, Token::exclamation_identifier, state.symbols.typeAliasDefinitions,
+      *this, state.symbols.typeAliasDefinitions,
       [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Type {
         // If we found a registered dialect, then ask it to parse the type.
         if (auto *dialect = ctx->getOrLoadDialect(dialectName)) {
           // Temporarily reset the lexer to let the dialect parse the type.
           const char *curLexerPos = getToken().getLoc().getPointer();
           resetToken(symbolData.data());
 
           // Parse the type.

mlir/lib/Parser/Lexer.h

public:
 
   /// Change the position of the lexer cursor.  The next token we lex will start
   /// at the designated point in the input.
   void resetPointer(const char *newPointer) { curPtr = newPointer; }
 
   /// Returns the start of the buffer.
   const char *getBufferBegin() { return curBuffer.data(); }
 
+  /// Return the code completion location of the lexer, or nullptr if there is
+  /// none.
+  const char *getCodeCompleteLoc() const { return codeCompleteLoc; }
+
 private:
   // Helpers.
   Token formToken(Token::Kind kind, const char *tokStart) {
     return Token(kind, StringRef(tokStart, curPtr - tokStart));
   }
 
   Token emitError(const char *loc, const Twine &message);
 

mlir/lib/Parser/Parser.h

public:
 
   /// Parse a comma separated list of elements that must have at least one entry
   /// in it.
   ParseResult
   parseCommaSeparatedList(function_ref<ParseResult()> parseElementFn) {
     return parseCommaSeparatedList(Delimiter::None, parseElementFn);
   }
 
-  ParseResult parseDialectSymbolBody(StringRef &body);
+  /// Parse the body of a dialect symbol, which starts and ends with <>'s, and
+  /// may be recursive. Return with the 'body' StringRef encompassing the entire
+  /// body. `isCodeCompletion` is set to true if the body contained a code
+  /// completion location, in which case the body is only populated up to the
+  /// completion.
+  ParseResult parseDialectSymbolBody(StringRef &body, bool &isCodeCompletion);
+  ParseResult parseDialectSymbolBody(StringRef &body) {
+    bool isCodeCompletion = false;
+    return parseDialectSymbolBody(body, isCodeCompletion);
+  }
 
   // We have two forms of parsing methods - those that return a non-null
   // pointer on success, and those that return a ParseResult to indicate whether
   // they returned a failure.  The second class fills in by-reference arguments
   // as the results of their action.
 
   //===--------------------------------------------------------------------===//
   // Error Handling
public:
 
   ParseResult codeCompleteDialectName();
   ParseResult codeCompleteOperationName(StringRef dialectName);
   ParseResult codeCompleteDialectOrElidedOpName(SMLoc loc);
   ParseResult codeCompleteStringDialectOrOperationName(StringRef name);
   ParseResult codeCompleteExpectedTokens(ArrayRef<StringRef> tokens);
   ParseResult codeCompleteOptionalTokens(ArrayRef<StringRef> tokens);
 
+  Attribute codeCompleteAttribute();
+  Type codeCompleteType();
+  Attribute
+  codeCompleteDialectSymbol(const llvm::StringMap<Attribute> &aliases);
+  Type codeCompleteDialectSymbol(const llvm::StringMap<Type> &aliases);
+
 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::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();
 }
 
+Attribute Parser::codeCompleteAttribute() {
+  state.codeCompleteContext->completeAttribute();
+  return {};
+}
+Type Parser::codeCompleteType() {
+  state.codeCompleteContext->completeType();
+  return {};
+}
+
+Attribute
+Parser::codeCompleteDialectSymbol(const llvm::StringMap<Attribute> &aliases) {
+  state.codeCompleteContext->completeDialectAttributeOrAlias(aliases);
+  return {};
+}
+Type Parser::codeCompleteDialectSymbol(const llvm::StringMap<Type> &aliases) {
+  state.codeCompleteContext->completeDialectTypeOrAlias(aliases);
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // OperationParser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for parsing operations and regions of
 /// operations.
 class OperationParser : public Parser {

mlir/lib/Parser/TypeParser.cpp

Type Parser::parseNonFunctionType() {
   // none-type
   case Token::kw_none:
     consumeToken(Token::kw_none);
     return builder.getNoneType();
 
   // extended type
   case Token::exclamation_identifier:
     return parseExtendedType();
+
+  // Handle completion of a dialect type.
+  case Token::code_complete:
+    if (getToken().isCodeCompletionFor(Token::exclamation_identifier))
+      return parseExtendedType();
+    return codeCompleteType();
   }
 }
 
 /// Parse a tensor type.
 ///
 ///   tensor-type ::= `tensor` `<` dimension-list type `>`
 ///   dimension-list ::= dimension-list-ranked | `*x`
 ///

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

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);
     }
   }
 
+  /// Signal a completion for an attribute.
+  void completeAttribute() override {
+    appendSimpleCompletions({"affine_set", "affine_map", "dense", "false",
+                             "loc", "opaque", "sparse", "true", "unit"},
+                            lsp::CompletionItemKind::Field);
+  }
+  void completeDialectAttributeOrAlias(
+      const llvm::StringMap<Attribute> &aliases) override {
+    completeDialectName();
+    completeAliases(aliases);
+  }
+
+  /// Signal a completion for a type.
+  void completeType() override {
+    appendSimpleCompletions({"memref", "tensor", "complex", "tuple", "vector",
+                             "bf16", "f16", "f32", "f64", "f80", "f128",
+                             "index", "none"},
+                            lsp::CompletionItemKind::Field);
+
+    lsp::CompletionItem item("iN", lsp::CompletionItemKind::Field);
+    item.insertText = "i";
+    completionList.items.emplace_back(item);
+  }
+  void
+  completeDialectTypeOrAlias(const llvm::StringMap<Type> &aliases) override {
+    completeDialectName();
+    completeAliases(aliases);
+  }
+
+  /// Add completion results for the given set of aliases.
+  template <typename T>
+  void completeAliases(const llvm::StringMap<T> &aliases) {
+    for (const auto &alias : aliases) {
+      lsp::CompletionItem item(alias.getKey(), lsp::CompletionItemKind::Field);
+      item.sortText = "0";
+      llvm::raw_string_ostream(item.detail) << "alias: " << alias.getValue();
+      completionList.items.emplace_back(item);
+    }
+  }
+
+  /// Add a set of simple completions that all have the same kind.
+  void appendSimpleCompletions(ArrayRef<StringRef> completions,
+                               lsp::CompletionItemKind kind) {
+    for (StringRef completion : completions)
+      completionList.items.emplace_back(completion, kind);
+  }
+
 private:
   lsp::CompletionList &completionList;
   MLIRContext *ctx;
 };
 } // namespace
 
 lsp::CompletionList
 MLIRDocument::getCodeCompletion(const lsp::URIForFile &uri,

mlir/test/lib/Dialect/Test/TestTypes.cpp

#include "TestTypeDefs.cpp.inc"
   registerDynamicType(getSingletonDynamicType(this));
   registerDynamicType(getPairDynamicType(this));
   registerDynamicType(getCustomAssemblyFormatDynamicType(this));
 }
 
 Type TestDialect::parseTestType(AsmParser &parser,
                                 SetVector<Type> &stack) const {
   StringRef typeTag;
-  if (failed(parser.parseKeyword(&typeTag)))
-    return Type();
-
   {
     Type genType;
-    auto parseResult = generatedTypeParser(parser, typeTag, genType);
+    auto parseResult = generatedTypeParser(parser, &typeTag, genType);
     if (parseResult.hasValue())
       return genType;
   }
 
   {
     Type dynType;
     auto parseResult = parseOptionalDynamicType(typeTag, parser, dynType);
     if (parseResult.hasValue()) {

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 %"
+  "text":"#attr = i32\n!alias = i32\nfunc.func private @foo(%arg: i32) -> i32 {\n%cast = \"builtin.unrealized_conversion_cast\"() : () -> (!pdl.value)\nreturn %"
 }}}
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
-  "position":{"line":0,"character":0}
+  "position":{"line":2,"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:             "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}
+  "position":{"line":3,"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}
+  "position":{"line":3,"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}
+  "position":{"line":4,"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:        "detail": "builtin.unrealized_conversion_cast: !pdl.value",
 // 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":1,"method":"textDocument/completion","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
-  "position":{"line":0,"character":10}
+  "position":{"line":2,"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": "nested",
 // CHECK-NEXT:        "sortText": "0"
 // CHECK-NEXT:      }
 // CHECK-NEXT:    ]
 // CHECK-NEXT:  }
 // -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":0,"character":8}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK:           {
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "false"
+// CHECK:           },
+// CHECK:           {
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "loc"
+// CHECK:           },
+// CHECK:           {
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "true"
+// CHECK:           },
+// CHECK:           {
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "unit"
+// CHECK:           }
+// CHECK:    ]
+// CHECK:  }
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":3,"character":56}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK:           {
+// CHECK:              "insertTextFormat": 1,
+// CHECK:              "kind": 5,
+// CHECK:             "label": "index"
+// CHECK:           },
+// CHECK:           {
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "none"
+// CHECK:           },
+// CHECK:           {
+// CHECK:             "insertText": "i",
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "iN"
+// CHECK:           }
+// CHECK-NEXT:    ]
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":3,"character":57}
+}}
+//      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": "alias: i32",
+// CHECK:             "insertTextFormat": 1,
+// CHECK:             "kind": 5,
+// CHECK:             "label": "alias",
+// CHECK:             "sortText": "0"
+// CHECK:           }
+// CHECK-NEXT:    ]
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":3,"character":61}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "attribute",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "operation",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "range",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "type",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "insertTextFormat": 1,
+// CHECK-NEXT:        "kind": 14,
+// CHECK-NEXT:        "label": "value",
+// CHECK-NEXT:        "sortText": "0"
+// CHECK-NEXT:      }
+// CHECK-NEXT:    ]
+// CHECK-NEXT:  }
+// -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-tblgen/attrdefs.td

 // DEF: #undef GET_ATTRDEF_LIST
 // DEF: ::test::SimpleAAttr,
 // DEF: ::test::CompoundAAttr,
 // DEF: ::test::IndexAttr,
 // DEF: ::test::SingleParameterAttr
 
 // DEF-LABEL: ::mlir::OptionalParseResult generatedAttributeParser(
 // DEF-SAME: ::mlir::AsmParser &parser,
-// DEF-SAME: ::llvm::StringRef mnemonic, ::mlir::Type type,
+// DEF-SAME: ::llvm::StringRef *mnemonic, ::mlir::Type type,
 // DEF-SAME: ::mlir::Attribute &value) {
-// DEF: if (mnemonic == ::test::CompoundAAttr::getMnemonic()) {
+// DEF: return ::mlir::AsmParser::KeywordSwitch<::mlir::OptionalParseResult>(parser)
+// DEF: .Case(::test::CompoundAAttr::getMnemonic()
 // DEF-NEXT: value = ::test::CompoundAAttr::parse(parser, type);
 // DEF-NEXT: return ::mlir::success(!!value);
-// DEF-NEXT: }
-// DEF-NEXT: if (mnemonic == ::test::IndexAttr::getMnemonic()) {
+// DEF-NEXT: })
+// DEF-NEXT: .Case(::test::IndexAttr::getMnemonic()
 // DEF-NEXT:   value = ::test::IndexAttr::parse(parser, type);
 // DEF-NEXT:   return ::mlir::success(!!value);
-// DEF: return {};
+// DEF: .Default([&](llvm::StringRef keyword, 
+// DEF-NEXT:   *mnemonic = keyword;
+// DEF-NEXT:   return llvm::None; 
 
 def Test_Dialect: Dialect {
 // DECL-NOT: TestDialect
 // DEF-NOT: TestDialect
     let name = "TestDialect";
     let cppNamespace = "::test";
 }
 

mlir/test/mlir-tblgen/default-type-attr-print-parser.td

 }
 
 // ATTR: namespace test {
 
 // ATTR: ::mlir::Attribute TestDialect::parseAttribute(::mlir::DialectAsmParser &parser,
 // ATTR:                                               ::mlir::Type type) const {
 // ATTR:   ::llvm::SMLoc typeLoc = parser.getCurrentLocation();
 // ATTR:   ::llvm::StringRef attrTag;
-// ATTR:   if (::mlir::failed(parser.parseKeyword(&attrTag)))
-// ATTR:     return {};
 // ATTR:   {
 // ATTR:     ::mlir::Attribute attr;
-// ATTR:     auto parseResult = generatedAttributeParser(parser, attrTag, type, attr);
+// ATTR:     auto parseResult = generatedAttributeParser(parser, &attrTag, type, attr);
 // ATTR:     if (parseResult.hasValue())
 // ATTR:       return attr;
 // ATTR:   }
 // ATTR:   parser.emitError(typeLoc) << "unknown attribute `"
 // ATTR:       << attrTag << "` in dialect `" << getNamespace() << "`";
 // ATTR:   return {}
 // ATTR: }
 
def TypeA : TestType<"TypeA"> {
   let mnemonic = "type_a";
 }
 
 // TYPE: namespace test {
 
 // TYPE: ::mlir::Type TestDialect::parseType(::mlir::DialectAsmParser &parser) const {
 // TYPE:   ::llvm::SMLoc typeLoc = parser.getCurrentLocation();
 // TYPE:   ::llvm::StringRef mnemonic;
-// TYPE:   if (parser.parseKeyword(&mnemonic))
-// TYPE:     return ::mlir::Type();
 // TYPE:   ::mlir::Type genType;
-// TYPE:   auto parseResult = generatedTypeParser(parser, mnemonic, genType);
+// TYPE:   auto parseResult = generatedTypeParser(parser, &mnemonic, genType);
 // TYPE:   if (parseResult.hasValue())
 // TYPE:     return genType;
 // TYPE:   parser.emitError(typeLoc) << "unknown  type `"
 // TYPE:       << mnemonic << "` in dialect `" << getNamespace() << "`";
 // TYPE:   return {};
 // TYPE: }
 
 // TYPE: void TestDialect::printType(::mlir::Type type,
 // TYPE:                     ::mlir::DialectAsmPrinter &printer) const {
 // TYPE:   if (::mlir::succeeded(generatedTypePrinter(type, printer)))
 // TYPE:     return;
 // TYPE: }
 
 // TYPE: } // namespace test

mlir/test/mlir-tblgen/typedefs.td

 // DEF: ::test::SimpleAType,
 // DEF: ::test::CompoundAType,
 // DEF: ::test::IndexType,
 // DEF: ::test::SingleParameterType,
 // DEF: ::test::IntegerType
 
 // DEF-LABEL: ::mlir::OptionalParseResult generatedTypeParser(
 // DEF-SAME: ::mlir::AsmParser &parser,
-// DEF-SAME: ::llvm::StringRef mnemonic,
+// DEF-SAME: ::llvm::StringRef *mnemonic,
 // DEF-SAME: ::mlir::Type &value) {
-// DEF: if (mnemonic == ::test::CompoundAType::getMnemonic()) {
+// DEF: .Case(::test::CompoundAType::getMnemonic()
 // DEF-NEXT:   value = ::test::CompoundAType::parse(parser);
 // DEF-NEXT:   return ::mlir::success(!!value);
-// DEF-NEXT: }
-// DEF-NEXT: if (mnemonic == ::test::IndexType::getMnemonic()) {
+// DEF-NEXT: })
+// DEF-NEXT: .Case(::test::IndexType::getMnemonic()
 // DEF-NEXT:   value = ::test::IndexType::parse(parser);
 // DEF-NEXT:   return ::mlir::success(!!value);
-// DEF: return {};
+// DEF: .Default([&](llvm::StringRef keyword, 
+// DEF-NEXT:   *mnemonic = keyword;
+// DEF-NEXT:   return llvm::None; 
 
 def Test_Dialect: Dialect {
 // DECL-NOT: TestDialect
   let name = "TestDialect";
   let cppNamespace = "::test";
 }
 
 class TestType<string name> : TypeDef<Test_Dialect, name> { }

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

 /// {1}: the optional code for the dynamic attribute parser dispatch.
 /// {2}: the optional code for the dynamic attribute printer dispatch.
 static const char *const dialectDefaultAttrPrinterParserDispatch = R"(
 /// Parse an attribute registered to this dialect.
 ::mlir::Attribute {0}::parseAttribute(::mlir::DialectAsmParser &parser,
                                       ::mlir::Type type) const {{
   ::llvm::SMLoc typeLoc = parser.getCurrentLocation();
   ::llvm::StringRef attrTag;
-  if (::mlir::failed(parser.parseKeyword(&attrTag)))
-    return {{};
   {{
     ::mlir::Attribute attr;
-    auto parseResult = generatedAttributeParser(parser, attrTag, type, attr);
+    auto parseResult = generatedAttributeParser(parser, &attrTag, type, attr);
     if (parseResult.hasValue())
       return attr;
   }
   {1}
   parser.emitError(typeLoc) << "unknown attribute `"
       << attrTag << "` in dialect `" << getNamespace() << "`";
   return {{};
 }

 /// {0}: the dialect fully qualified class name.
 /// {1}: the optional code for the dynamic type parser dispatch.
 /// {2}: the optional code for the dynamic type printer dispatch.
 static const char *const dialectDefaultTypePrinterParserDispatch = R"(
 /// Parse a type registered to this dialect.
 ::mlir::Type {0}::parseType(::mlir::DialectAsmParser &parser) const {{
   ::llvm::SMLoc typeLoc = parser.getCurrentLocation();
   ::llvm::StringRef mnemonic;
-  if (parser.parseKeyword(&mnemonic))
-    return ::mlir::Type();
   ::mlir::Type genType;
-  auto parseResult = generatedTypeParser(parser, mnemonic, genType);
+  auto parseResult = generatedTypeParser(parser, &mnemonic, genType);
   if (parseResult.hasValue())
     return genType;
   {1}
   parser.emitError(typeLoc) << "unknown  type `"
       << mnemonic << "` in dialect `" << getNamespace() << "`";
   return {{};
 }
 /// Print a type registered to this dialect.

 void DefGenerator::emitParsePrintDispatch(ArrayRef<AttrOrTypeDef> defs) {
   if (llvm::none_of(defs, [](const AttrOrTypeDef &def) {
         return def.getMnemonic().has_value();
       })) {
     return;
   }
   // Declare the parser.
   SmallVector<MethodParameter> params = {{"::mlir::AsmParser &", "parser"},
-                                         {"::llvm::StringRef", "mnemonic"}};
+                                         {"::llvm::StringRef *", "mnemonic"}};
   if (isAttrGenerator)
     params.emplace_back("::mlir::Type", "type");
   params.emplace_back(strfmt("::mlir::{0} &", valueType), "value");
   Method parse("::mlir::OptionalParseResult",
                strfmt("generated{0}Parser", valueType), Method::StaticInline,
                std::move(params));
   // Declare the printer.
   Method printer("::mlir::LogicalResult",
                  strfmt("generated{0}Printer", valueType), Method::StaticInline,
                  {{strfmt("::mlir::{0}", valueType), "def"},
                   {"::mlir::AsmPrinter &", "printer"}});
 
-  // The parser dispatch is just a list of if-elses, matching on the mnemonic
-  // and calling the def's parse function.
+  // The parser dispatch uses a KeywordSwitch, matching on the mnemonic and
+  // calling the def's parse function.
+  parse.body() << "  return "
+                  "::mlir::AsmParser::KeywordSwitch<::mlir::"
+                  "OptionalParseResult>(parser)\n";
   const char *const getValueForMnemonic =
-      R"(  if (mnemonic == {0}::getMnemonic()) {{
+      R"(    .Case({0}::getMnemonic(), [&](llvm::StringRef, llvm::SMLoc) {{
       value = {0}::{1};
       return ::mlir::success(!!value);
-  }
+    })
 )";
+
   // The printer dispatch uses llvm::TypeSwitch to find and call the correct
   // printer.
   printer.body() << "  return ::llvm::TypeSwitch<::mlir::" << valueType
                  << ", ::mlir::LogicalResult>(def)";
   const char *const printValue = R"(    .Case<{0}>([&](auto t) {{
       printer << {0}::getMnemonic();{1}
       return ::mlir::success();
     })
for (auto &def : defs) {
     parse.body() << llvm::formatv(getValueForMnemonic, defClass, parseOrGet);
 
     // If the def has no parameters and no printer, just print the mnemonic.
     StringRef printDef = "";
     if (hasParserPrinterDecl)
       printDef = "\nt.print(printer);";
     printer.body() << llvm::formatv(printValue, defClass, printDef);
   }
-  parse.body() << "  return {};";
+  parse.body() << "    .Default([&](llvm::StringRef keyword, llvm::SMLoc) {\n"
+                  "      *mnemonic = keyword;\n"
+                  "      return llvm::None;\n"
+                  "    });";
   printer.body() << "    .Default([](auto) { return ::mlir::failure(); });";
 
   raw_indented_ostream indentedOs(os);
   parse.writeDeclTo(indentedOs);
   printer.writeDeclTo(indentedOs);
 }
 
 bool DefGenerator::emitDefs(StringRef selectedDialect) {