[TableGen] Refactor TableGenParseFile to no longer use a callback

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.

Differential Revision: https://reviews.llvm.org/D125277

llvm/include/llvm/TableGen/Parser.h

@@ -18,21 +18,16 @@
 #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
 

llvm/lib/TableGen/Parser.cpp

@@ -9,29 +9,31 @@
 #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);
+  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());
+
   TGParser Parser(SrcMgr, /*Macros=*/None, Records);
-  if (Parser.ParseFile())
-    return true;
+  bool ParseResult = Parser.ParseFile();
 
-  // Invoke the provided handler function.
-  if (ParserFn(Records))
-    return true;
-
-  // After parsing, reset the tablegen data.
+  // After parsing, reclaim the source manager buffers from TableGen's global
+  // manager.
+  InputSrcMgr.takeSourceBuffersFrom(SrcMgr);
   SrcMgr = SourceMgr();
-  return false;
+  return ParseResult;
 }

llvm/unittests/TableGen/ParserEntryPointTest.cpp

@@ -9,6 +9,7 @@
 #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"
@@ -24,16 +25,16 @@ TEST(Parser, SanityTest) {
     }
   )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;
-  };
+  SourceMgr SrcMgr;
+  SrcMgr.AddNewSourceBuffer(
+      MemoryBuffer::getMemBuffer(SimpleTdSource, "test_buffer"), SMLoc());
 
-  bool ProcessResult = TableGenParseFile(
-      MemoryBuffer::getMemBuffer(SimpleTdSource, "test_buffer"),
-      /*IncludeDirs=*/{}, ProcessFn);
+  RecordKeeper Records;
+  bool ProcessResult = TableGenParseFile(SrcMgr, Records);
   EXPECT_FALSE(ProcessResult);
+
+  Record *Foo = Records.getDef("Foo");
+  Optional<StringRef> Field = Foo->getValueAsOptionalString("strField");
+  EXPECT_TRUE(Field.hasValue());
+  EXPECT_EQ(Field.getValue(), "value");
 }

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

@@ -722,6 +722,18 @@ 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 {
@@ -731,7 +743,7 @@ LogicalResult Parser::parseTdInclude(StringRef filename, llvm::SMRange fileLoc,
   } 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(
@@ -741,26 +753,18 @@ LogicalResult Parser::parseTdInclude(StringRef filename, llvm::SMRange fileLoc,
       },
       &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 + "`");
-
-  auto processFn = [&](llvm::RecordKeeper &records) {
-    processTdIncludeRecords(records, 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))
+  // Parse the tablegen file.
+  llvm::RecordKeeper tdRecords;
+  if (llvm::TableGenParseFile(tdSrcMgr, tdRecords))
     return failure();
 
+  // 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(tdSrcMgr, fileLoc.End);
   return success();
 }