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

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