1 files changed, 42 insertions, 12 deletions

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: {