6 files changed, 125 insertions, 38 deletions

diff --git a/clang-tools-extra/clang-query/Query.cpp b/clang-tools-extra/clang-query/Query.cpp
index dcfa80c5277..65a68edea9d 100644
--- a/clang-tools-extra/clang-query/Query.cpp
+++ b/clang-tools-extra/clang-query/Query.cpp
@@ -54,7 +54,7 @@ struct CollectBoundNodes : MatchFinder::MatchCallback {
   }
 };
 
-}
+}  // namespace
 
 bool MatchQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
   unsigned MatchCount = 0;
@@ -124,6 +124,15 @@ bool MatchQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
   return true;
 }
 
+bool LetQuery::run(llvm::raw_ostream &OS, QuerySession &QS) const {
+  if (Value) {
+    QS.NamedValues[Name] = Value;
+  } else {
+    QS.NamedValues.erase(Name);
+  }
+  return true;
+}
+
 #ifndef _MSC_VER
 const QueryKind SetQueryKind<bool>::value;
 const QueryKind SetQueryKind<OutputKind>::value;
diff --git a/clang-tools-extra/clang-query/Query.h b/clang-tools-extra/clang-query/Query.h
index 60514f73d10..f45057016d4 100644
--- a/clang-tools-extra/clang-query/Query.h
+++ b/clang-tools-extra/clang-query/Query.h
@@ -28,9 +28,10 @@ enum QueryKind {
   QK_Invalid,
   QK_NoOp,
   QK_Help,
+  QK_Let,
   QK_Match,
   QK_SetBool,
-  QK_SetOutputKind
+  QK_SetOutputKind,
 };
 
 class QuerySession;
@@ -86,6 +87,17 @@ struct MatchQuery : Query {
   static bool classof(const Query *Q) { return Q->Kind == QK_Match; }
 };
 
+struct LetQuery : Query {
+  LetQuery(StringRef Name, const ast_matchers::dynamic::VariantValue &Value)
+      : Query(QK_Let), Name(Name), Value(Value) {}
+  bool run(llvm::raw_ostream &OS, QuerySession &QS) const override;
+
+  std::string Name;
+  ast_matchers::dynamic::VariantValue Value;
+
+  static bool classof(const Query *Q) { return Q->Kind == QK_Let; }
+};
+
 template <typename T> struct SetQueryKind {};
 
 template <> struct SetQueryKind<bool> {
diff --git a/clang-tools-extra/clang-query/QueryParser.cpp b/clang-tools-extra/clang-query/QueryParser.cpp
index 3a0cd8bcc69..fb0f105f783 100644
--- a/clang-tools-extra/clang-query/QueryParser.cpp
+++ b/clang-tools-extra/clang-query/QueryParser.cpp
@@ -132,12 +132,16 @@ QueryRef QueryParser::endQuery(QueryRef Q) {
   return Q;
 }
 
+namespace {
+
 enum ParsedQueryKind {
   PQK_Invalid,
   PQK_NoOp,
   PQK_Help,
+  PQK_Let,
   PQK_Match,
-  PQK_Set
+  PQK_Set,
+  PQK_Unlet,
 };
 
 enum ParsedQueryVariable {
@@ -146,16 +150,52 @@ enum ParsedQueryVariable {
   PQV_BindRoot
 };
 
+QueryRef makeInvalidQueryFromDiagnostics(const Diagnostics &Diag) {
+  std::string ErrStr;
+  llvm::raw_string_ostream OS(ErrStr);
+  Diag.printToStreamFull(OS);
+  return new InvalidQuery(OS.str());
+}
+
+class QuerySessionSema : public Parser::RegistrySema {
+public:
+  QuerySessionSema(const QuerySession &QS) : QS(QS) {}
+
+  ast_matchers::dynamic::VariantValue getNamedValue(StringRef Name) override {
+    return QS.NamedValues.lookup(Name);
+  }
+
+private:
+  const QuerySession &QS;
+};
+
+}  // namespace
+
+QueryRef QueryParser::completeMatcherExpression() {
+  std::vector<MatcherCompletion> Comps = Parser::completeExpression(
+      StringRef(Begin, End - Begin), CompletionPos - Begin);
+  for (std::vector<MatcherCompletion>::iterator I = Comps.begin(),
+                                                E = Comps.end();
+       I != E; ++I) {
+    Completions.push_back(LineEditor::Completion(I->TypedText, I->MatcherDecl));
+  }
+  return QueryRef();
+}
+
 QueryRef QueryParser::doParse() {
   StringRef CommandStr;
   ParsedQueryKind QKind = lexOrCompleteWord<ParsedQueryKind>(CommandStr)
                               .Case("", PQK_NoOp)
                               .Case("help", PQK_Help)
                               .Case("m", PQK_Match, /*IsCompletion=*/false)
+                              .Case("let", PQK_Let)
                               .Case("match", PQK_Match)
                               .Case("set", PQK_Set)
+                              .Case("unlet", PQK_Unlet)
                               .Default(PQK_Invalid);
 
+  QuerySessionSema S(QS);
+
   switch (QKind) {
   case PQK_NoOp:
     return new NoOpQuery;
@@ -163,29 +203,36 @@ QueryRef QueryParser::doParse() {
   case PQK_Help:
     return endQuery(new HelpQuery);
 
+  case PQK_Let: {
+    StringRef Name = lexWord();
+
+    if (Name.empty())
+      return new InvalidQuery("expected variable name");
+
+    if (CompletionPos)
+      return completeMatcherExpression();
+
+    Diagnostics Diag;
+    ast_matchers::dynamic::VariantValue Value;
+    if (!Parser::parseExpression(StringRef(Begin, End - Begin), &S, &Value,
+                                 &Diag)) {
+      return makeInvalidQueryFromDiagnostics(Diag);
+    }
+
+    return new LetQuery(Name, Value);
+  }
+
   case PQK_Match: {
-    if (CompletionPos) {
-      std::vector<MatcherCompletion> Comps = Parser::completeExpression(
-          StringRef(Begin, End - Begin), CompletionPos - Begin);
-      for (std::vector<MatcherCompletion>::iterator I = Comps.begin(),
-                                                    E = Comps.end();
-           I != E; ++I) {
-        Completions.push_back(
-            LineEditor::Completion(I->TypedText, I->MatcherDecl));
-      }
-      return QueryRef();
-    } else {
-      Diagnostics Diag;
-      Optional<DynTypedMatcher> Matcher =
-          Parser::parseMatcherExpression(StringRef(Begin, End - Begin), &Diag);
-      if (!Matcher) {
-        std::string ErrStr;
-        llvm::raw_string_ostream OS(ErrStr);
-        Diag.printToStreamFull(OS);
-        return new InvalidQuery(OS.str());
-      }
-      return new MatchQuery(*Matcher);
+    if (CompletionPos)
+      return completeMatcherExpression();
+
+    Diagnostics Diag;
+    Optional<DynTypedMatcher> Matcher = Parser::parseMatcherExpression(
+        StringRef(Begin, End - Begin), &S, &Diag);
+    if (!Matcher) {
+      return makeInvalidQueryFromDiagnostics(Diag);
     }
+    return new MatchQuery(*Matcher);
   }
 
   case PQK_Set: {
@@ -214,6 +261,15 @@ QueryRef QueryParser::doParse() {
     return endQuery(Q);
   }
 
+  case PQK_Unlet: {
+    StringRef Name = lexWord();
+
+    if (Name.empty())
+      return new InvalidQuery("expected variable name");
+
+    return endQuery(new LetQuery(Name, {}));
+  }
+
   case PQK_Invalid:
     return new InvalidQuery("unknown command: " + CommandStr);
   }
@@ -221,13 +277,13 @@ QueryRef QueryParser::doParse() {
   llvm_unreachable("Invalid query kind");
 }
 
-QueryRef QueryParser::parse(StringRef Line) {
-  return QueryParser(Line).doParse();
+QueryRef QueryParser::parse(StringRef Line, const QuerySession &QS) {
+  return QueryParser(Line, QS).doParse();
 }
 
-std::vector<LineEditor::Completion> QueryParser::complete(StringRef Line,
-                                                          size_t Pos) {
-  QueryParser P(Line);
+std::vector<LineEditor::Completion>
+QueryParser::complete(StringRef Line, size_t Pos, const QuerySession &QS) {
+  QueryParser P(Line, QS);
   P.CompletionPos = Line.data() + Pos;
 
   P.doParse();
diff --git a/clang-tools-extra/clang-query/QueryParser.h b/clang-tools-extra/clang-query/QueryParser.h
index b448c907dc2..5ebeaf2c473 100644
--- a/clang-tools-extra/clang-query/QueryParser.h
+++ b/clang-tools-extra/clang-query/QueryParser.h
@@ -11,6 +11,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANG_QUERY_QUERY_PARSER_H
 
 #include "Query.h"
+#include "QuerySession.h"
 #include "llvm/LineEditor/LineEditor.h"
 #include <stddef.h>
 
@@ -24,19 +25,20 @@ public:
   /// Parse \a Line as a query.
   ///
   /// \return A QueryRef representing the query, which may be an InvalidQuery.
-  static QueryRef parse(StringRef Line);
+  static QueryRef parse(StringRef Line, const QuerySession &QS);
 
   /// Compute a list of completions for \a Line assuming a cursor at
   /// \param Pos characters past the start of \a Line, ordered from most
   /// likely to least likely.
   ///
   /// \return A vector of completions for \a Line.
-  static std::vector<llvm::LineEditor::Completion> complete(StringRef Line,
-                                                            size_t Pos);
+  static std::vector<llvm::LineEditor::Completion>
+  complete(StringRef Line, size_t Pos, const QuerySession &QS);
 
 private:
-  QueryParser(StringRef Line)
-      : Begin(Line.data()), End(Line.data() + Line.size()), CompletionPos(0) {}
+  QueryParser(StringRef Line, const QuerySession &QS)
+      : Begin(Line.data()), End(Line.data() + Line.size()), CompletionPos(0),
+        QS(QS) {}
 
   StringRef lexWord();
 
@@ -45,6 +47,7 @@ private:
 
   QueryRef parseSetBool(bool QuerySession::*Var);
   QueryRef parseSetOutputKind();
+  QueryRef completeMatcherExpression();
 
   QueryRef endQuery(QueryRef Q);
 
@@ -59,6 +62,8 @@ private:
 
   const char *CompletionPos;
   std::vector<llvm::LineEditor::Completion> Completions;
+
+  const QuerySession &QS;
 };
 
 } // namespace query
diff --git a/clang-tools-extra/clang-query/QuerySession.h b/clang-tools-extra/clang-query/QuerySession.h
index fdbd9eae018..b4361003bec 100644
--- a/clang-tools-extra/clang-query/QuerySession.h
+++ b/clang-tools-extra/clang-query/QuerySession.h
@@ -11,7 +11,9 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANG_QUERY_QUERY_SESSION_H
 
 #include "Query.h"
+#include "clang/ASTMatchers/Dynamic/VariantValue.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringMap.h"
 
 namespace clang {
 
@@ -28,6 +30,7 @@ public:
   llvm::ArrayRef<ASTUnit *> ASTs;
   OutputKind OutKind;
   bool BindRoot;
+  llvm::StringMap<ast_matchers::dynamic::VariantValue> NamedValues;
 };
 
 } // namespace query
diff --git a/clang-tools-extra/clang-query/tool/ClangQuery.cpp b/clang-tools-extra/clang-query/tool/ClangQuery.cpp
index 92865fb3d85..d95892618ce 100644
--- a/clang-tools-extra/clang-query/tool/ClangQuery.cpp
+++ b/clang-tools-extra/clang-query/tool/ClangQuery.cpp
@@ -95,7 +95,7 @@ int main(int argc, const char **argv) {
     for (cl::list<std::string>::iterator I = Commands.begin(),
                                          E = Commands.end();
          I != E; ++I) {
-      QueryRef Q = QueryParser::parse(I->c_str());
+      QueryRef Q = QueryParser::parse(I->c_str(), QS);
       if (!Q->run(llvm::outs(), QS))
         return 1;
     }
@@ -112,16 +112,18 @@ int main(int argc, const char **argv) {
         std::string Line;
         std::getline(Input, Line);
 
-        QueryRef Q = QueryParser::parse(Line.c_str());
+        QueryRef Q = QueryParser::parse(Line.c_str(), QS);
         if (!Q->run(llvm::outs(), QS))
           return 1;
       }
     }
   } else {
     LineEditor LE("clang-query");
-    LE.setListCompleter(QueryParser::complete);
+    LE.setListCompleter([&QS](StringRef Line, size_t Pos) {
+      return QueryParser::complete(Line, Pos, QS);
+    });
     while (llvm::Optional<std::string> Line = LE.readLine()) {
-      QueryRef Q = QueryParser::parse(*Line);
+      QueryRef Q = QueryParser::parse(*Line, QS);
       Q->run(llvm::outs(), QS);
     }
   }