4 files changed, 68 insertions, 19 deletions

diff --git a/clang-tools-extra/clang-query/Query.cpp b/clang-tools-extra/clang-query/Query.cpp
index 675fd968f46..8eafc5eed75 100644
--- a/clang-tools-extra/clang-query/Query.cpp
+++ b/clang-tools-extra/clang-query/Query.cpp
@@ -101,9 +101,24 @@ bool MatchQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
     Finder.matchAST(AST->getASTContext());
 
     if (QS.PrintMatcher) {
-      std::string prefixText = "Matcher: ";
-      OS << "\n  " << prefixText << Source << "\n";
-      OS << "  " << std::string(prefixText.size() + Source.size(), '=') << '\n';
+      SmallVector<StringRef, 4> Lines;
+      Source.split(Lines, "\n");
+      auto FirstLine = Lines[0];
+      Lines.erase(Lines.begin(), Lines.begin() + 1);
+      while (!Lines.empty() && Lines.back().empty()) {
+        Lines.resize(Lines.size() - 1);
+      }
+      unsigned MaxLength = FirstLine.size();
+      std::string PrefixText = "Matcher: ";
+      OS << "\n  " << PrefixText << FirstLine;
+
+      for (auto Line : Lines) {
+        OS << "\n" << std::string(PrefixText.size() + 2, ' ') << Line;
+        MaxLength = std::max<int>(MaxLength, Line.rtrim().size());
+      }
+
+      OS << "\n"
+         << "  " << std::string(PrefixText.size() + MaxLength, '=') << "\n\n";
     }
 
     for (auto MI = Matches.begin(), ME = Matches.end(); MI != ME; ++MI) {
diff --git a/clang-tools-extra/clang-query/Query.h b/clang-tools-extra/clang-query/Query.h
index 56af486984e..78bcbc79cdf 100644
--- a/clang-tools-extra/clang-query/Query.h
+++ b/clang-tools-extra/clang-query/Query.h
@@ -44,6 +44,7 @@ struct Query : llvm::RefCountedBase<Query> {
   /// \return false if an error occurs, otherwise return true.
   virtual bool run(llvm::raw_ostream &OS, QuerySession &QS) const = 0;
 
+  StringRef RemainingContent;
   const QueryKind Kind;
 };
 
diff --git a/clang-tools-extra/clang-query/QueryParser.cpp b/clang-tools-extra/clang-query/QueryParser.cpp
index 4da2f5da79d..c94a5c8c5a0 100644
--- a/clang-tools-extra/clang-query/QueryParser.cpp
+++ b/clang-tools-extra/clang-query/QueryParser.cpp
@@ -26,7 +26,10 @@ namespace query {
 // is found before End, return StringRef().  Begin is adjusted to exclude the
 // lexed region.
 StringRef QueryParser::lexWord() {
-  Line = Line.ltrim();
+  Line = Line.drop_while([this](char c) {
+    // Don't trim newlines.
+    return StringRef(" \t\v\f\r").contains(c);
+  });
 
   if (Line.empty())
     // Even though the Line is empty, it contains a pointer and
@@ -34,12 +37,12 @@ StringRef QueryParser::lexWord() {
     // code completion.
     return Line;
 
-  if (Line.front() == '#') {
-    Line = {};
-    return StringRef();
-  }
+  StringRef Word;
+  if (Line.front() == '#')
+    Word = Line.substr(0, 1);
+  else
+    Word = Line.take_until(isWhitespace);
 
-  StringRef Word = Line.take_until(isWhitespace);
   Line = Line.drop_front(Word.size());
   return Word;
 }
@@ -125,9 +128,25 @@ template <typename QueryType> QueryRef QueryParser::parseSetOutputKind() {
 }
 
 QueryRef QueryParser::endQuery(QueryRef Q) {
-  const StringRef Extra = Line;
-  if (!lexWord().empty())
-    return new InvalidQuery("unexpected extra input: '" + Extra + "'");
+  StringRef Extra = Line;
+  StringRef ExtraTrimmed = Extra.drop_while(
+      [](char c) { return StringRef(" \t\v\f\r").contains(c); });
+
+  if ((!ExtraTrimmed.empty() && ExtraTrimmed[0] == '\n') ||
+      (ExtraTrimmed.size() >= 2 && ExtraTrimmed[0] == '\r' &&
+       ExtraTrimmed[1] == '\n'))
+    Q->RemainingContent = Extra;
+  else {
+    StringRef TrailingWord = lexWord();
+    if (TrailingWord.front() == '#') {
+      Line = Line.drop_until([](char c) { return c == '\n'; });
+      Line = Line.drop_front();
+      return endQuery(Q);
+    }
+    if (!TrailingWord.empty()) {
+      return new InvalidQuery("unexpected extra input: '" + Extra + "'");
+    }
+  }
   return Q;
 }
 
@@ -193,7 +212,11 @@ QueryRef QueryParser::doParse() {
   switch (QKind) {
   case PQK_Comment:
   case PQK_NoOp:
-    return new NoOpQuery;
+    Line = Line.drop_until([](char c) { return c == '\n'; });
+    Line = Line.drop_front();
+    if (Line.empty())
+      return new NoOpQuery;
+    return doParse();
 
   case PQK_Help:
     return endQuery(new HelpQuery);
@@ -217,7 +240,9 @@ QueryRef QueryParser::doParse() {
       return makeInvalidQueryFromDiagnostics(Diag);
     }
 
-    return new LetQuery(Name, Value);
+    auto *Q = new LetQuery(Name, Value);
+    Q->RemainingContent = Line;
+    return Q;
   }
 
   case PQK_Match: {
@@ -226,12 +251,17 @@ QueryRef QueryParser::doParse() {
 
     Diagnostics Diag;
     auto MatcherSource = Line.trim();
+    auto OrigMatcherSource = MatcherSource;
     Optional<DynTypedMatcher> Matcher = Parser::parseMatcherExpression(
         MatcherSource, nullptr, &QS.NamedValues, &Diag);
     if (!Matcher) {
       return makeInvalidQueryFromDiagnostics(Diag);
     }
-    return new MatchQuery(MatcherSource, *Matcher);
+    auto ActualSource = OrigMatcherSource.slice(0, OrigMatcherSource.size() -
+                                                       MatcherSource.size());
+    auto *Q = new MatchQuery(ActualSource, *Matcher);
+    Q->RemainingContent = MatcherSource;
+    return Q;
   }
 
   case PQK_Set: {
diff --git a/clang-tools-extra/clang-query/tool/ClangQuery.cpp b/clang-tools-extra/clang-query/tool/ClangQuery.cpp
index 80e1c602796..5cfa0acf912 100644
--- a/clang-tools-extra/clang-query/tool/ClangQuery.cpp
+++ b/clang-tools-extra/clang-query/tool/ClangQuery.cpp
@@ -69,13 +69,16 @@ bool runCommandsInFile(const char *ExeName, std::string const &FileName,
     llvm::errs() << ExeName << ": cannot open " << FileName << "\n";
     return 1;
   }
-  while (Input.good()) {
-    std::string Line;
-    std::getline(Input, Line);
 
-    QueryRef Q = QueryParser::parse(Line, QS);
+  std::string FileContent((std::istreambuf_iterator<char>(Input)),
+                          std::istreambuf_iterator<char>());
+
+  StringRef FileContentRef(FileContent);
+  while (!FileContentRef.empty()) {
+    QueryRef Q = QueryParser::parse(FileContentRef, QS);
     if (!Q->run(llvm::outs(), QS))
       return true;
+    FileContentRef = Q->RemainingContent;
   }
   return false;
 }