69 lines · cpp
1//===- MatchFinder.cpp - --------------------------------------------------===//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 contains the method definitions for the `MatchFinder` class10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Query/Matcher/MatchFinder.h"14namespace mlir::query::matcher {15 16MatchFinder::MatchResult::MatchResult(Operation *rootOp,17 std::vector<Operation *> matchedOps)18 : rootOp(rootOp), matchedOps(std::move(matchedOps)) {}19 20std::vector<MatchFinder::MatchResult>21MatchFinder::collectMatches(Operation *root, DynMatcher matcher) const {22 std::vector<MatchResult> results;23 llvm::SetVector<Operation *> tempStorage;24 root->walk([&](Operation *subOp) {25 if (matcher.match(subOp)) {26 MatchResult match;27 match.rootOp = subOp;28 match.matchedOps.push_back(subOp);29 results.push_back(std::move(match));30 } else if (matcher.match(subOp, tempStorage)) {31 results.emplace_back(subOp, std::vector<Operation *>(tempStorage.begin(),32 tempStorage.end()));33 }34 tempStorage.clear();35 });36 return results;37}38 39void MatchFinder::printMatch(llvm::raw_ostream &os, QuerySession &qs,40 Operation *op) const {41 auto fileLoc = cast<FileLineColLoc>(op->getLoc());42 SMLoc smloc = qs.getSourceManager().FindLocForLineAndColumn(43 qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());44 llvm::SMDiagnostic diag =45 qs.getSourceManager().GetMessage(smloc, llvm::SourceMgr::DK_Note, "");46 diag.print("", os, true, false, true);47}48 49void MatchFinder::printMatch(llvm::raw_ostream &os, QuerySession &qs,50 Operation *op, const std::string &binding) const {51 auto fileLoc = cast<FileLineColLoc>(op->getLoc());52 auto smloc = qs.getSourceManager().FindLocForLineAndColumn(53 qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());54 qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,55 "\"" + binding + "\" binds here");56}57 58std::vector<Operation *>59MatchFinder::flattenMatchedOps(std::vector<MatchResult> &matches) const {60 std::vector<Operation *> newVector;61 for (auto &result : matches) {62 newVector.insert(newVector.end(), result.matchedOps.begin(),63 result.matchedOps.end());64 }65 return newVector;66}67 68} // namespace mlir::query::matcher69