mlir/lib/Support/ToolUtilities.cpp

   1 //===- ToolUtilities.cpp - MLIR Tool Utilities ----------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file defines common utilities for implementing MLIR tools.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "mlir/Support/ToolUtilities.h"
  14 #include "mlir/Support/LLVM.h"
  15 #include "mlir/Support/LogicalResult.h"
  16 #include "llvm/Support/SourceMgr.h"
  17 #include "llvm/Support/raw_ostream.h"
  18
  19 using namespace mlir;
  20
  21 LogicalResult
  22 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
  23                             ChunkBufferHandler processChunkBuffer,
  24                             raw_ostream &os, bool enableSplitting,
  25                             bool insertMarkerInOutput) {
  26   // If splitting is disabled, we process the full input buffer.
  27   if (!enableSplitting)
  28     return processChunkBuffer(std::move(originalBuffer), os);
  29
  30   const char splitMarkerConst[] = "// -----";
  31   StringRef splitMarker(splitMarkerConst);
  32   const int splitMarkerLen = splitMarker.size();
  33
  34   auto *origMemBuffer = originalBuffer.get();
  35   SmallVector<StringRef, 8> rawSourceBuffers;
  36   const int checkLen = 2;
  37   // Split dropping the last checkLen chars to enable flagging near misses.
  38   origMemBuffer->getBuffer().split(rawSourceBuffers,
  39                                    splitMarker.drop_back(checkLen));
  40   if (rawSourceBuffers.empty())
  41     return success();
  42
  43   // Add the original buffer to the source manager.
  44   llvm::SourceMgr fileSourceMgr;
  45   fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), SMLoc());
  46
  47   // Flag near misses by iterating over all the sub-buffers found when splitting
  48   // with the prefix of the splitMarker. Use a sliding window where we only add
  49   // a buffer as a sourceBuffer if terminated by a full match of the
  50   // splitMarker, else flag a warning (if near miss) and extend the size of the
  51   // buffer under consideration.
  52   SmallVector<StringRef, 8> sourceBuffers;
  53   StringRef prev;
  54   for (auto buffer : rawSourceBuffers) {
  55     if (prev.empty()) {
  56       prev = buffer;
  57       continue;
  58     }
  59
  60     // Check that suffix is as expected and doesn't have any dash post.
  61     bool expectedSuffix = buffer.starts_with(splitMarker.take_back(checkLen)) &&
  62                           buffer.size() > checkLen && buffer[checkLen] != '0';
  63     if (expectedSuffix) {
  64       sourceBuffers.push_back(prev);
  65       prev = buffer.drop_front(checkLen);
  66     } else {
  67       // TODO: Consider making this a failure.
  68       auto splitLoc = SMLoc::getFromPointer(buffer.data());
  69       fileSourceMgr.PrintMessage(llvm::errs(), splitLoc,
  70                                  llvm::SourceMgr::DK_Warning,
  71                                  "near miss with file split marker");
  72       prev = StringRef(prev.data(),
  73                        prev.size() + splitMarkerLen - checkLen + buffer.size());
  74     }
  75   }
  76   if (!prev.empty())
  77     sourceBuffers.push_back(prev);
  78
  79   // Process each chunk in turn.
  80   bool hadFailure = false;
  81   auto interleaveFn = [&](StringRef subBuffer) {
  82     auto splitLoc = SMLoc::getFromPointer(subBuffer.data());
  83     unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
  84     auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy(
  85         subBuffer, Twine("within split at ") +
  86                        origMemBuffer->getBufferIdentifier() + ":" +
  87                        Twine(splitLine) + " offset ");
  88     if (failed(processChunkBuffer(std::move(subMemBuffer), os)))
  89       hadFailure = true;
  90   };
  91   llvm::interleave(sourceBuffers, os, interleaveFn,
  92                    insertMarkerInOutput ? "\n// -----\n" : "");
  93
  94   // If any fails, then return a failure of the tool.
  95   return failure(hadFailure);
  96 }