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