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(A&: rawSourceBuffers, |
39 | Separator: splitMarker.drop_back(N: 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(F: std::move(originalBuffer), IncludeLoc: 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(Prefix: splitMarker.take_back(N: checkLen)) && |
62 | buffer.size() > checkLen && buffer[checkLen] != '0'; |
63 | if (expectedSuffix) { |
64 | sourceBuffers.push_back(Elt: prev); |
65 | prev = buffer.drop_front(N: checkLen); |
66 | } else { |
67 | // TODO: Consider making this a failure. |
68 | auto splitLoc = SMLoc::getFromPointer(Ptr: buffer.data()); |
69 | fileSourceMgr.PrintMessage(OS&: llvm::errs(), Loc: splitLoc, |
70 | Kind: llvm::SourceMgr::DK_Warning, |
71 | Msg: "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(Elt: prev); |
78 | |
79 | // Process each chunk in turn. |
80 | bool hadFailure = false; |
81 | auto interleaveFn = [&](StringRef subBuffer) { |
82 | auto splitLoc = SMLoc::getFromPointer(Ptr: subBuffer.data()); |
83 | unsigned splitLine = fileSourceMgr.getLineAndColumn(Loc: splitLoc).first; |
84 | auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy( |
85 | InputData: subBuffer, BufferName: Twine("within split at " ) + |
86 | origMemBuffer->getBufferIdentifier() + ":" + |
87 | Twine(splitLine) + " offset " ); |
88 | if (failed(result: processChunkBuffer(std::move(subMemBuffer), os))) |
89 | hadFailure = true; |
90 | }; |
91 | llvm::interleave(c: sourceBuffers, os, each_fn: interleaveFn, |
92 | separator: insertMarkerInOutput ? "\n// -----\n" : "" ); |
93 | |
94 | // If any fails, then return a failure of the tool. |
95 | return failure(isFailure: hadFailure); |
96 | } |
97 | |