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

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