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
splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,ChunkBufferHandler processChunkBuffer,raw_ostream & os,bool enableSplitting,bool insertMarkerInOutput)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