mlir/lib/Query/QueryParser.cpp

   1 //===---- QueryParser.cpp - mlir-query command parser ---------------------===//
   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 #include "QueryParser.h"
  10 #include "llvm/ADT/StringSwitch.h"
  11
  12 namespace mlir::query {
  13
  14 // Lex any amount of whitespace followed by a "word" (any sequence of
  15 // non-whitespace characters) from the start of region [begin,end).  If no word
  16 // is found before end, return StringRef(). begin is adjusted to exclude the
  17 // lexed region.
  18 llvm::StringRef QueryParser::lexWord() {
  19   // Don't trim newlines.
  20   line = line.ltrim(" \t\v\f\r");
  21
  22   if (line.empty())
  23     // Even though the line is empty, it contains a pointer and
  24     // a (zero) length. The pointer is used in the LexOrCompleteWord
  25     // code completion.
  26     return line;
  27
  28   llvm::StringRef word;
  29   if (line.front() == '#') {
  30     word = line.substr(0, 1);
  31   } else {
  32     word = line.take_until([](char c) {
  33       // Don't trim newlines.
  34       return llvm::StringRef(" \t\v\f\r").contains(c);
  35     });
  36   }
  37
  38   line = line.drop_front(word.size());
  39   return word;
  40 }
  41
  42 // This is the StringSwitch-alike used by LexOrCompleteWord below. See that
  43 // function for details.
  44 template <typename T>
  45 struct QueryParser::LexOrCompleteWord {
  46   llvm::StringRef word;
  47   llvm::StringSwitch<T> stringSwitch;
  48
  49   QueryParser *queryParser;
  50   // Set to the completion point offset in word, or StringRef::npos if
  51   // completion point not in word.
  52   size_t wordCompletionPos;
  53
  54   // Lexes a word and stores it in word. Returns a LexOrCompleteword<T> object
  55   // that can be used like a llvm::StringSwitch<T>, but adds cases as possible
  56   // completions if the lexed word contains the completion point.
  57   LexOrCompleteWord(QueryParser *queryParser, llvm::StringRef &outWord)
  58       : word(queryParser->lexWord()), stringSwitch(word),
  59         queryParser(queryParser), wordCompletionPos(llvm::StringRef::npos) {
  60     outWord = word;
  61     if (queryParser->completionPos &&
  62         queryParser->completionPos <= word.data() + word.size()) {
  63       if (queryParser->completionPos < word.data())
  64         wordCompletionPos = 0;
  65       else
  66         wordCompletionPos = queryParser->completionPos - word.data();
  67     }
  68   }
  69
  70   LexOrCompleteWord &Case(llvm::StringLiteral caseStr, const T &value,
  71                           bool isCompletion = true) {
  72
  73     if (wordCompletionPos == llvm::StringRef::npos)
  74       stringSwitch.Case(caseStr, value);
  75     else if (!caseStr.empty() && isCompletion &&
  76              wordCompletionPos <= caseStr.size() &&
  77              caseStr.substr(0, wordCompletionPos) ==
  78                  word.substr(0, wordCompletionPos)) {
  79
  80       queryParser->completions.emplace_back(
  81           (caseStr.substr(wordCompletionPos) + " ").str(),
  82           std::string(caseStr));
  83     }
  84     return *this;
  85   }
  86
  87   T Default(T value) { return stringSwitch.Default(value); }
  88 };
  89
  90 QueryRef QueryParser::endQuery(QueryRef queryRef) {
  91   llvm::StringRef extra = line;
  92   llvm::StringRef extraTrimmed = extra.ltrim(" \t\v\f\r");
  93
  94   if (extraTrimmed.starts_with('\n') || extraTrimmed.starts_with("\r\n"))
  95     queryRef->remainingContent = extra;
  96   else {
  97     llvm::StringRef trailingWord = lexWord();
  98     if (trailingWord.starts_with('#')) {
  99       line = line.drop_until([](char c) { return c == '\n'; });
 100       line = line.drop_while([](char c) { return c == '\n'; });
 101       return endQuery(queryRef);
 102     }
 103     if (!trailingWord.empty()) {
 104       return new InvalidQuery("unexpected extra input: '" + extra + "'");
 105     }
 106   }
 107   return queryRef;
 108 }
 109
 110 namespace {
 111
 112 enum class ParsedQueryKind {
 113   Invalid,
 114   Comment,
 115   NoOp,
 116   Help,
 117   Match,
 118   Quit,
 119 };
 120
 121 QueryRef
 122 makeInvalidQueryFromDiagnostics(const matcher::internal::Diagnostics &diag) {
 123   std::string errStr;
 124   llvm::raw_string_ostream os(errStr);
 125   diag.print(os);
 126   return new InvalidQuery(errStr);
 127 }
 128 } // namespace
 129
 130 QueryRef QueryParser::completeMatcherExpression() {
 131   std::vector<matcher::MatcherCompletion> comps =
 132       matcher::internal::Parser::completeExpression(
 133           line, completionPos - line.begin(), qs.getRegistryData(),
 134           &qs.namedValues);
 135   for (const auto &comp : comps) {
 136     completions.emplace_back(comp.typedText, comp.matcherDecl);
 137   }
 138   return QueryRef();
 139 }
 140
 141 QueryRef QueryParser::doParse() {
 142
 143   llvm::StringRef commandStr;
 144   ParsedQueryKind qKind =
 145       LexOrCompleteWord<ParsedQueryKind>(this, commandStr)
 146           .Case("", ParsedQueryKind::NoOp)
 147           .Case("#", ParsedQueryKind::Comment, /*isCompletion=*/false)
 148           .Case("help", ParsedQueryKind::Help)
 149           .Case("m", ParsedQueryKind::Match, /*isCompletion=*/false)
 150           .Case("match", ParsedQueryKind::Match)
 151           .Case("q", ParsedQueryKind::Quit, /*IsCompletion=*/false)
 152           .Case("quit", ParsedQueryKind::Quit)
 153           .Default(ParsedQueryKind::Invalid);
 154
 155   switch (qKind) {
 156   case ParsedQueryKind::Comment:
 157   case ParsedQueryKind::NoOp:
 158     line = line.drop_until([](char c) { return c == '\n'; });
 159     line = line.drop_while([](char c) { return c == '\n'; });
 160     if (line.empty())
 161       return new NoOpQuery;
 162     return doParse();
 163
 164   case ParsedQueryKind::Help:
 165     return endQuery(new HelpQuery);
 166
 167   case ParsedQueryKind::Quit:
 168     return endQuery(new QuitQuery);
 169
 170   case ParsedQueryKind::Match: {
 171     if (completionPos) {
 172       return completeMatcherExpression();
 173     }
 174
 175     matcher::internal::Diagnostics diag;
 176     auto matcherSource = line.ltrim();
 177     auto origMatcherSource = matcherSource;
 178     std::optional<matcher::DynMatcher> matcher =
 179         matcher::internal::Parser::parseMatcherExpression(
 180             matcherSource, qs.getRegistryData(), &qs.namedValues, &diag);
 181     if (!matcher) {
 182       return makeInvalidQueryFromDiagnostics(diag);
 183     }
 184     auto actualSource = origMatcherSource.substr(0, origMatcherSource.size() -
 185                                                         matcherSource.size());
 186     QueryRef query = new MatchQuery(actualSource, *matcher);
 187     query->remainingContent = matcherSource;
 188     return query;
 189   }
 190
 191   case ParsedQueryKind::Invalid:
 192     return new InvalidQuery("unknown command: " + commandStr);
 193   }
 194
 195   llvm_unreachable("Invalid query kind");
 196 }
 197
 198 QueryRef QueryParser::parse(llvm::StringRef line, const QuerySession &qs) {
 199   return QueryParser(line, qs).doParse();
 200 }
 201
 202 std::vector<llvm::LineEditor::Completion>
 203 QueryParser::complete(llvm::StringRef line, size_t pos,
 204                       const QuerySession &qs) {
 205   QueryParser queryParser(line, qs);
 206   queryParser.completionPos = line.data() + pos;
 207
 208   queryParser.doParse();
 209   return queryParser.completions;
 210 }
 211
 212 } // namespace mlir::query