From 5890717e7098489ff953a9f96d429a8a02477b2d Mon Sep 17 00:00:00 2001
From: Aaron Ballman <aaron@aaronballman.com>
Date: Thu, 6 Aug 2015 11:56:57 +0000
Subject: [PATCH] Add the "quit" command as a way to terminate clang-query
 interactive sessions.

llvm-svn: 244206
---
 clang-tools-extra/clang-query/Query.cpp           | 9 ++++++++-
 clang-tools-extra/clang-query/Query.h             | 9 +++++++++
 clang-tools-extra/clang-query/QueryParser.cpp     | 5 +++++
 clang-tools-extra/clang-query/QuerySession.h      | 3 ++-
 clang-tools-extra/clang-query/tool/ClangQuery.cpp | 2 ++
 5 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/clang-tools-extra/clang-query/Query.cpp b/clang-tools-extra/clang-query/Query.cpp
index d5d755157df3..74eb6ea44a85 100644
--- a/clang-tools-extra/clang-query/Query.cpp
+++ b/clang-tools-extra/clang-query/Query.cpp
@@ -44,7 +44,14 @@ bool HelpQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
         "  set output (diag|print|dump)      "
         "Set whether to print bindings as diagnostics,\n"
         "                                    "
-        "AST pretty prints or AST dumps.\n\n";
+        "AST pretty prints or AST dumps.\n"
+        "  quit                              "
+        "Terminates the query session.\n\n";
+  return true;
+}
+
+bool QuitQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
+  QS.Terminate = true;
   return true;
 }
 
diff --git a/clang-tools-extra/clang-query/Query.h b/clang-tools-extra/clang-query/Query.h
index f45057016d46..109336a9ddc4 100644
--- a/clang-tools-extra/clang-query/Query.h
+++ b/clang-tools-extra/clang-query/Query.h
@@ -32,6 +32,7 @@ enum QueryKind {
   QK_Match,
   QK_SetBool,
   QK_SetOutputKind,
+  QK_Quit
 };
 
 class QuerySession;
@@ -76,6 +77,14 @@ struct HelpQuery : Query {
   static bool classof(const Query *Q) { return Q->Kind == QK_Help; }
 };
 
+/// Query for "quit".
+struct QuitQuery : Query {
+  QuitQuery() : Query(QK_Quit) {}
+  bool run(llvm::raw_ostream &OS, QuerySession &QS) const override;
+
+  static bool classof(const Query *Q) { return Q->Kind == QK_Quit; }
+};
+
 /// Query for "match MATCHER".
 struct MatchQuery : Query {
   MatchQuery(const ast_matchers::dynamic::DynTypedMatcher &Matcher)
diff --git a/clang-tools-extra/clang-query/QueryParser.cpp b/clang-tools-extra/clang-query/QueryParser.cpp
index 7df7dd2844b9..e25b69700b05 100644
--- a/clang-tools-extra/clang-query/QueryParser.cpp
+++ b/clang-tools-extra/clang-query/QueryParser.cpp
@@ -142,6 +142,7 @@ enum ParsedQueryKind {
   PQK_Match,
   PQK_Set,
   PQK_Unlet,
+  PQK_Quit
 };
 
 enum ParsedQueryVariable {
@@ -181,6 +182,7 @@ QueryRef QueryParser::doParse() {
                               .Case("match", PQK_Match)
                               .Case("set", PQK_Set)
                               .Case("unlet", PQK_Unlet)
+                              .Case("quit", PQK_Quit)
                               .Default(PQK_Invalid);
 
   switch (QKind) {
@@ -190,6 +192,9 @@ QueryRef QueryParser::doParse() {
   case PQK_Help:
     return endQuery(new HelpQuery);
 
+  case PQK_Quit:
+    return endQuery(new QuitQuery);
+
   case PQK_Let: {
     StringRef Name = lexWord();
 
diff --git a/clang-tools-extra/clang-query/QuerySession.h b/clang-tools-extra/clang-query/QuerySession.h
index e6a26acd2519..162289f8c3a3 100644
--- a/clang-tools-extra/clang-query/QuerySession.h
+++ b/clang-tools-extra/clang-query/QuerySession.h
@@ -25,11 +25,12 @@ namespace query {
 class QuerySession {
 public:
   QuerySession(llvm::ArrayRef<std::unique_ptr<ASTUnit>> ASTs)
-      : ASTs(ASTs), OutKind(OK_Diag), BindRoot(true) {}
+      : ASTs(ASTs), OutKind(OK_Diag), BindRoot(true), Terminate(false) {}
 
   llvm::ArrayRef<std::unique_ptr<ASTUnit>> ASTs;
   OutputKind OutKind;
   bool BindRoot;
+  bool Terminate;
   llvm::StringMap<ast_matchers::dynamic::VariantValue> NamedValues;
 };
 
diff --git a/clang-tools-extra/clang-query/tool/ClangQuery.cpp b/clang-tools-extra/clang-query/tool/ClangQuery.cpp
index df7190d4013b..f1c97ce423b9 100644
--- a/clang-tools-extra/clang-query/tool/ClangQuery.cpp
+++ b/clang-tools-extra/clang-query/tool/ClangQuery.cpp
@@ -111,6 +111,8 @@ int main(int argc, const char **argv) {
       QueryRef Q = QueryParser::parse(*Line, QS);
       Q->run(llvm::outs(), QS);
       llvm::outs().flush();
+      if (QS.Terminate)
+        break;
     }
   }