[TableGen] Refactor TableGenParseFile to no longer use a callback

[rriddle]

Summary

Now that TableGen no longer relies on global Record state, we can allow
for the client to own the RecordKeeper and SourceMgr. Given that TableGen
internally still relies on the global llvm::SrcMgr, this method unfortunately
still isn't thread-safe.

[lattner]

nice, this is a more logical API

[nhaehnle]

Thank you for doing this! Not a detailed review but the interface and approach LGTM.

llvm/include/llvm/TableGen/Parser.h

 #ifndef LLVM_TABLEGEN_PARSER_H
 #define LLVM_TABLEGEN_PARSER_H
 
 #include "llvm/ADT/STLExtras.h"
 #include <string>
 #include <vector>
 
 namespace llvm {
-class MemoryBuffer;
 class RecordKeeper;
+class SourceMgr;
 
-/// Peform the tablegen action using the given set of parsed records. Returns
-/// true on error, false otherwise.
-using TableGenParserFn = function_ref<bool(RecordKeeper &)>;
-
-/// Parse the given input buffer containing a tablegen file, invoking the
-/// provided parser function with the set of parsed records. All tablegen state
-/// is reset after the provided parser function is invoked, i.e., the provided
-/// parser function should not maintain references to any tablegen constructs
-/// after executing. Returns true on failure, false otherwise.
-bool TableGenParseFile(std::unique_ptr<MemoryBuffer> Buffer,
-                       std::vector<std::string> IncludeDirs,
-                       TableGenParserFn ParserFn);
+/// Parse the TableGen file defined within the main buffer of the given
+/// SourceMgr. On success, populates the provided RecordKeeper with the parsed
+/// records and returns false. On failure, returns true.
+///
+/// NOTE: TableGen currently relies on global state within a given parser
+///       invocation, so this function is not thread-safe.
+bool TableGenParseFile(SourceMgr &InputSrcMgr, RecordKeeper &Records);
 
 } // end namespace llvm
 
 #endif // LLVM_TABLEGEN_PARSER_H

llvm/lib/TableGen/Parser.cpp

 //===- Parser.cpp - Top-Level TableGen Parser implementation --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/TableGen/Parser.h"
 #include "TGParser.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 
 using namespace llvm;
 
-bool llvm::TableGenParseFile(std::unique_ptr<MemoryBuffer> Buffer,
-                             std::vector<std::string> IncludeDirs,
-                             TableGenParserFn ParserFn) {
-  RecordKeeper Records;
-  Records.saveInputFilename(Buffer->getBufferIdentifier().str());
-
+bool llvm::TableGenParseFile(SourceMgr &InputSrcMgr, RecordKeeper &Records) {
+  // Initialize the global TableGen source manager by temporarily taking control
+  // of the input buffer in `SrcMgr`. This is kind of a hack, but allows for
+  // preserving TableGen's current awkward diagnostic behavior. If we can remove
+  // this reliance, we could drop all of this.
   SrcMgr = SourceMgr();
-  SrcMgr.AddNewSourceBuffer(std::move(Buffer), SMLoc());
-  SrcMgr.setIncludeDirs(IncludeDirs);
-  TGParser Parser(SrcMgr, /*Macros=*/None, Records);
-  if (Parser.ParseFile())
-    return true;
+  SrcMgr.takeSourceBuffersFrom(InputSrcMgr);
+  SrcMgr.setIncludeDirs(InputSrcMgr.getIncludeDirs());
+  SrcMgr.setDiagHandler(InputSrcMgr.getDiagHandler(),
+                        InputSrcMgr.getDiagContext());
+
+  // Setup the record keeper and try to parse the file.
+  auto *MainFileBuffer = SrcMgr.getMemoryBuffer(SrcMgr.getMainFileID());
+  Records.saveInputFilename(MainFileBuffer->getBufferIdentifier().str());
 
-  // Invoke the provided handler function.
-  if (ParserFn(Records))
-    return true;
+  TGParser Parser(SrcMgr, /*Macros=*/None, Records);
+  bool ParseResult = Parser.ParseFile();
 
-  // After parsing, reset the tablegen data.
+  // After parsing, reclaim the source managers from TableGens global manager.
+  InputSrcMgr.takeSourceBuffersFrom(SrcMgr);
   SrcMgr = SourceMgr();
-  return false;
+
+  // If we succeeded, return the record keeper. Otherwise, return None.
+  return ParseResult;
 }

llvm/unittests/TableGen/ParserEntryPointTest.cpp

 //===- unittest/TableGen/ParserEntryPointTest.cpp - Parser tests ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/TableGen/Parser.h"
 #include "llvm/TableGen/Record.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 TEST(Parser, SanityTest) {
   // Simple TableGen source file with a single record.
   const char *SimpleTdSource = R"td(
     def Foo {
       string strField = "value";
     }
   )td";
 
-  auto ProcessFn = [&](const RecordKeeper &Records) {
-    Record *Foo = Records.getDef("Foo");
-    Optional<StringRef> Field = Foo->getValueAsOptionalString("strField");
-    EXPECT_TRUE(Field.hasValue());
-    EXPECT_EQ(Field.getValue(), "value");
-    return false;
-  };
-
-  bool ProcessResult = TableGenParseFile(
-      MemoryBuffer::getMemBuffer(SimpleTdSource, "test_buffer"),
-      /*IncludeDirs=*/{}, ProcessFn);
-  EXPECT_FALSE(ProcessResult);
+  SourceMgr SrcMgr;
+  SrcMgr.AddNewSourceBuffer(
+      MemoryBuffer::getMemBuffer(SimpleTdSource, "test_buffer"), SMLoc());
+
+  RecordKeeper Records;
+  bool ProcessResult = TableGenParseFile(SrcMgr, Records);
+  EXPECT_FALSE(ProcessResult);
+
 }

mlir/lib/Tools/PDLL/Parser/Parser.cpp

LogicalResult Parser::parseInclude(SmallVectorImpl<ast::Decl *> &decls) {
   return emitError(fileLoc,
                    "expected include filename to end with `.pdll` or `.td`");
 }
 
 LogicalResult Parser::parseTdInclude(StringRef filename, llvm::SMRange fileLoc,
                                      SmallVectorImpl<ast::Decl *> &decls) {
   llvm::SourceMgr &parserSrcMgr = lexer.getSourceMgr();
 
+  // Use the source manager to open the file, but don't yet add it.
+  std::string includedFile;
+  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> includeBuffer =
+      parserSrcMgr.OpenIncludeFile(filename.str(), includedFile);
+  if (!includeBuffer)
+    return emitError(fileLoc, "unable to open include file `" + filename + "`");
+
+  // Setup the source manager for parsing the tablegen file.
+  llvm::SourceMgr tdSrcMgr;
+  tdSrcMgr.AddNewSourceBuffer(std::move(*includeBuffer), SMLoc());
+  tdSrcMgr.setIncludeDirs(parserSrcMgr.getIncludeDirs());
+
   // This class provides a context argument for the llvm::SourceMgr diagnostic
   // handler.
   struct DiagHandlerContext {
     Parser &parser;
     StringRef filename;
     llvm::SMRange loc;
   } handlerContext{*this, filename, fileLoc};
 
   // Set the diagnostic handler for the tablegen source manager.
-  llvm::SrcMgr.setDiagHandler(
+  tdSrcMgr.setDiagHandler(
       [](const llvm::SMDiagnostic &diag, void *rawHandlerContext) {
         auto *ctx = reinterpret_cast<DiagHandlerContext *>(rawHandlerContext);
         (void)ctx->parser.emitError(
             ctx->loc,
             llvm::formatv("error while processing include file `{0}`: {1}",
                           ctx->filename, diag.getMessage()));
       },
       &handlerContext);
 
-  // Use the source manager to open the file, but don't yet add it.
-  std::string includedFile;
-  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> includeBuffer =
-      parserSrcMgr.OpenIncludeFile(filename.str(), includedFile);
-  if (!includeBuffer)
-    return emitError(fileLoc, "unable to open include file `" + filename + "`");
+  // Parse the tablegen file.
+  llvm::RecordKeeper tdRecords;
+  if (llvm::TableGenParseFile(tdSrcMgr, tdRecords))
+    return failure();
 
-  auto processFn = [&](llvm::RecordKeeper &records) {
-    processTdIncludeRecords(records, decls);
+  // Process the parsed records.
+  processTdIncludeRecords(tdRecords, decls);
 
-    // After we are done processing, move all of the tablegen source buffers to
-    // the main parser source mgr. This allows for directly using source
-    // locations from the .td files without needing to remap them.
-    parserSrcMgr.takeSourceBuffersFrom(llvm::SrcMgr, fileLoc.End);
-    return false;
-  };
-  if (llvm::TableGenParseFile(std::move(*includeBuffer),
-                              parserSrcMgr.getIncludeDirs(), processFn))
-    return failure();
-
+  // the main parser source mgr. This allows for directly using source locations
+  // from the .td files without needing to remap them.
+  parserSrcMgr.takeSourceBuffersFrom(tdSrcMgr, fileLoc.End);
   return success();
 }
 
 void Parser::processTdIncludeRecords(llvm::RecordKeeper &tdRecords,
                                      SmallVectorImpl<ast::Decl *> &decls) {
   // Return the length kind of the given value.
   auto getLengthKind = [](const auto &value) {
     if (value.isOptional())