Add new 'let' command to bind arbitrary values into constants.

Summary:
Add new 'let' command to bind arbitrary values into constants.
These constants can then be used in the matcher expressions.

Reviewers: pcc

CC: cfe-commits

Differential Revision: http://reviews.llvm.org/D3383

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@206984 91177308-0d34-0410-b5e6-96231b3b80d8

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;

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

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,46 +150,89 @@ 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;
 
   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,20 +261,29 @@ 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);
   }
 
   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();

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 @@ class QueryParser {
   /// 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 @@ class QueryParser {
 
   QueryRef parseSetBool(bool QuerySession::*Var);
   QueryRef parseSetOutputKind();
+  QueryRef completeMatcherExpression();
 
   QueryRef endQuery(QueryRef Q);
 
@@ -59,6 +62,8 @@ class QueryParser {
 
   const char *CompletionPos;
   std::vector<llvm::LineEditor::Completion> Completions;
+
+  const QuerySession &QS;
 };
 
 } // namespace query

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 @@ class QuerySession {
   llvm::ArrayRef<ASTUnit *> ASTs;
   OutputKind OutKind;
   bool BindRoot;
+  llvm::StringMap<ast_matchers::dynamic::VariantValue> NamedValues;
 };
 
 } // namespace query

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);
     }
   }

unittests/clang-query/QueryEngineTest.cpp

@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Query.h"
+#include "QueryParser.h"
 #include "QuerySession.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/ASTMatchers/Dynamic/VariantValue.h"
@@ -24,20 +25,22 @@ using namespace clang::ast_matchers::dynamic;
 using namespace clang::query;
 using namespace clang::tooling;
 
-TEST(Query, Basic) {
-  std::unique_ptr<ASTUnit> FooAST(
-      buildASTFromCode("void foo1(void) {}\nvoid foo2(void) {}", "foo.cc"));
-  ASSERT_TRUE(FooAST.get());
-  std::unique_ptr<ASTUnit> BarAST(
-      buildASTFromCode("void bar1(void) {}\nvoid bar2(void) {}", "bar.cc"));
-  ASSERT_TRUE(BarAST.get());
+class QueryEngineTest : public ::testing::Test {
+protected:
+  QueryEngineTest() {}
 
-  ASTUnit *ASTs[] = { FooAST.get(), BarAST.get() };
+  std::unique_ptr<ASTUnit> FooAST{
+      buildASTFromCode("void foo1(void) {}\nvoid foo2(void) {}", "foo.cc")};
+  std::unique_ptr<ASTUnit> BarAST{
+      buildASTFromCode("void bar1(void) {}\nvoid bar2(void) {}", "bar.cc")};
+  ASTUnit *ASTs[2]{FooAST.get(), BarAST.get()};
+  QuerySession S{ASTs};
 
   std::string Str;
-  llvm::raw_string_ostream OS(Str);
-  QuerySession S(ASTs);
+  llvm::raw_string_ostream OS{Str};
+};
 
+TEST_F(QueryEngineTest, Basic) {
   DynTypedMatcher FnMatcher = functionDecl();
   DynTypedMatcher FooMatcher = functionDecl(hasName("foo1"));
 
@@ -108,3 +111,28 @@ TEST(Query, Basic) {
 
   EXPECT_EQ("Not a valid top-level matcher.\n", OS.str());
 }
+
+TEST_F(QueryEngineTest, LetAndMatch) {
+  EXPECT_TRUE(QueryParser::parse("let x \"foo1\"", S)->run(OS, S));
+  EXPECT_EQ("", OS.str());
+  Str.clear();
+
+  EXPECT_TRUE(QueryParser::parse("let y hasName(x)", S)->run(OS, S));
+  EXPECT_EQ("", OS.str());
+  Str.clear();
+
+  EXPECT_TRUE(QueryParser::parse("match functionDecl(y)", S)->run(OS, S));
+  EXPECT_TRUE(OS.str().find("foo.cc:1:1: note: \"root\" binds here") !=
+              std::string::npos);
+  EXPECT_TRUE(OS.str().find("1 match.") != std::string::npos);
+  Str.clear();
+
+  EXPECT_TRUE(QueryParser::parse("unlet x", S)->run(OS, S));
+  EXPECT_EQ("", OS.str());
+  Str.clear();
+
+  EXPECT_FALSE(QueryParser::parse("let y hasName(x)", S)->run(OS, S));
+  EXPECT_EQ("1:2: Error parsing argument 1 for matcher hasName.\n"
+            "1:10: Value not found: x\n", OS.str());
+  Str.clear();
+}