diff --git a/llvm/include/llvm/TableGen/Parser.h b/llvm/include/llvm/TableGen/Parser.h --- a/llvm/include/llvm/TableGen/Parser.h +++ b/llvm/include/llvm/TableGen/Parser.h @@ -18,21 +18,16 @@ #include 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; - -/// 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 Buffer, - std::vector 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 diff --git a/llvm/lib/TableGen/Parser.cpp b/llvm/lib/TableGen/Parser.cpp --- a/llvm/lib/TableGen/Parser.cpp +++ b/llvm/lib/TableGen/Parser.cpp @@ -14,24 +14,28 @@ using namespace llvm; -bool llvm::TableGenParseFile(std::unique_ptr Buffer, - std::vector 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()); - // Invoke the provided handler function. - if (ParserFn(Records)) - return true; + // Setup the record keeper and try to parse the file. + auto *MainFileBuffer = SrcMgr.getMemoryBuffer(SrcMgr.getMainFileID()); + Records.saveInputFilename(MainFileBuffer->getBufferIdentifier().str()); - // After parsing, reset the tablegen data. + TGParser Parser(SrcMgr, /*Macros=*/None, Records); + bool ParseResult = Parser.ParseFile(); + + // 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; } diff --git a/llvm/unittests/TableGen/ParserEntryPointTest.cpp b/llvm/unittests/TableGen/ParserEntryPointTest.cpp --- a/llvm/unittests/TableGen/ParserEntryPointTest.cpp +++ b/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 @@ } )td"; - auto ProcessFn = [&](const RecordKeeper &Records) { - Record *Foo = Records.getDef("Foo"); - Optional 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 Field = Foo->getValueAsOptionalString("strField"); + EXPECT_TRUE(Field.hasValue()); + EXPECT_EQ(Field.getValue(), "value"); } diff --git a/mlir/lib/Tools/PDLL/Parser/Parser.cpp b/mlir/lib/Tools/PDLL/Parser/Parser.cpp --- a/mlir/lib/Tools/PDLL/Parser/Parser.cpp +++ b/mlir/lib/Tools/PDLL/Parser/Parser.cpp @@ -722,6 +722,18 @@ SmallVectorImpl &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> 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 @@ } 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(rawHandlerContext); (void)ctx->parser.emitError( @@ -741,26 +753,18 @@ }, &handlerContext); - // Use the source manager to open the file, but don't yet add it. - std::string includedFile; - llvm::ErrorOr> 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(); }