Add support to the SourceMgrDiagnosticHandlers to auto-load new source files...

--

PiperOrigin-RevId: 247681779

diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index b456bc1..9ab29c7 100644 (file)
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -448,9 +448,8 @@ public:
   void emitDiagnostic(Location loc, Twine message, DiagnosticSeverity kind);
 
 protected:
-  /// Get a memory buffer for the given file, or the main file of the source
-  /// manager if one doesn't exist. If no file is available, nullptr is
-  /// returned.
+  /// Get a memory buffer for the given file, or nullptr if no file is
+  /// available.
   const llvm::MemoryBuffer *getBufferForFile(StringRef filename);
 
   /// The source manager that we are wrapping.

diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index c06539d..10eaed8 100644 (file)
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -313,10 +313,11 @@ SourceMgrDiagnosticHandler::getBufferForFile(StringRef filename) {
       return filenameToBuf[filename] = buf;
   }
 
-  // Otherwise, default to the main file buffer, if it exists.
+  // Otherwise, try to load the source file.
   const llvm::MemoryBuffer *newBuf = nullptr;
-  if (mgr.getNumBuffers() != 0)
-    newBuf = mgr.getMemoryBuffer(mgr.getMainFileID());
+  std::string ignored;
+  if (auto newBufID = mgr.AddIncludeFile(filename, llvm::SMLoc(), ignored))
+    newBuf = mgr.getMemoryBuffer(newBufID);
   return filenameToBuf[filename] = newBuf;
 }