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
19using namespace mlir;
20
21LogicalResult
22mlir::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

source code of mlir/lib/Support/ToolUtilities.cpp