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