[clangd] Add ClangdServer::customAction() extension point

[sammccall]

Summary

This lets basic AST-consuming actions be defined outside clangdserver.
(it essentially exposes TUScheduler::runWithAST).

Two use cases for now:

• replace ClangdServer::dumpAST, which doesn't really belong in the public interface (a followup patch will do this)
• allow embedders to add/experiment with extra features (e.g. I know some who want to try crazy AST+ML code completion...)

Maybe in the future we should use this sort of mechanism to make ClangdServer
less of a monolith, but that's not in scope for now.
This would probably require a richer and more carefully-thought-out API.

[kadircet]

mostly looks good to me, as discussed offline I would rather expose AST within a thread, both to keep the API changes to a minimum in the future and possibly indicating the "view-ness" of the exposed structs like AST more explicitly.

[sammccall]

Oops sorry, misread "mostly lgtm" and committed.

Changed to use the "InputsAndAST" struct to provide the rest of the info.
Happy to address any further comments...

[kadircet]

no I am totally fine with the change. I just wanted to give Adam some time to take a look too :D

[sammccall]

Yup sorry @adamcz, I was just confused. Likewise, happy to change to a different structure or revert entirely if you have concerns.

41d0edd54e29e994fa7d40961a38e8fca27addac eliminates dumpAST based on this, but is easy to revert too.

[adamcz]

LGTM

clang-tools-extra/clangd/ClangdServer.h

void semanticRanges(PathRef File, const std::vector<Position> &Pos,
                       Callback<std::vector<SelectionRange>> CB);
 
   /// Get all document links in a file.
   void documentLinks(PathRef File, Callback<std::vector<DocumentLink>> CB);
 
   void semanticHighlights(PathRef File,
                           Callback<std::vector<HighlightingToken>>);
 
+  /// Runs an arbitrary action that has access to the AST of the specified file.
+  /// The action will execute on one of ClangdServer's internal threads.
+  /// The AST is only valid for the duration of the callback.
+  /// As with other actions, the file must have been opened.
+  void customAction(PathRef File, llvm::StringRef Name,
+                    Callback<InputsAndAST> Action);
+
   /// Returns estimated memory usage and other statistics for each of the
   /// currently open files.
   /// Overall memory usage of clangd may be significantly more than reported
   /// here, as this metric does not account (at least) for:
   ///   - memory occupied by static and dynamic index,
   ///   - memory required for in-flight requests,
   /// FIXME: those metrics might be useful too, we should add them.
   llvm::StringMap<TUScheduler::FileStats> fileStats() const;

clang-tools-extra/clangd/ClangdServer.cpp

auto Action =
         if (!InpAST)
           return CB(InpAST.takeError());
         CB(clangd::getSemanticHighlightings(InpAST->AST));
       };
   WorkScheduler.runWithAST("SemanticHighlights", File, std::move(Action),
                            TUScheduler::InvalidateOnUpdate);
 }
 
+void ClangdServer::customAction(PathRef File, llvm::StringRef Name,
+                                Callback<InputsAndAST> Action) {
+  WorkScheduler.runWithAST(Name, File, std::move(Action));
+}
+
 llvm::StringMap<TUScheduler::FileStats> ClangdServer::fileStats() const {
   return WorkScheduler.fileStats();
 }
 
 Context ClangdServer::createProcessingContext(PathRef File) const {
   if (!ConfigProvider)
     return Context::current().clone();
 

clang-tools-extra/clangd/unittests/ClangdTests.cpp

 #include "ClangdLSPServer.h"
 #include "ClangdServer.h"
 #include "CodeComplete.h"
 #include "ConfigFragment.h"
 #include "GlobalCompilationDatabase.h"
 #include "Matchers.h"
 #include "SyncAPI.h"
 #include "TestFS.h"
+#include "TestTU.h"
 #include "URI.h"
 #include "support/Path.h"
 #include "support/Threading.h"
 #include "clang/Config/config.h"
 #include "clang/Sema/CodeCompleteConsumer.h"
 #include "clang/Tooling/ArgumentsAdjusters.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Regex.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
 #include <chrono>
 #include <iostream>
 #include <random>
 #include <string>
 #include <thread>
TEST_F(ClangdVFSTest, FallbackWhenWaitingForCompileCommand) {
   ASSERT_TRUE(Server.blockUntilIdleForTest());
   EXPECT_THAT(cantFail(runCodeComplete(Server, FooCpp, Code.point(),
                                        clangd::CodeCompleteOptions()))
                   .Completions,
               ElementsAre(AllOf(Field(&CodeCompletion::Name, "xyz"),
                                 Field(&CodeCompletion::Scope, "ns::"))));
 }
 
+TEST(ClangdServerTest, CustomAction) {
+  OverlayCDB CDB(/*Base=*/nullptr);
+  MockFS FS;
+  ClangdServer Server(CDB, FS, ClangdServer::optsForTest());
+
+  Server.addDocument(testPath("foo.cc"), "void x();");
+  Decl::Kind XKind = Decl::TranslationUnit;
+  EXPECT_THAT_ERROR(runCustomAction(Server, testPath("foo.cc"),
+                                    [&](InputsAndAST AST) {
+                                      XKind = findDecl(AST.AST, "x").getKind();
+                                    }),
+                    llvm::Succeeded());
+  EXPECT_EQ(XKind, Decl::Function);
+}
+
 // Tests fails when built with asan due to stack overflow. So skip running the
 // test as a workaround.
 #if !defined(__has_feature) || !__has_feature(address_sanitizer)
 TEST_F(ClangdVFSTest, TestStackOverflow) {
   MockFS FS;
   ErrorCheckingCallbacks DiagConsumer;
   MockCompilationDatabase CDB;
   ClangdServer Server(CDB, FS, ClangdServer::optsForTest(), &DiagConsumer);

clang-tools-extra/clangd/unittests/SyncAPI.h

 
 llvm::Expected<std::vector<SelectionRange>>
 runSemanticRanges(ClangdServer &Server, PathRef File,
                   const std::vector<Position> &Pos);
 
 llvm::Expected<llvm::Optional<clangd::Path>>
 runSwitchHeaderSource(ClangdServer &Server, PathRef File);
 
+llvm::Error runCustomAction(ClangdServer &Server, PathRef File,
+                            llvm::function_ref<void(InputsAndAST)>);
+
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_SYNCAPI_H

clang-tools-extra/clangd/unittests/SyncAPI.cpp

 
 llvm::Expected<llvm::Optional<clangd::Path>>
 runSwitchHeaderSource(ClangdServer &Server, PathRef File) {
   llvm::Optional<llvm::Expected<llvm::Optional<clangd::Path>>> Result;
   Server.switchSourceHeader(File, capture(Result));
   return std::move(*Result);
 }
 
+llvm::Error runCustomAction(ClangdServer &Server, PathRef File,
+                            llvm::function_ref<void(InputsAndAST)> Action) {
+  llvm::Error Result = llvm::Error::success();
+  Notification Done;
+  Server.customAction(File, "Custom", [&](llvm::Expected<InputsAndAST> AST) {
+    if (!AST)
+      Result = AST.takeError();
+    else
+      Action(*AST);
+    Done.notify();
+  });
+  Done.wait();
+  return Result;
+}
+
 } // namespace clangd
 } // namespace clang