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.startswith(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 }
97