Add support in the SourceMgrDiagnosticHandler for handling source managers with...

--

PiperOrigin-RevId: 247482733

diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index 9156e54..e5d2207 100644 (file)
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -26,10 +26,12 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/Twine.h"
 #include <functional>
 
 namespace llvm {
+class MemoryBuffer;
 class SMLoc;
 class SourceMgr;
 } // end namespace llvm
@@ -441,11 +443,18 @@ public:
   void emitDiagnostic(Location loc, Twine message, DiagnosticSeverity kind);
 
 private:
+  /// Get a memory buffer for the given file, or the main file of the source
+  /// manager if one doesn't exist.
+  const llvm::MemoryBuffer &getBufferForFile(StringRef filename);
+
   /// Convert a location into the given memory buffer into an SMLoc.
   llvm::SMLoc convertLocToSMLoc(Location loc);
 
   /// The source manager that we are wrapping.
   llvm::SourceMgr &mgr;
+
+  /// Mapping between file name and buffer pointer.
+  llvm::StringMap<const llvm::MemoryBuffer *> filenameToBuf;
 };
 } // namespace mlir
 

diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index a265171..a6ae074 100644 (file)
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -275,7 +275,26 @@ void SourceMgrDiagnosticHandler::emitDiagnostic(Location loc, Twine message,
   mgr.PrintMessage(convertLocToSMLoc(loc), getDiagKind(kind), message);
 }
 
-/// Convert a location into the given memory buffer into an SMLoc.
+/// Get a memory buffer for the given file, or nullptr if one is not found.
+const llvm::MemoryBuffer &
+SourceMgrDiagnosticHandler::getBufferForFile(StringRef filename) {
+  // Check for an existing mapping to the buffer id for this file.
+  auto bufferIt = filenameToBuf.find(filename);
+  if (bufferIt != filenameToBuf.end())
+    return *bufferIt->second;
+
+  // Look for a buffer in the manager that has this filename.
+  for (unsigned i = 1, e = mgr.getNumBuffers() + 1; i != e; ++i) {
+    auto *buf = mgr.getMemoryBuffer(i);
+    if (buf->getBufferIdentifier() == filename)
+      return *(filenameToBuf[filename] = buf);
+  }
+
+  return *(filenameToBuf[filename] = mgr.getMemoryBuffer(mgr.getMainFileID()));
+}
+
+/// Get a memory buffer for the given file, or the main file of the source
+/// manager if one doesn't exist. This always returns non-null.
 llvm::SMLoc SourceMgrDiagnosticHandler::convertLocToSMLoc(Location loc) {
   auto fileLoc = loc.dyn_cast<FileLineColLoc>();
 
@@ -283,7 +302,8 @@ llvm::SMLoc SourceMgrDiagnosticHandler::convertLocToSMLoc(Location loc) {
   if (!fileLoc)
     return llvm::SMLoc();
 
-  auto *membuf = mgr.getMemoryBuffer(mgr.getMainFileID());
+  // Get the buffer for this filename.
+  auto &membuf = getBufferForFile(fileLoc->getFilename());
 
   // TODO: This should really be upstreamed to be a method on llvm::SourceMgr.
   // Doing so would allow it to use the offset cache that is already maintained
@@ -292,8 +312,8 @@ llvm::SMLoc SourceMgrDiagnosticHandler::convertLocToSMLoc(Location loc) {
   unsigned columnNo = fileLoc->getColumn();
 
   // Scan for the correct line number.
-  const char *position = membuf->getBufferStart();
-  const char *end = membuf->getBufferEnd();
+  const char *position = membuf.getBufferStart();
+  const char *end = membuf.getBufferEnd();
 
   // We start counting line and column numbers from 1.
   --lineNo;
@@ -319,7 +339,7 @@ llvm::SMLoc SourceMgrDiagnosticHandler::convertLocToSMLoc(Location loc) {
   // If the line/column counter was invalid, return a pointer to the start of
   // the buffer.
   if (lineNo || position + columnNo > end)
-    return llvm::SMLoc::getFromPointer(membuf->getBufferStart());
+    return llvm::SMLoc::getFromPointer(membuf.getBufferStart());
 
   // Otherwise return the right pointer.
   return llvm::SMLoc::getFromPointer(position + columnNo);