[Clangd] Add AST search capabilities from clang-query

clang-tools-extra/clangd/CMakeLists.txt

@@ -60,6 +60,7 @@ endif()
 
 include_directories(BEFORE "${CMAKE_CURRENT_BINARY_DIR}/../clang-tidy")
 include_directories(BEFORE "${CMAKE_CURRENT_SOURCE_DIR}/../include-cleaner/include")
+include_directories(BEFORE "${CMAKE_CURRENT_SOURCE_DIR}/../clang-query")
 
 add_clang_library(clangDaemon STATIC
   AST.cpp
@@ -183,6 +184,7 @@ target_link_libraries(clangDaemon
   ${LLVM_PTHREAD_LIB}
 
   clangIncludeCleaner
+  clangQuery
   clangTidy
   clangTidyUtils
 

clang-tools-extra/clangd/ClangdLSPServer.cpp

@@ -76,6 +76,7 @@ std::optional<int64_t> decodeVersion(llvm::StringRef Encoded) {
 const llvm::StringLiteral ApplyFixCommand = "clangd.applyFix";
 const llvm::StringLiteral ApplyTweakCommand = "clangd.applyTweak";
 const llvm::StringLiteral ApplyRenameCommand = "clangd.applyRename";
+constexpr llvm::StringLiteral SearchASTMethod = "textDocument/searchAST";
 
 CodeAction toCodeAction(const ClangdServer::CodeActionResult::Rename &R,
                         const URIForFile &File) {
@@ -638,6 +639,9 @@ void ClangdLSPServer::onInitialize(const InitializeParams &Params,
       {"workspaceSymbolProvider", true},
       {"referencesProvider", true},
       {"astProvider", true}, // clangd extension
+      {"astSearchProvider",
+       llvm::json::Object{{"search", true},
+                          {"replace", false}}}, // clangd extension
       {"typeHierarchyProvider", true},
       // Unfortunately our extension made use of the same capability name as the
       // standard. Advertise this capability to tell clients that implement our
@@ -852,6 +856,16 @@ void ClangdLSPServer::onCommandApplyRename(const RenameParams &R,
   });
 }
 
+void ClangdLSPServer::onMethodSearchAST(const SearchASTArgs &Args,
+                                        Callback<llvm::json::Value> Reply) {
+  Server->findAST(Args, [Reply = std::move(Reply)](
+                            llvm::Expected<BoundASTNodes> BoundNodes) mutable {
+    if (!BoundNodes)
+      return Reply(BoundNodes.takeError());
+    return Reply(*BoundNodes);
+  });
+}
+
 void ClangdLSPServer::applyEdit(WorkspaceEdit WE, llvm::json::Value Success,
                                 Callback<llvm::json::Value> Reply) {
   ApplyWorkspaceEditParams Edit;
@@ -1728,6 +1742,7 @@ void ClangdLSPServer::bindMethods(LSPBinder &Bind,
   Bind.command(ApplyFixCommand, this, &ClangdLSPServer::onCommandApplyEdit);
   Bind.command(ApplyTweakCommand, this, &ClangdLSPServer::onCommandApplyTweak);
   Bind.command(ApplyRenameCommand, this, &ClangdLSPServer::onCommandApplyRename);
+  Bind.method(SearchASTMethod, this, &ClangdLSPServer::onMethodSearchAST);
 
   ApplyWorkspaceEdit = Bind.outgoingMethod("workspace/applyEdit");
   PublishDiagnostics = Bind.outgoingNotification("textDocument/publishDiagnostics");

clang-tools-extra/clangd/ClangdLSPServer.h

@@ -186,6 +186,7 @@ class ClangdLSPServer : private ClangdServer::Callbacks,
   void onCommandApplyEdit(const WorkspaceEdit &, Callback<llvm::json::Value>);
   void onCommandApplyTweak(const TweakArgs &, Callback<llvm::json::Value>);
   void onCommandApplyRename(const RenameParams &, Callback<llvm::json::Value>);
+  void onMethodSearchAST(const SearchASTArgs &, Callback<llvm::json::Value>);
 
   /// Outgoing LSP calls.
   LSPBinder::OutgoingMethod<ApplyWorkspaceEditParams,

clang-tools-extra/clangd/ClangdServer.cpp

@@ -810,6 +810,42 @@ void ClangdServer::locateSymbolAt(PathRef File, Position Pos,
   WorkScheduler->runWithAST("Definitions", File, std::move(Action));
 }
 
+void ClangdServer::findAST(SearchASTArgs const &Args,
+                           Callback<BoundASTNodes> CB) {
+  auto Action = [Args, CB = std::move(CB)](
+                    llvm::Expected<InputsAndAST> InpAST) mutable {
+    if (!InpAST)
+      return CB(InpAST.takeError());
+    auto BoundNodes = clangd::locateASTQuery(InpAST->AST, Args);
+    if (!BoundNodes)
+      return CB(BoundNodes.takeError());
+    if (BoundNodes->empty())
+      return CB(error("No matching AST nodes found"));
+
+    auto &&AST = InpAST->AST;
+    // Convert BoundNodes to a vector of vectors to ASTNode's.
+    BoundASTNodes Result;
+    Result.reserve(BoundNodes->size());
+    for (auto &&BN : *BoundNodes) {
+      auto &&Map = BN.getMap();
+      BoundASTNodes::value_type BAN;
+      for (const auto &[Key, Value] : Map) {
+        BAN.emplace(Key, dumpAST(Value, AST.getTokens(), AST.getASTContext()));
+      }
+      if (BAN.empty())
+        continue;
+      Result.push_back(std::move(BAN));
+    }
+    if (Result.empty()) {
+      return CB(error("No AST nodes found for the query"));
+    }
+    CB(std::move(Result));
+  };
+
+  WorkScheduler->runWithAST("Definitions", Args.textDocument.uri.file(),
+                            std::move(Action));
+}
+
 void ClangdServer::switchSourceHeader(
     PathRef Path, Callback<std::optional<clangd::Path>> CB) {
   // We want to return the result as fast as possible, strategy is:

clang-tools-extra/clangd/ClangdServer.h

@@ -30,11 +30,14 @@
 #include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/ThreadsafeFS.h"
+#include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/ASTMatchers/ASTMatchers.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include <functional>
+#include <map>
 #include <memory>
 #include <optional>
 #include <string>
@@ -260,6 +263,8 @@ class ClangdServer {
   void locateSymbolAt(PathRef File, Position Pos,
                       Callback<std::vector<LocatedSymbol>> CB);
 
+  void findAST(const SearchASTArgs &Args, Callback<BoundASTNodes> CB);
+
   /// Switch to a corresponding source file when given a header file, and vice
   /// versa.
   void switchSourceHeader(PathRef Path,

clang-tools-extra/clangd/Protocol.cpp

@@ -13,6 +13,7 @@
 #include "Protocol.h"
 #include "URI.h"
 #include "support/Logger.h"
+#include "clang/AST/ASTTypeTraits.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/StringExtras.h"
@@ -1650,6 +1651,18 @@ bool fromJSON(const llvm::json::Value &Params, SelectionRangeParams &S,
          O.map("positions", S.positions);
 }
 
+bool fromJSON(const llvm::json::Value &Params, SearchASTArgs &Args,
+              llvm::json::Path P) {
+  llvm::json::ObjectMapper O(Params, P);
+  return O && O.map("query", Args.searchQuery) &&
+         O.map("textDocument", Args.textDocument)
+      // && O.map("bindRoot", Args.bindRoot); TODO: add bindRoot to extend this
+      // feature
+      // && O.map("traversalKind", Args.tk); TODO: add traversalKind to extend
+      // this feature
+      ;
+}
+
 llvm::json::Value toJSON(const SelectionRange &Out) {
   if (Out.parent) {
     return llvm::json::Object{{"range", Out.range},

clang-tools-extra/clangd/Protocol.h

@@ -26,6 +26,7 @@
 #include "URI.h"
 #include "index/SymbolID.h"
 #include "support/MemoryTree.h"
+#include "clang/AST/ASTTypeTraits.h"
 #include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/JSON.h"
@@ -1451,6 +1452,19 @@ struct RenameParams {
 bool fromJSON(const llvm::json::Value &, RenameParams &, llvm::json::Path);
 llvm::json::Value toJSON(const RenameParams &);
 
+struct SearchASTArgs {
+  std::string searchQuery;
+  TextDocumentIdentifier textDocument;
+
+  // Todo (extend feature): make them members and modifiable:
+  /// wheter the whole query is shown
+  static auto constexpr BindRoot = true;
+  /// Simplify things for users; default for now.
+  static auto constexpr Tk = TraversalKind::TK_IgnoreUnlessSpelledInSource;
+};
+bool fromJSON(const llvm::json::Value &, SearchASTArgs &, llvm::json::Path);
+using BoundASTNodes = std::vector<std::map<std::string, struct ASTNode>>;
+
 struct PrepareRenameResult {
   /// Range of the string to rename.
   Range range;

clang-tools-extra/clangd/XRefs.cpp

@@ -14,6 +14,8 @@
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "Quality.h"
+#include "Query.h"
+#include "QuerySession.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "URI.h"
@@ -41,6 +43,10 @@
 #include "clang/AST/StmtCXX.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/AST/Type.h"
+#include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/ASTMatchers/Dynamic/Diagnostics.h"
+#include "clang/ASTMatchers/Dynamic/Parser.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
@@ -52,6 +58,7 @@
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Index/USRGeneration.h"
 #include "clang/Lex/Lexer.h"
+#include "clang/Parse/Parser.h"
 #include "clang/Sema/HeuristicResolver.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -66,6 +73,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
+#include <algorithm>
+#include <cmath>
 #include <optional>
 #include <string>
 #include <vector>
@@ -773,6 +782,57 @@ const syntax::Token *findNearbyIdentifier(const SpelledWord &Word,
   return BestTok;
 }
 
+auto locateASTQuery(ParsedAST &AST, SearchASTArgs const &Query)
+    -> llvm::Expected<std::vector<ast_matchers::BoundNodes>> {
+  using namespace ast_matchers;
+  using namespace ast_matchers::dynamic;
+  using ast_matchers::dynamic::Parser;
+
+  Diagnostics Diag;
+  auto MatcherSource = llvm::StringRef(Query.searchQuery).ltrim();
+
+  std::optional<DynTypedMatcher> Matcher = Parser::parseMatcherExpression(
+      MatcherSource,
+      nullptr /* is this sema instance usefull, to reduce overhead?*/,
+      nullptr /*we currently don't support let*/, &Diag);
+  if (!Matcher) {
+    return error("Not a valid top-level matcher: {}.", Diag.toString());
+  }
+
+  struct CollectBoundNodes : MatchFinder::MatchCallback {
+    std::vector<BoundNodes> *Bindings;
+    CollectBoundNodes(std::vector<BoundNodes> &Bindings)
+        : Bindings(&Bindings) {}
+    void run(const MatchFinder::MatchResult &Result) override {
+      Bindings->push_back(Result.Nodes);
+    }
+  };
+
+  DynTypedMatcher MaybeBoundMatcher = *Matcher;
+  if (Query.BindRoot) {
+    std::optional<DynTypedMatcher> M = Matcher->tryBind("root");
+    if (M)
+      MaybeBoundMatcher = *M;
+  }
+  std::vector<BoundNodes> Matches;
+  CollectBoundNodes Collect(Matches);
+
+  MatchFinder::MatchFinderOptions Opt;
+  Opt.IgnoreSystemHeaders = true;
+  MatchFinder Finder{Opt};
+  if (!Finder.addDynamicMatcher(MaybeBoundMatcher, &Collect)) {
+    return error("Can't add matcher.");
+  }
+
+  ASTContext &Ctx = AST.getASTContext();
+
+  auto OldTK = Ctx.getParentMapContext().getTraversalKind();
+  Ctx.getParentMapContext().setTraversalKind(Query.Tk);
+  Finder.matchAST(Ctx);
+  Ctx.getParentMapContext().setTraversalKind(OldTK);
+  return Matches;
+}
+
 std::vector<LocatedSymbol> locateSymbolAt(ParsedAST &AST, Position Pos,
                                           const SymbolIndex *Index) {
   const auto &SM = AST.getSourceManager();

clang-tools-extra/clangd/XRefs.h

@@ -19,6 +19,7 @@
 #include "index/SymbolID.h"
 #include "support/Path.h"
 #include "clang/AST/ASTTypeTraits.h"
+#include "clang/ASTMatchers/ASTMatchers.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
@@ -32,6 +33,15 @@ class TokenBuffer;
 namespace clangd {
 class ParsedAST;
 
+struct LocatedAST {
+  ast_matchers::BoundNodes &AST;
+};
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &, const LocatedAST &);
+
+auto locateASTQuery(ParsedAST &AST, SearchASTArgs const &)
+    -> llvm::Expected<std::vector<ast_matchers::BoundNodes>>;
+
 // Describes where a symbol is declared and defined (as far as clangd knows).
 // There are three cases:
 //  - a declaration only, no definition is known (e.g. only header seen)

clang-tools-extra/clangd/test/find-in-ast.test

@@ -0,0 +1,37 @@
+# RUN: clangd -lit-test < %s | FileCheck -strict-whitespace --dump-input always %s
+void bob();
+void f() {
+  bob();
+}
+---
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{"textDocument": {"foldingRange": {"lineFoldingOnly": true}}},"trace":"off"}}
+---
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"languageId":"cpp","text":"void bob();\nvoid f() {\n  bob();\n}\n","uri":"test:///foo.cpp","version":1}}}
+---
+{"id":1,"jsonrpc":"2.0","method":"textDocument/searchAST","params":{"textDocument":{"uri":"test:///foo.cpp"},"query":"declRefExpr(to(namedDecl(hasName(\"bob\"))))"}}
+#      CHECK:  "id": 1,
+# CHECK-NEXT:  "jsonrpc": "2.0",
+# CHECK-NEXT:  "result": [
+# CHECK-NEXT:    {
+# CHECK-NEXT:      "root": {
+# CHECK-NEXT:        "arcana": "DeclRefExpr {{.*}} 'void ()' lvalue Function {{.*}} 'bob' 'void ()'",
+# CHECK-NEXT:        "detail": "bob",
+# CHECK-NEXT:        "kind": "DeclRef",
+# CHECK-NEXT:        "range": {
+# CHECK-NEXT:          "end": {
+# CHECK-NEXT:            "character": 5,
+# CHECK-NEXT:            "line": 2
+# CHECK-NEXT:          },
+# CHECK-NEXT:          "start": {
+# CHECK-NEXT:            "character": 2,
+# CHECK-NEXT:            "line": 2
+# CHECK-NEXT:          }
+# CHECK-NEXT:        },
+# CHECK-NEXT:        "role": "expression"
+# CHECK-NEXT:      }
+# CHECK-NEXT:    }
+# CHECK-NEXT:  ]
+---
+{"jsonrpc":"2.0","id":5,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

clang-tools-extra/clangd/test/initialize-params.test

@@ -6,6 +6,10 @@
 # CHECK-NEXT:  "result": {
 # CHECK-NEXT:    "capabilities": {
 # CHECK-NEXT:      "astProvider": true,
+# CHECK-NEXT:      "astSearchProvider": {
+# CHECK-NEXT:        "replace": false,
+# CHECK-NEXT:        "search": true
+# CHECK-NEXT:      },
 # CHECK-NEXT:      "callHierarchyProvider": true,
 # CHECK-NEXT:      "clangdInlayHintsProvider": true,
 # CHECK-NEXT:      "codeActionProvider": true,

llvm/include/llvm/Support/JSON.h

@@ -111,6 +111,10 @@ class Object {
   // (using std::pair forces extra copies).
   struct KV;
   explicit Object(std::initializer_list<KV> Properties);
+  template <typename Collection> explicit Object(Collection &&C) {
+    for (auto &&P : C)
+      M.insert(P);
+  }
 
   iterator begin() { return M.begin(); }
   const_iterator begin() const { return M.begin(); }