[mlir] Add an AsmParserConfig class to start encapsulating parser options

[rriddle]

Summary

We are currently growing quite a few parameters to the various parser
commands, and more are to be added going forward. This commit
adds a new AsmParserConfig class to simplify method definitions, and
make it easier to inject information into the parser.

mlir/include/mlir/Parser/Parser.h

inline OwningOpRef<ContainerOpT> constructContainerOpForParserIfNecessary(
   // After splicing, verify just this operation to ensure it can properly
   // contain the operations inside of it.
   if (failed(op.verifyInvariants()))
     return OwningOpRef<ContainerOpT>();
   return opRef;
 }
 } // namespace detail
 
+/// This class represents a configuration for the MLIR assembly parser. It
+/// contains all of the necessary state to parse a textual MLIR source file.
+class AsmParserConfig {
+public:
+  AsmParserConfig(MLIRContext *context, AsmParserState *asmState = nullptr)
+      : context(context), asmState(asmState) {
+    assert(context && "expected valid MLIR context");
+  }
+
+  /// Return the MLIRContext to be used when parsing.
+  MLIRContext *getContext() const { return context; }
+
+  /// Return asm parser state to be used when parsing.
+  AsmParserState *getAsmParserState() const { return asmState; }
+
+private:
+  MLIRContext *context;
+  AsmParserState *asmState;
+};
+
 /// This parses the file specified by the indicated SourceMgr 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. 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,
-                              MLIRContext *context,
-                              LocationAttr *sourceFileLoc = nullptr,
-                              AsmParserState *asmState = nullptr);
+                              const AsmParserConfig &config,
+                              LocationAttr *sourceFileLoc = 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.
-LogicalResult parseSourceFile(llvm::StringRef filename, Block *block,
-                              MLIRContext *context,
+LogicalResult parseSourceFile(StringRef filename, Block *block,
+                              const AsmParserConfig &config,
                               LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the file specified by the indicated filename using the provided
 /// SourceMgr 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. 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(llvm::StringRef filename,
-                              llvm::SourceMgr &sourceMgr, Block *block,
-                              MLIRContext *context,
-                              LocationAttr *sourceFileLoc = nullptr,
-                              AsmParserState *asmState = nullptr);
+LogicalResult parseSourceFile(StringRef filename, llvm::SourceMgr &sourceMgr,
+                              Block *block, const AsmParserConfig &config,
+                              LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the IR string 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.
-LogicalResult parseSourceString(llvm::StringRef sourceStr, Block *block,
-                                MLIRContext *context,
+LogicalResult parseSourceString(StringRef sourceStr, Block *block,
+                                const AsmParserConfig &config,
                                 LocationAttr *sourceFileLoc = nullptr);
 
+namespace detail {
+/// The internal implementation of the templated `parseSourceFile` methods
+/// below, that simply forwards to the non-templated version.
+template <typename ContainerOpT, typename... ParserArgs>
+inline OwningOpRef<ContainerOpT> parseSourceFile(const AsmParserConfig &config,
+                                                 ParserArgs &&...args) {
+  LocationAttr sourceFileLoc;
+  Block block;
+  if (failed(parseSourceFile(std::forward<ParserArgs>(args)..., &block, config,
+                             &sourceFileLoc)))
+    return OwningOpRef<ContainerOpT>();
+  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
+      &block, config.getContext(), sourceFileLoc);
+}
+} // namespace detail
+
 /// This parses the file specified by the indicated SourceMgr. If the source IR
 /// contained a single instance of `ContainerOpT`, it is returned. Otherwise, a
 /// new instance of `ContainerOpT` is constructed containing all of the parsed
 /// operations. If parsing was not successful, null is returned and an error
 /// message is emitted through the error handler registered in the context, and
 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the
 /// `SingleBlockImplicitTerminator` trait.
 template <typename ContainerOpT>
 inline OwningOpRef<ContainerOpT>
-parseSourceFile(const llvm::SourceMgr &sourceMgr, MLIRContext *context) {
-  LocationAttr sourceFileLoc;
-  Block block;
-  if (failed(parseSourceFile(sourceMgr, &block, context, &sourceFileLoc)))
-    return OwningOpRef<ContainerOpT>();
-  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
-      &block, context, sourceFileLoc);
+parseSourceFile(const llvm::SourceMgr &sourceMgr,
+                const AsmParserConfig &config) {
+  return detail::parseSourceFile<ContainerOpT>(config, sourceMgr);
 }
 
 /// This parses the file specified by the indicated filename. If the source IR
 /// contained a single instance of `ContainerOpT`, it is returned. Otherwise, a
 /// new instance of `ContainerOpT` is constructed containing all of the parsed
 /// operations. If parsing was not successful, null is returned and an error
 /// message is emitted through the error handler registered in the context, and
 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the
 /// `SingleBlockImplicitTerminator` trait.
 template <typename ContainerOpT>
-inline OwningOpRef<ContainerOpT> parseSourceFile(llvm::StringRef filename,
-                                                 MLIRContext *context) {
-  LocationAttr sourceFileLoc;
-  Block block;
-  if (failed(parseSourceFile(filename, &block, context, &sourceFileLoc)))
-    return OwningOpRef<ContainerOpT>();
-  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
-      &block, context, sourceFileLoc);
+inline OwningOpRef<ContainerOpT>
+parseSourceFile(StringRef filename, const AsmParserConfig &config) {
+  return detail::parseSourceFile<ContainerOpT>(config, filename);
 }
 
 /// This parses the file specified by the indicated filename using the provided
 /// SourceMgr. If the source IR contained a single instance of `ContainerOpT`,
 /// it is returned. Otherwise, a new instance of `ContainerOpT` is constructed
 /// containing all of the parsed operations. If parsing was not successful, null
 /// is returned and an error message is emitted through the error handler
 /// registered in the context, and failure is returned. `ContainerOpT` is
 /// required to have a single region containing a single block, and must
 /// implement the `SingleBlockImplicitTerminator` trait.
 template <typename ContainerOpT>
-inline OwningOpRef<ContainerOpT> parseSourceFile(llvm::StringRef filename,
-                                                 llvm::SourceMgr &sourceMgr,
-                                                 MLIRContext *context) {
-  LocationAttr sourceFileLoc;
-  Block block;
-  if (failed(parseSourceFile(filename, sourceMgr, &block, context,
-                             &sourceFileLoc)))
-    return OwningOpRef<ContainerOpT>();
-  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
-      &block, context, sourceFileLoc);
+inline OwningOpRef<ContainerOpT>
+parseSourceFile(llvm::StringRef filename, llvm::SourceMgr &sourceMgr,
+                const AsmParserConfig &config) {
+  return detail::parseSourceFile<ContainerOpT>(config, filename, sourceMgr);
 }
 
 /// This parses the provided string containing MLIR. If the source IR contained
 /// a single instance of `ContainerOpT`, it is returned. Otherwise, a new
 /// instance of `ContainerOpT` is constructed containing all of the parsed
 /// operations. If parsing was not successful, null is returned and an error
 /// message is emitted through the error handler registered in the context, and
 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the
 /// `SingleBlockImplicitTerminator` trait.
 template <typename ContainerOpT>
-inline OwningOpRef<ContainerOpT> parseSourceString(llvm::StringRef sourceStr,
-                                                   MLIRContext *context) {
+inline OwningOpRef<ContainerOpT>
+parseSourceString(StringRef sourceStr, const AsmParserConfig &config) {
   LocationAttr sourceFileLoc;
   Block block;
-  if (failed(parseSourceString(sourceStr, &block, context, &sourceFileLoc)))
+  if (failed(parseSourceString(sourceStr, &block, config, &sourceFileLoc)))
     return OwningOpRef<ContainerOpT>();
   return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
-      &block, context, sourceFileLoc);
+      &block, config.getContext(), sourceFileLoc);
 }
 
 /// This parses a single MLIR attribute to an MLIR context if it was valid.  If
 /// not, an error message is emitted through a new SourceMgrDiagnosticHandler
 /// constructed from a new SourceMgr with a single a MemoryBuffer wrapping
 /// `attrStr`. If the passed `attrStr` has additional tokens that were not part
 /// of the type, an error is emitted.
 // TODO: Improve diagnostic reporting.

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;
-  ParserState state(sourceMgr, context, symbolState, /*asmState=*/nullptr);
+  AsmParserConfig config(context);
+  ParserState state(sourceMgr, config, symbolState);
   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/AsmParserImpl.h

if (atToken.isNot(Token::at_identifier))
       return failure();
 
     result = getBuilder().getStringAttr(atToken.getSymbolReference());
     attrs.push_back(getBuilder().getNamedAttr(attrName, result));
     parser.consumeToken();
 
     // If we are populating the assembly parser state, record this as a symbol
     // reference.
-    if (parser.getState().asmState) {
-      parser.getState().asmState->addUses(SymbolRefAttr::get(result),
-                                          atToken.getLocRange());
-    }
+    if (auto *asmState = parser.getState().config.getAsmParserState())
+      asmState->addUses(SymbolRefAttr::get(result), atToken.getLocRange());
     return success();
   }
 
   //===--------------------------------------------------------------------===//
   // Type Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a type.

mlir/lib/Parser/AttributeParser.cpp

return type ? StringAttr::get(val, type)
                 : StringAttr::get(getContext(), val);
   }
 
   // Parse a symbol reference attribute.
   case Token::at_identifier: {
     // When populating the parser state, this is a list of locations for all of
     // the nested references.
     SmallVector<SMRange> referenceLocations;
-    if (state.asmState)
+    if (state.config.getAsmParserState())
       referenceLocations.push_back(getToken().getLocRange());
 
     // Parse the top-level reference.
     std::string nameStr = getToken().getSymbolReference();
     consumeToken(Token::at_identifier);
 
     // Parse any nested references.
     std::vector<FlatSymbolRefAttr> nestedRefs;
while (getToken().is(Token::colon)) {
       auto curLoc = getToken().getLoc();
       if (getToken().isNot(Token::at_identifier)) {
         emitError(curLoc, "expected nested symbol reference identifier");
         return Attribute();
       }
 
       // If we are populating the assembly state, add the location for this
       // reference.
-      if (state.asmState)
+      if (state.config.getAsmParserState())
         referenceLocations.push_back(getToken().getLocRange());
 
       std::string nameStr = getToken().getSymbolReference();
       consumeToken(Token::at_identifier);
       nestedRefs.push_back(SymbolRefAttr::get(getContext(), nameStr));
     }
     SymbolRefAttr symbolRefAttr =
         SymbolRefAttr::get(getContext(), nameStr, nestedRefs);
 
     // If we are populating the assembly state, record this symbol reference.
-    if (state.asmState)
-      state.asmState->addUses(symbolRefAttr, referenceLocations);
+    if (auto *asmState = state.config.getAsmParserState())
+      asmState->addUses(symbolRefAttr, referenceLocations);
     return symbolRefAttr;
   }
 
   // Parse a 'unit' attribute.
   case Token::kw_unit:
     consumeToken(Token::kw_unit);
     return builder.getUnitAttr();
 

mlir/lib/Parser/DialectSymbolParser.cpp

static Symbol parseExtendedSymbol(Parser &p, Token::Kind identifierTok,
   p.getState().symbols.nestedParserLocs.pop_back();
   return sym;
 }
 
 /// Parses a symbol, of type 'T', and returns it if parsing was successful. If
 /// parsing failed, nullptr is returned. The number of bytes read from the input
 /// string is returned in 'numRead'.
 template <typename T, typename ParserFn>
-static T parseSymbol(StringRef inputStr, MLIRContext *context,
+static T parseSymbol(StringRef inputStr, const AsmParserConfig &config,
                      SymbolState &symbolState, ParserFn &&parserFn,
                      size_t *numRead = nullptr) {
   SourceMgr sourceMgr;
   auto memBuffer = MemoryBuffer::getMemBuffer(
       inputStr, /*BufferName=*/"<mlir_parser_buffer>",
       /*RequiresNullTerminator=*/false);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
-  ParserState state(sourceMgr, context, symbolState, /*asmState=*/nullptr);
+  ParserState state(sourceMgr, config, symbolState);
   Parser parser(state);
 
   Token startTok = parser.getToken();
   T symbol = parserFn(parser);
   if (!symbol)
     return T();
 
   // If 'numRead' is valid, then provide the number of bytes that were read.

 /// 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 *context = getContext();
   Attribute attr = parseExtendedSymbol<Attribute>(
       *this, Token::hash_identifier, state.symbols.attributeAliasDefinitions,
-      [&](StringRef dialectName, StringRef symbolData,
-          SMLoc loc) -> Attribute {
+      [&](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 =
                 builder.getContext()->getOrLoadDialect(dialectName)) {
           return parseSymbol<Attribute>(
-              symbolData, state.context, state.symbols, [&](Parser &parser) {
+              symbolData, context, state.symbols, [&](Parser &parser) {
                 CustomDialectAsmParser customParser(symbolData, parser);
                 return dialect->parseAttribute(customParser, attrType);
               });
         }
 
         // Otherwise, form a new opaque attribute.
         return OpaqueAttr::getChecked(
             [&] { return emitError(loc); },
-            StringAttr::get(state.context, dialectName), symbolData,
-            attrType ? attrType : NoneType::get(state.context));
+            StringAttr::get(context, dialectName), symbolData,
+            attrType ? attrType : NoneType::get(context));
       });
 
   // Ensure that the attribute has the same type as requested.
   if (attr && type && attr.getType() != type) {
     emitError("attribute type different than expected: expected ")
         << type << ", but got " << attr.getType();
     return nullptr;
   }
   return attr;
 }
 
 /// Parse an extended type.
 ///
 ///   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 *context = getContext();
   return parseExtendedSymbol<Type>(
       *this, Token::exclamation_identifier, state.symbols.typeAliasDefinitions,
-      [&](StringRef dialectName, StringRef symbolData,
-          SMLoc loc) -> Type {
+      [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Type {
         // If we found a registered dialect, then ask it to parse the type.
-        auto *dialect = state.context->getOrLoadDialect(dialectName);
+        auto *dialect = context->getOrLoadDialect(dialectName);
 
         if (dialect) {
           return parseSymbol<Type>(
-              symbolData, state.context, state.symbols, [&](Parser &parser) {
+              symbolData, context, state.symbols, [&](Parser &parser) {
                 CustomDialectAsmParser customParser(symbolData, parser);
                 return dialect->parseType(customParser);
               });
         }
 
         // Otherwise, form a new opaque type.
-        return OpaqueType::getChecked(
-            [&] { return emitError(loc); },
-            StringAttr::get(state.context, dialectName), symbolData);
+        return OpaqueType::getChecked([&] { return emitError(loc); },
+                                      StringAttr::get(context, dialectName),
+                                      symbolData);
       });
 }
 
 //===----------------------------------------------------------------------===//
 // mlir::parseAttribute/parseType
 //===----------------------------------------------------------------------===//
 
 /// Parses a symbol, of type 'T', and returns it if parsing was successful. If

mlir/lib/Parser/Parser.h

 /// types. It is intended to be subclassed by specialized subparsers that
 /// include state.
 class Parser {
 public:
   using Delimiter = OpAsmParser::Delimiter;
 
   Builder builder;
 
-  Parser(ParserState &state) : builder(state.context), state(state) {}
+  Parser(ParserState &state)
+      : builder(state.config.getContext()), state(state) {}
 
   // Helper methods to get stuff from the parser-global state.
   ParserState &getState() const { return state; }
-  MLIRContext *getContext() const { return state.context; }
+  MLIRContext *getContext() const { return builder.getContext(); }
   const llvm::SourceMgr &getSourceMgr() { return state.lex.getSourceMgr(); }
 
   /// Parse a comma-separated list of elements up until the specified end token.
   ParseResult
   parseCommaSeparatedListUntil(Token::Kind rightToken,
                                function_ref<ParseResult()> parseElement,
                                bool allowEmptyList = true);
 

mlir/lib/Parser/Parser.cpp

 MLIR_DEFINE_EXPLICIT_TYPE_ID(OperationParser::DeferredLocInfo *)
 
 OperationParser::OperationParser(ParserState &state, ModuleOp topLevelOp)
     : Parser(state), opBuilder(topLevelOp.getRegion()), topLevelOp(topLevelOp) {
   // The top level operation starts a new name scope.
   pushSSANameScope(/*isIsolated=*/true);
 
   // If we are populating the parser state, prepare it for parsing.
-  if (state.asmState)
-    state.asmState->initialize(topLevelOp);
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->initialize(topLevelOp);
 }
 
 OperationParser::~OperationParser() {
   for (auto &fwd : forwardRefPlaceholders) {
     // Drop all uses of undefined forward declared reference and destroy
     // defining operation.
     fwd.first.dropAllUses();
     fwd.first.getDefiningOp()->destroy();
ParseResult OperationParser::finalize() {
   if (failed(popSSANameScope()))
     return failure();
 
   // Verify that the parsed operations are valid.
   if (failed(verify(topLevelOp)))
     return failure();
 
   // If we are populating the parser state, finalize the top-level operation.
-  if (state.asmState)
-    state.asmState->finalize(topLevelOp);
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->finalize(topLevelOp);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SSA Value Handling
 //===----------------------------------------------------------------------===//
 
 void OperationParser::pushSSANameScope(bool isIsolated) {
if (auto existing = entries[useInfo.number].value) {
     // the actual definition instead, delete the forward ref, and remove it
     // from our set of forward references we track.
     existing.replaceAllUsesWith(value);
     existing.getDefiningOp()->destroy();
     forwardRefPlaceholders.erase(existing);
 
     // If a definition of the value already exists, replace it in the assembly
     // state.
-    if (state.asmState)
-      state.asmState->refineDefinition(existing, value);
+    if (auto *asmState = state.config.getAsmParserState())
+      asmState->refineDefinition(existing, value);
   }
 
   /// Record this definition for the current scope.
   entries[useInfo.number] = {value, useInfo.location};
   recordDefinition(useInfo.name);
   return success();
 }
 

 /// Given an unbound reference to an SSA value and its type, return the value
 /// it specifies.  This returns null on failure.
 Value OperationParser::resolveSSAUse(UnresolvedOperand useInfo, Type type) {
   auto &entries = getSSAValueEntry(useInfo.name);
 
   // Functor used to record the use of the given value if the assembly state
   // field is populated.
   auto maybeRecordUse = [&](Value value) {
-    if (state.asmState)
-      state.asmState->addUses(value, useInfo.location);
+    if (auto *asmState = state.config.getAsmParserState())
+      asmState->addUses(value, useInfo.location);
     return value;
   };
 
   // If we have already seen a value of this name, return it.
   if (useInfo.number < entries.size() && entries[useInfo.number].value) {
     Value result = entries[useInfo.number].value;
     // Check that the type matches the other uses.
     if (result.getType() == type)
if (!resultIDs.empty()) {
     if (op->getNumResults() == 0)
       return emitError(loc, "cannot name an operation with no results");
     if (numExpectedResults != op->getNumResults())
       return emitError(loc, "operation defines ")
              << op->getNumResults() << " results but was provided "
              << numExpectedResults << " to bind";
 
     // Add this operation to the assembly state if it was provided to populate.
-    if (state.asmState) {
+    if (auto *asmState = state.config.getAsmParserState()) {
       unsigned resultIt = 0;
       SmallVector<std::pair<unsigned, SMLoc>> asmResultGroups;
       asmResultGroups.reserve(resultIDs.size());
       for (ResultRecord &record : resultIDs) {
         asmResultGroups.emplace_back(resultIt, std::get<2>(record));
         resultIt += std::get<1>(record);
       }
-      state.asmState->finalizeOperationDefinition(
-          op, nameTok.getLocRange(), /*endLoc=*/getToken().getLoc(),
+      asmState->finalizeOperationDefinition(op, nameTok.getLocRange(),
+                                            /*endLoc=*/getToken().getLoc(),
                                             asmResultGroups);
     }
 
     // Add definitions for each of the result groups.
     unsigned opResI = 0;
     for (ResultRecord &resIt : resultIDs) {
       for (unsigned subRes : llvm::seq<unsigned>(0, std::get<1>(resIt))) {
         if (addDefinition({std::get<2>(resIt), std::get<0>(resIt), subRes},
                           op->getResult(opResI++)))
           return failure();
       }
     }
 
     // Add this operation to the assembly state if it was provided to populate.
-  } else if (state.asmState) {
-    state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange(),
-                                                /*endLoc=*/getToken().getLoc());
+  } else if (auto *asmState = state.config.getAsmParserState()) {
+    asmState->finalizeOperationDefinition(op, nameTok.getLocRange(),
   }
 
   return success();
 }
 
 /// Parse a single operation successor.
 ///
 ///   successor ::= block-id
if (!getContext()->getLoadedDialect(dialectName) &&
       emitError("operation being parsed with an unregistered dialect. If "
                 "this is intended, please use -allow-unregistered-dialect "
                 "with the MLIR tool used");
       return nullptr;
     }
   }
 
   // If we are populating the parser state, start a new operation definition.
-  if (state.asmState)
-    state.asmState->startOperationDefinition(result.name);
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->startOperationDefinition(result.name);
 
   if (parseGenericOperationAfterOpName(result))
     return nullptr;
 
   // Create the operation and try to parse a location for it.
   Operation *op = opBuilder.create(result);
   if (parseTrailingLocationSpecifier(op))
     return nullptr;
   return op;
 }
 
 Operation *OperationParser::parseGenericOperation(Block *insertBlock,
                                                   Block::iterator insertPt) {
   Token nameToken = getToken();
 
   OpBuilder::InsertionGuard restoreInsertionPoint(opBuilder);
   opBuilder.setInsertionPoint(insertBlock, insertPt);
   Operation *op = parseGenericOperation();
   if (!op)
     return nullptr;
 
   // If we are populating the parser asm state, finalize this operation
   // definition.
-  if (state.asmState)
-    state.asmState->finalizeOperationDefinition(op, nameToken.getLocRange(),
-                                                /*endLoc=*/getToken().getLoc());
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->finalizeOperationDefinition(op, nameToken.getLocRange(),
   return op;
 }
 
 namespace {
 class CustomOpAsmParser : public AsmParserImpl<OpAsmParser> {
 public:
   CustomOpAsmParser(
       SMLoc nameLoc, ArrayRef<OperationParser::ResultRecord> resultIDs,
OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
   llvm::PrettyStackTraceFormat fmt("MLIR Parser: custom op parser '%s'",
                                    opNameInfo->getIdentifier().data());
 
   // Get location information for the operation.
   auto srcLocation = getEncodedSourceLocation(opLoc);
   OperationState opState(srcLocation, *opNameInfo);
 
   // If we are populating the parser state, start a new operation definition.
-  if (state.asmState)
-    state.asmState->startOperationDefinition(opState.name);
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->startOperationDefinition(opState.name);
 
   // Have the op implementation take a crack and parsing this.
   CleanupOpStateRegions guard{opState};
   CustomOpAsmParser opAsmParser(opLoc, resultIDs, parseAssemblyFn,
                                 isIsolatedFromAbove, opName, *this);
   if (opAsmParser.parseOperation(opState))
     return nullptr;
 
ParseResult OperationParser::parseRegion(Region &region,
                                          ArrayRef<Argument> entryArguments,
                                          bool isIsolatedNameScope) {
   // Parse the '{'.
   Token lBraceTok = getToken();
   if (parseToken(Token::l_brace, "expected '{' to begin a region"))
     return failure();
 
   // If we are populating the parser state, start a new region definition.
-  if (state.asmState)
-    state.asmState->startRegionDefinition();
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->startRegionDefinition();
 
   // Parse the region body.
   if ((!entryArguments.empty() || getToken().isNot(Token::r_brace)) &&
       parseRegionBody(region, lBraceTok.getLoc(), entryArguments,
                       isIsolatedNameScope)) {
     return failure();
   }
   consumeToken(Token::r_brace);
 
   // If we are populating the parser state, finalize this region.
-  if (state.asmState)
-    state.asmState->finalizeRegionDefinition();
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->finalizeRegionDefinition();
 
   return success();
 }
 
 ParseResult OperationParser::parseRegionBody(Region &region, SMLoc startLoc,
                                              ArrayRef<Argument> entryArguments,
                                              bool isIsolatedNameScope) {
   auto currentPt = opBuilder.saveInsertionPoint();
 
   // Push a new named value scope.
   pushSSANameScope(isIsolatedNameScope);
 
   // Parse the first block directly to allow for it to be unnamed.
   auto owningBlock = std::make_unique<Block>();
   Block *block = owningBlock.get();
 
   // If this block is not defined in the source file, add a definition for it
   // now in the assembly state. Blocks with a name will be defined when the name
   // is parsed.
-  if (state.asmState && getToken().isNot(Token::caret_identifier))
-    state.asmState->addDefinition(block, startLoc);
+  if (state.config.getAsmParserState() &&
+      getToken().isNot(Token::caret_identifier))
+    state.config.getAsmParserState()->addDefinition(block, startLoc);
 
   // Add arguments to the entry block if we had the form with explicit names.
   if (!entryArguments.empty() && !entryArguments[0].ssaName.name.empty()) {
     // If we had named arguments, then don't allow a block name.
     if (getToken().is(Token::caret_identifier))
       return emitError("invalid block name in region with named arguments");
 
     for (auto &entryArg : entryArguments) {
       auto &argInfo = entryArg.ssaName;
 
       // Ensure that the argument was not already defined.
       if (auto defLoc = getReferenceLoc(argInfo.name, argInfo.number)) {
         return emitError(argInfo.location, "region entry argument '" +
                                                argInfo.name +
                                                "' is already in use")
                    .attachNote(getEncodedSourceLocation(*defLoc))
                << "previously referenced here";
       }
       Location loc = entryArg.sourceLoc.hasValue()
                          ? entryArg.sourceLoc.getValue()
                          : getEncodedSourceLocation(argInfo.location);
       BlockArgument arg = block->addArgument(entryArg.type, loc);
 
       // Add a definition of this arg to the assembly state if provided.
-      if (state.asmState)
-        state.asmState->addDefinition(arg, argInfo.location);
+      if (auto *asmState = state.config.getAsmParserState())
+        asmState->addDefinition(arg, argInfo.location);
 
       // Record the definition for this argument.
       if (addDefinition(argInfo, arg))
         return failure();
     }
   }
 
   if (parseBlock(block))
if (!blockAndLoc.block) {
     // removed once defined, so if we are defining a existing block and it is
     // not a forward declaration, then it is a redeclaration. Fail if the block
     // was already defined.
   } else if (!eraseForwardRef(blockAndLoc.block)) {
     return emitError(nameLoc, "redefinition of block '") << name << "'";
   }
 
   // Populate the high level assembly state if necessary.
-  if (state.asmState)
-    state.asmState->addDefinition(blockAndLoc.block, nameLoc);
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->addDefinition(blockAndLoc.block, nameLoc);
 
   block = blockAndLoc.block;
 
   // If an argument list is present, parse it.
   if (getToken().is(Token::l_paren))
     if (parseOptionalBlockArgList(block))
       return failure();
 

 Block *OperationParser::getBlockNamed(StringRef name, SMLoc loc) {
   BlockDefinition &blockDef = getBlockInfoByName(name);
   if (!blockDef.block) {
     blockDef = {new Block(), loc};
     insertForwardRef(blockDef.block, blockDef.loc);
   }
 
   // Populate the high level assembly state if necessary.
-  if (state.asmState)
-    state.asmState->addUses(blockDef.block, loc);
+  if (auto *asmState = state.config.getAsmParserState())
+    asmState->addUses(blockDef.block, loc);
 
   return blockDef.block;
 }
 
 /// Parse a (possibly empty) list of SSA operands with types as block arguments
 /// enclosed in parentheses.
 ///
 ///   value-id-and-type-list ::= value-id-and-type (`,` ssa-id-and-type)*
return parseSSADefOrUseAndType(
 
           // If the argument has an explicit loc(...) specifier, parse and apply
           // it.
           if (parseTrailingLocationSpecifier(arg))
             return failure();
 
           // Mark this block argument definition in the parser state if it was
           // provided.
-          if (state.asmState)
-            state.asmState->addDefinition(arg, useInfo.location);
+          if (auto *asmState = state.config.getAsmParserState())
+            asmState->addDefinition(arg, useInfo.location);
 
           return addDefinition(useInfo, arg);
         });
   });
 }
 
 //===----------------------------------------------------------------------===//
 // Top-level entity parsing.
case Token::exclamation_identifier:
       break;
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::parseSourceFile(const llvm::SourceMgr &sourceMgr,
-                                    Block *block, MLIRContext *context,
-                                    LocationAttr *sourceFileLoc,
-                                    AsmParserState *asmState) {
+                                    Block *block, const AsmParserConfig &config,
+                                    LocationAttr *sourceFileLoc) {
   const auto *sourceBuf = sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
 
-  Location parserLoc = FileLineColLoc::get(
-      context, sourceBuf->getBufferIdentifier(), /*line=*/0, /*column=*/0);
+  Location parserLoc =
+      FileLineColLoc::get(config.getContext(), sourceBuf->getBufferIdentifier(),
+                          /*line=*/0, /*column=*/0);
   if (sourceFileLoc)
     *sourceFileLoc = parserLoc;
 
   SymbolState aliasState;
-  ParserState state(sourceMgr, context, aliasState, asmState);
+  ParserState state(sourceMgr, config, aliasState);
   return TopLevelOperationParser(state).parse(block, parserLoc);
 }
 
 LogicalResult mlir::parseSourceFile(llvm::StringRef filename, Block *block,
-                                    MLIRContext *context,
+                                    const AsmParserConfig &config,
                                     LocationAttr *sourceFileLoc) {
   llvm::SourceMgr sourceMgr;
-  return parseSourceFile(filename, sourceMgr, block, context, sourceFileLoc);
+  return parseSourceFile(filename, sourceMgr, block, config, sourceFileLoc);
 }
 
 LogicalResult mlir::parseSourceFile(llvm::StringRef filename,
                                     llvm::SourceMgr &sourceMgr, Block *block,
-                                    MLIRContext *context,
-                                    LocationAttr *sourceFileLoc,
-                                    AsmParserState *asmState) {
+                                    const AsmParserConfig &config,
+                                    LocationAttr *sourceFileLoc) {
   if (sourceMgr.getNumBuffers() != 0) {
     // TODO: Extend to support multiple buffers.
-    return emitError(mlir::UnknownLoc::get(context),
+    return emitError(mlir::UnknownLoc::get(config.getContext()),
                      "only main buffer parsed at the moment");
   }
   auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(filename);
   if (std::error_code error = fileOrErr.getError())
-    return emitError(mlir::UnknownLoc::get(context),
+    return emitError(mlir::UnknownLoc::get(config.getContext()),
                      "could not open input file " + filename);
 
   // Load the MLIR source file.
   sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), SMLoc());
-  return parseSourceFile(sourceMgr, block, context, sourceFileLoc, asmState);
+  return parseSourceFile(sourceMgr, block, config, sourceFileLoc);
 }
 
 LogicalResult mlir::parseSourceString(llvm::StringRef sourceStr, Block *block,
-                                      MLIRContext *context,
+                                      const AsmParserConfig &config,
                                       LocationAttr *sourceFileLoc) {
   auto memBuffer = MemoryBuffer::getMemBuffer(sourceStr);
   if (!memBuffer)
     return failure();
 
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
-  return parseSourceFile(sourceMgr, block, context, sourceFileLoc);
+  return parseSourceFile(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, MLIRContext *ctx,
-              SymbolState &symbols, AsmParserState *asmState)
-      : context(ctx), lex(sourceMgr, ctx), curToken(lex.lexToken()),
-        symbols(symbols), parserDepth(symbols.nestedParserLocs.size()),
-        asmState(asmState) {
+  ParserState(const llvm::SourceMgr &sourceMgr, const AsmParserConfig &config,
+              SymbolState &symbols)
+      : config(config), lex(sourceMgr, config.getContext()),
+        curToken(lex.lexToken()), symbols(symbols),
+        parserDepth(symbols.nestedParserLocs.size()) {
     // Set the top level lexer for the symbol state if one doesn't exist.
     if (!symbols.topLevelLexer)
       symbols.topLevelLexer = &lex;
   }
   ~ParserState() {
     // Reset the top level lexer if it refers the lexer in our state.
     if (symbols.topLevelLexer == &lex)
       symbols.topLevelLexer = nullptr;
   }
   ParserState(const ParserState &) = delete;
   void operator=(const ParserState &) = delete;
 
-  /// The context we're parsing into.
-  MLIRContext *const context;
+  /// The configuration of the parser.
+  const AsmParserConfig &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;
 
   /// The depth of this parser in the nested parsing stack.
   size_t parserDepth;
 
-  /// An optional pointer to a struct containing high level parser state to be
-  /// populated during parsing.
-  AsmParserState *asmState;
-
   // 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/TypeParser.cpp

Type Parser::parseMemRefType() {
   auto parseElt = [&]() -> ParseResult {
     // Check for AffineMap as offset/strides.
     if (getToken().is(Token::kw_offset)) {
       int64_t offset;
       SmallVector<int64_t, 4> strides;
       if (failed(parseStridedLayout(offset, strides)))
         return failure();
       // Construct strided affine map.
-      AffineMap map =
-          makeStridedLinearLayoutMap(strides, offset, state.context);
+      AffineMap map = makeStridedLinearLayoutMap(strides, offset, getContext());
       layout = AffineMapAttr::get(map);
     } else {
       // Either it is MemRefLayoutAttrInterface or memory space attribute.
       Attribute attr = parseAttribute();
       if (!attr)
         return failure();
 
       if (attr.isa<MemRefLayoutAttrInterface>()) {

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

MLIRDocument::MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
     return;
   }
 
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
-  if (failed(parseSourceFile(sourceMgr, &parsedIR, &context, nullptr,
-                             &asmState))) {
+  if (failed(parseSourceFile(sourceMgr, &parsedIR,
+                             AsmParserConfig(&context, &asmState)))) {
     // If parsing failed, clear out any of the current state.
     parsedIR.clear();
     asmState = AsmParserState();
     return;
   }
 }
 
 //===----------------------------------------------------------------------===//