Revert "[mlir] Report line number from file rather than split (llvm#150982)"

This reverts commit b6a98b9.

Author: hanhanW

mlir/include/mlir/IR/Diagnostics.h

@@ -639,10 +639,6 @@ class SourceMgrDiagnosticVerifierHandler : public SourceMgrDiagnosticHandler {
   /// verified correctly, failure otherwise.
   LogicalResult verify();
 
-  /// Register this handler with the given context. This is intended for use
-  /// with the splitAndProcessBuffer function.
-  void registerInContext(MLIRContext *ctx);
-
 private:
   /// Process a single diagnostic.
   void process(Diagnostic &diag);

mlir/include/mlir/Support/ToolUtilities.h

@@ -21,16 +21,10 @@
 
 namespace llvm {
 class MemoryBuffer;
-class MemoryBufferRef;
 } // namespace llvm
 
 namespace mlir {
-// A function that processes a chunk of a buffer and writes the result to an
-// output stream.
 using ChunkBufferHandler = function_ref<LogicalResult(
-    std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
-    const llvm::MemoryBufferRef &sourceBuffer, raw_ostream &os)>;
-using NoSourceChunkBufferHandler = function_ref<LogicalResult(
     std::unique_ptr<llvm::MemoryBuffer> chunkBuffer, raw_ostream &os)>;
 
 extern inline const char *const kDefaultSplitMarker = "// -----";
@@ -51,15 +45,6 @@ splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
                       ChunkBufferHandler processChunkBuffer, raw_ostream &os,
                       llvm::StringRef inputSplitMarker = kDefaultSplitMarker,
                       llvm::StringRef outputSplitMarker = "");
-
-/// Same as above, but for case where the original buffer is not used while
-/// processing the chunk.
-LogicalResult
-splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
-                      NoSourceChunkBufferHandler processChunkBuffer,
-                      raw_ostream &os,
-                      llvm::StringRef inputSplitMarker = kDefaultSplitMarker,
-                      llvm::StringRef outputSplitMarker = "");
 } // namespace mlir
 
 #endif // MLIR_SUPPORT_TOOLUTILITIES_H

mlir/lib/IR/Diagnostics.cpp

@@ -821,7 +821,15 @@ SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
   for (unsigned i = 0, e = mgr.getNumBuffers(); i != e; ++i)
     (void)impl->computeExpectedDiags(out, mgr, mgr.getMemoryBuffer(i + 1));
 
-  registerInContext(ctx);
+  // Register a handler to verify the diagnostics.
+  setHandler([&](Diagnostic &diag) {
+    // Process the main diagnostics.
+    process(diag);
+
+    // Process each of the notes.
+    for (auto &note : diag.getNotes())
+      process(note);
+  });
 }
 
 SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
@@ -854,17 +862,6 @@ LogicalResult SourceMgrDiagnosticVerifierHandler::verify() {
   return impl->status;
 }
 
-void SourceMgrDiagnosticVerifierHandler::registerInContext(MLIRContext *ctx) {
-  ctx->getDiagEngine().registerHandler([&](Diagnostic &diag) {
-    // Process the main diagnostics.
-    process(diag);
-
-    // Process each of the notes.
-    for (auto &note : diag.getNotes())
-      process(note);
-  });
-}
-
 /// Process a single diagnostic.
 void SourceMgrDiagnosticVerifierHandler::process(Diagnostic &diag) {
   return process(diag.getLocation(), diag.str(), diag.getSeverity());

mlir/lib/Support/ToolUtilities.cpp

@@ -14,8 +14,6 @@
 #include "mlir/Support/LLVM.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
-#include <string>
-#include <utility>
 
 using namespace mlir;
 
@@ -24,18 +22,18 @@ mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
                             ChunkBufferHandler processChunkBuffer,
                             raw_ostream &os, llvm::StringRef inputSplitMarker,
                             llvm::StringRef outputSplitMarker) {
-  llvm::MemoryBufferRef originalBufferRef = originalBuffer->getMemBufferRef();
   // If splitting is disabled, we process the full input buffer.
   if (inputSplitMarker.empty())
-    return processChunkBuffer(std::move(originalBuffer), originalBufferRef, os);
+    return processChunkBuffer(std::move(originalBuffer), os);
 
   const int inputSplitMarkerLen = inputSplitMarker.size();
 
+  auto *origMemBuffer = originalBuffer.get();
   SmallVector<StringRef, 8> rawSourceBuffers;
   const int checkLen = 2;
   // Split dropping the last checkLen chars to enable flagging near misses.
-  originalBufferRef.getBuffer().split(rawSourceBuffers,
-                                      inputSplitMarker.drop_back(checkLen));
+  origMemBuffer->getBuffer().split(rawSourceBuffers,
+                                   inputSplitMarker.drop_back(checkLen));
   if (rawSourceBuffers.empty())
     return success();
 
@@ -81,17 +79,11 @@ mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
   auto interleaveFn = [&](StringRef subBuffer) {
     auto splitLoc = SMLoc::getFromPointer(subBuffer.data());
     unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
-    std::string name((Twine("within split at ") +
-                      originalBufferRef.getBufferIdentifier() + ":" +
-                      Twine(splitLine) + " offset ")
-                         .str());
-    // Use MemoryBufferRef to avoid copying the buffer & keep at same location
-    // relative to the original buffer.
-    auto subMemBuffer =
-        llvm::MemoryBuffer::getMemBuffer(llvm::MemoryBufferRef(subBuffer, name),
-                                         /*RequiresNullTerminator=*/false);
-    if (failed(
-            processChunkBuffer(std::move(subMemBuffer), originalBufferRef, os)))
+    auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy(
+        subBuffer, Twine("within split at ") +
+                       origMemBuffer->getBufferIdentifier() + ":" +
+                       Twine(splitLine) + " offset ");
+    if (failed(processChunkBuffer(std::move(subMemBuffer), os)))
       hadFailure = true;
   };
   llvm::interleave(sourceBuffers, os, interleaveFn,
@@ -100,16 +92,3 @@ mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
   // If any fails, then return a failure of the tool.
   return failure(hadFailure);
 }
-
-LogicalResult
-mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
-                            NoSourceChunkBufferHandler processChunkBuffer,
-                            raw_ostream &os, llvm::StringRef inputSplitMarker,
-                            llvm::StringRef outputSplitMarker) {
-  auto process = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
-                     const llvm::MemoryBufferRef &, raw_ostream &os) {
-    return processChunkBuffer(std::move(chunkBuffer), os);
-  };
-  return splitAndProcessBuffer(std::move(originalBuffer), process, os,
-                               inputSplitMarker, outputSplitMarker);
-}

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

@@ -508,29 +508,20 @@ performActions(raw_ostream &os,
 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
-static LogicalResult
-processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
-              llvm::MemoryBufferRef sourceBuffer,
-              const MlirOptMainConfig &config, DialectRegistry &registry,
-              SourceMgrDiagnosticVerifierHandler *verifyHandler,
-              llvm::ThreadPoolInterface *threadPool) {
+static LogicalResult processBuffer(raw_ostream &os,
+                                   std::unique_ptr<MemoryBuffer> ownedBuffer,
+                                   const MlirOptMainConfig &config,
+                                   DialectRegistry &registry,
+                                   llvm::ThreadPoolInterface *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   auto sourceMgr = std::make_shared<SourceMgr>();
-  // Add the original buffer to the source manager to use for determining
-  // locations.
-  sourceMgr->AddNewSourceBuffer(
-      llvm::MemoryBuffer::getMemBuffer(sourceBuffer,
-                                       /*RequiresNullTerminator=*/false),
-      SMLoc());
   sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   if (threadPool)
     context.setThreadPool(*threadPool);
-  if (verifyHandler)
-    verifyHandler->registerInContext(&context);
 
   StringRef irdlFile = config.getIrdlFile();
   if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, context)))
@@ -554,12 +545,17 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
     return performActions(os, sourceMgr, &context, config);
   }
 
+  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(
+      *sourceMgr, &context, config.verifyDiagnosticsLevel());
+
   // Do any processing requested by command line flags.  We don't care whether
   // these actions succeed or fail, we only care what diagnostics they produce
   // and whether they match our expectations.
   (void)performActions(os, sourceMgr, &context, config);
 
-  return success();
+  // Verify the diagnostic handler to make sure that each of the diagnostics
+  // matched.
+  return sourceMgrHandler.verify();
 }
 
 std::pair<std::string, std::string>
@@ -628,31 +624,14 @@ LogicalResult mlir::MlirOptMain(llvm::raw_ostream &outputStream,
   if (threadPoolCtx.isMultithreadingEnabled())
     threadPool = &threadPoolCtx.getThreadPool();
 
-  SourceMgr sourceMgr;
-  sourceMgr.AddNewSourceBuffer(
-      llvm::MemoryBuffer::getMemBuffer(buffer->getMemBufferRef(),
-                                       /*RequiresNullTerminator=*/false),
-      SMLoc());
-  // Note: this creates a verifier handler independent of the the flag set, as
-  // internally if the flag is not set, a new scoped diagnostic handler is
-  // created which would intercept the diagnostics and verify them.
-  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(
-      sourceMgr, &threadPoolCtx, config.verifyDiagnosticsLevel());
   auto chunkFn = [&](std::unique_ptr<MemoryBuffer> chunkBuffer,
-                     llvm::MemoryBufferRef sourceBuffer, raw_ostream &os) {
-    return processBuffer(
-        os, std::move(chunkBuffer), sourceBuffer, config, registry,
-        config.shouldVerifyDiagnostics() ? &sourceMgrHandler : nullptr,
-        threadPool);
+                     raw_ostream &os) {
+    return processBuffer(os, std::move(chunkBuffer), config, registry,
+                         threadPool);
   };
-  LogicalResult status = splitAndProcessBuffer(
-      llvm::MemoryBuffer::getMemBuffer(buffer->getMemBufferRef(),
-                                       /*RequiresNullTerminator=*/false),
-      chunkFn, outputStream, config.inputSplitMarker(),
-      config.outputSplitMarker());
-  if (config.shouldVerifyDiagnostics() && failed(sourceMgrHandler.verify()))
-    status = failure();
-  return status;
+  return splitAndProcessBuffer(std::move(buffer), chunkFn, outputStream,
+                               config.inputSplitMarker(),
+                               config.outputSplitMarker());
 }
 
 LogicalResult mlir::MlirOptMain(int argc, char **argv,

mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp

@@ -135,13 +135,6 @@ LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
   // Processes the memory buffer with a new MLIRContext.
   auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
                            raw_ostream &os) {
-    // Many of the translations expect a null-terminated buffer while splitting
-    // the buffer does not guarantee null-termination. Make a copy of the buffer
-    // to ensure null-termination.
-    if (!ownedBuffer->getBuffer().ends_with('\0')) {
-      ownedBuffer = llvm::MemoryBuffer::getMemBufferCopy(
-          ownedBuffer->getBuffer(), ownedBuffer->getBufferIdentifier());
-    }
     // Temporary buffers for chained translation processing.
     std::string dataIn;
     std::string dataOut;

mlir/test/IR/diagnostic-nosplit.mlir

@@ -6,10 +6,10 @@ func.func private @foo()
 
 // -----
 
-// expected-error@-9 {{source must contain a single top-level operation, found: 2}}
+// expected-error@-3 {{source must contain a single top-level operation, found: 2}}
 func.func private @bar()
 func.func private @baz()
 
 // -----
 
-// expected-error@-15 {{source must contain a single top-level operation, found: 0}}
+// expected-error@-3 {{source must contain a single top-level operation, found: 0}}

mlir/test/IR/top-level.mlir

@@ -201,12 +201,6 @@ int main(int argc, char **argv) {
   llvm::raw_string_ostream outputStrOS(outputStr);
   auto processFn = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
                        raw_ostream &os) {
-    // Split does not guarantee null-termination. Make a copy of the buffer to
-    // ensure null-termination.
-    if (!chunkBuffer->getBuffer().ends_with('\0')) {
-      chunkBuffer = llvm::MemoryBuffer::getMemBufferCopy(
-          chunkBuffer->getBuffer(), chunkBuffer->getBufferIdentifier());
-    }
     return processBuffer(os, std::move(chunkBuffer), outputType, includeDirs,
                          dumpODS, includedFiles);
   };