116 lines · cpp
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-exception6//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#include <string>18#include <utility>19 20using namespace mlir;21 22LogicalResult23mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,24 ChunkBufferHandler processChunkBuffer,25 raw_ostream &os, llvm::StringRef inputSplitMarker,26 llvm::StringRef outputSplitMarker) {27 llvm::MemoryBufferRef originalBufferRef = originalBuffer->getMemBufferRef();28 // If splitting is disabled, we process the full input buffer.29 if (inputSplitMarker.empty())30 return processChunkBuffer(std::move(originalBuffer), originalBufferRef, os);31 32 const int inputSplitMarkerLen = inputSplitMarker.size();33 34 SmallVector<StringRef, 8> rawSourceBuffers;35 const int checkLen = 2;36 // Split dropping the last checkLen chars to enable flagging near misses.37 originalBufferRef.getBuffer().split(rawSourceBuffers,38 inputSplitMarker.drop_back(checkLen));39 if (rawSourceBuffers.empty())40 return success();41 42 // Add the original buffer to the source manager.43 llvm::SourceMgr fileSourceMgr;44 fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), SMLoc());45 46 // Flag near misses by iterating over all the sub-buffers found when splitting47 // with the prefix of the splitMarker. Use a sliding window where we only add48 // a buffer as a sourceBuffer if terminated by a full match of the49 // splitMarker, else flag a warning (if near miss) and extend the size of the50 // buffer under consideration.51 SmallVector<StringRef, 8> sourceBuffers;52 StringRef prev;53 for (auto buffer : rawSourceBuffers) {54 if (prev.empty()) {55 prev = buffer;56 continue;57 }58 59 // Check that suffix is as expected and doesn't have any dash post.60 bool expectedSuffix =61 buffer.starts_with(inputSplitMarker.take_back(checkLen)) &&62 buffer.size() > checkLen && buffer[checkLen] != '0';63 if (expectedSuffix) {64 sourceBuffers.push_back(prev);65 prev = buffer.drop_front(checkLen);66 } else {67 // TODO: Consider making this a failure.68 auto splitLoc = SMLoc::getFromPointer(buffer.data());69 fileSourceMgr.PrintMessage(llvm::errs(), splitLoc,70 llvm::SourceMgr::DK_Warning,71 "near miss with file split marker");72 prev = StringRef(prev.data(), prev.size() + inputSplitMarkerLen -73 checkLen + buffer.size());74 }75 }76 if (!prev.empty())77 sourceBuffers.push_back(prev);78 79 // Process each chunk in turn.80 bool hadFailure = false;81 auto interleaveFn = [&](StringRef subBuffer) {82 auto splitLoc = SMLoc::getFromPointer(subBuffer.data());83 unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;84 std::string name((Twine("within split at ") +85 originalBufferRef.getBufferIdentifier() + ":" +86 Twine(splitLine) + " offset ")87 .str());88 // Use MemoryBufferRef to avoid copying the buffer & keep at same location89 // relative to the original buffer.90 auto subMemBuffer =91 llvm::MemoryBuffer::getMemBuffer(llvm::MemoryBufferRef(subBuffer, name),92 /*RequiresNullTerminator=*/false);93 if (failed(94 processChunkBuffer(std::move(subMemBuffer), originalBufferRef, os)))95 hadFailure = true;96 };97 llvm::interleave(sourceBuffers, os, interleaveFn,98 (llvm::Twine(outputSplitMarker) + "\n").str());99 100 // If any fails, then return a failure of the tool.101 return failure(hadFailure);102}103 104LogicalResult105mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,106 NoSourceChunkBufferHandler processChunkBuffer,107 raw_ostream &os, llvm::StringRef inputSplitMarker,108 llvm::StringRef outputSplitMarker) {109 auto process = [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,110 const llvm::MemoryBufferRef &, raw_ostream &os) {111 return processChunkBuffer(std::move(chunkBuffer), os);112 };113 return splitAndProcessBuffer(std::move(originalBuffer), process, os,114 inputSplitMarker, outputSplitMarker);115}116