[clangd] Make AST-based signals available to runWithPreamble.

Many useful signals can be derived from a valid AST which is regularly updated by
the ASTWorker. `runWithPreamble` does not have access to the ParsedAST
but it can be provided access to some signals derived from a (possibly
stale) AST.

Differential Revision: https://reviews.llvm.org/D94424

clang-tools-extra/clangd/ASTSignals.cpp

@@ -0,0 +1,42 @@
+//===--- ASTSignals.cpp - LSP server -----------------------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "ASTSignals.h"
+#include "AST.h"
+#include "FindTarget.h"
+
+namespace clang {
+namespace clangd {
+ASTSignals ASTSignals::derive(const ParsedAST &AST) {
+  ASTSignals Signals;
+  const SourceManager &SM = AST.getSourceManager();
+  findExplicitReferences(AST.getASTContext(), [&](ReferenceLoc Ref) {
+    for (const NamedDecl *ND : Ref.Targets) {
+      if (!isInsideMainFile(Ref.NameLoc, SM))
+        continue;
+      SymbolID ID = getSymbolID(ND);
+      if (!ID)
+        continue;
+      unsigned &SymbolCount = Signals.ReferencedSymbols[ID];
+      SymbolCount++;
+      // Process namespace only when we see the symbol for the first time.
+      if (SymbolCount != 1)
+        continue;
+      if (const auto *NSD = dyn_cast<NamespaceDecl>(ND->getDeclContext())) {
+        if (NSD->isAnonymousNamespace())
+          continue;
+        std::string NS = printNamespaceScope(*NSD);
+        if (!NS.empty())
+          Signals.RelatedNamespaces[NS]++;
+      }
+    }
+  });
+  return Signals;
+}
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/ASTSignals.h

@@ -0,0 +1,39 @@
+//===--- ASTSignals.h - LSP server -------------------------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTSIGNALS_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTSIGNALS_H
+
+#include "ParsedAST.h"
+#include "index/SymbolID.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace clang {
+namespace clangd {
+
+/// Signals derived from a valid AST of a file.
+/// Provides information that can only be extracted from the AST to actions that
+/// can't access an AST. The signals are computed and updated asynchronously by
+/// the ASTWorker and thus they are always stale and also can be absent.
+/// Example usage: Information about the declarations used in a file affects
+/// code-completion ranking in that file.
+struct ASTSignals {
+  /// Number of occurrences of each symbol present in the file.
+  llvm::DenseMap<SymbolID, unsigned> ReferencedSymbols;
+  /// Namespaces whose symbols are used in the file, and the number of such
+  /// distinct symbols.
+  llvm::StringMap<unsigned> RelatedNamespaces;
+
+  static ASTSignals derive(const ParsedAST &AST);
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTSIGNALS_H

clang-tools-extra/clangd/CMakeLists.txt

@@ -46,6 +46,7 @@ include_directories(BEFORE "${CMAKE_CURRENT_BINARY_DIR}/../clang-tidy")
 
 add_clang_library(clangDaemon
   AST.cpp
+  ASTSignals.cpp
   ClangdLSPServer.cpp
   ClangdServer.cpp
   CodeComplete.cpp

clang-tools-extra/clangd/TUScheduler.cpp

@@ -392,7 +392,8 @@ public:
              TUScheduler::ASTActionInvalidation);
   bool blockUntilIdle(Deadline Timeout) const;
 
-  std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
+  std::shared_ptr<const PreambleData> getPossiblyStalePreamble(
+      std::shared_ptr<const ASTSignals> *ASTSignals = nullptr) const;
 
   /// Used to inform ASTWorker about a new preamble build by PreambleThread.
   /// Diagnostics are only published through this callback. This ensures they
@@ -437,6 +438,8 @@ private:
   void generateDiagnostics(std::unique_ptr<CompilerInvocation> Invocation,
                            ParseInputs Inputs, std::vector<Diag> CIDiags);
 
+  void updateASTSignals(ParsedAST &AST);
+
   // Must be called exactly once on processing thread. Will return after
   // stop() is called on a separate thread and all pending requests are
   // processed.
@@ -499,6 +502,7 @@ private:
   /// Signalled whenever a new request has been scheduled or processing of a
   /// request has completed.
   mutable std::condition_variable RequestsCV;
+  std::shared_ptr<const ASTSignals> LatestASTSignals; /* GUARDED_BY(Mutex) */
   /// Latest build preamble for current TU.
   /// None means no builds yet, null means there was an error while building.
   /// Only written by ASTWorker's thread.
@@ -830,6 +834,16 @@ void ASTWorker::updatePreamble(std::unique_ptr<CompilerInvocation> CI,
   RequestsCV.notify_all();
 }
 
+void ASTWorker::updateASTSignals(ParsedAST &AST) {
+  auto Signals = std::make_shared<const ASTSignals>(ASTSignals::derive(AST));
+  // Existing readers of ASTSignals will have their copy preserved until the
+  // read is completed. The last reader deletes the old ASTSignals.
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    std::swap(LatestASTSignals, Signals);
+  }
+}
+
 void ASTWorker::generateDiagnostics(
     std::unique_ptr<CompilerInvocation> Invocation, ParseInputs Inputs,
     std::vector<Diag> CIDiags) {
@@ -908,6 +922,7 @@ void ASTWorker::generateDiagnostics(
   if (*AST) {
     trace::Span Span("Running main AST callback");
     Callbacks.onMainAST(FileName, **AST, RunPublish);
+    updateASTSignals(**AST);
   } else {
     // Failed to build the AST, at least report diagnostics from the
     // command line if there were any.
@@ -925,9 +940,11 @@ void ASTWorker::generateDiagnostics(
   }
 }
 
-std::shared_ptr<const PreambleData>
-ASTWorker::getPossiblyStalePreamble() const {
+std::shared_ptr<const PreambleData> ASTWorker::getPossiblyStalePreamble(
+    std::shared_ptr<const ASTSignals> *ASTSignals) const {
   std::lock_guard<std::mutex> Lock(Mutex);
+  if (ASTSignals)
+    *ASTSignals = LatestASTSignals;
   return LatestPreamble ? *LatestPreamble : nullptr;
 }
 
@@ -1364,38 +1381,40 @@ void TUScheduler::runWithPreamble(llvm::StringRef Name, PathRef File,
   if (!PreambleTasks) {
     trace::Span Tracer(Name);
     SPAN_ATTACH(Tracer, "file", File);
+    std::shared_ptr<const ASTSignals> Signals;
     std::shared_ptr<const PreambleData> Preamble =
-        It->second->Worker->getPossiblyStalePreamble();
+        It->second->Worker->getPossiblyStalePreamble(&Signals);
     WithContext WithProvidedContext(Opts.ContextProvider(File));
     Action(InputsAndPreamble{It->second->Contents,
                              It->second->Worker->getCurrentCompileCommand(),
-                             Preamble.get()});
+                             Preamble.get(), Signals.get()});
     return;
   }
 
   std::shared_ptr<const ASTWorker> Worker = It->second->Worker.lock();
-  auto Task =
-      [Worker, Consistency, Name = Name.str(), File = File.str(),
-       Contents = It->second->Contents,
-       Command = Worker->getCurrentCompileCommand(),
-       Ctx = Context::current().derive(kFileBeingProcessed, std::string(File)),
-       Action = std::move(Action), this]() mutable {
-        std::shared_ptr<const PreambleData> Preamble;
-        if (Consistency == PreambleConsistency::Stale) {
-          // Wait until the preamble is built for the first time, if preamble
-          // is required. This avoids extra work of processing the preamble
-          // headers in parallel multiple times.
-          Worker->waitForFirstPreamble();
-        }
-        Preamble = Worker->getPossiblyStalePreamble();
+  auto Task = [Worker, Consistency, Name = Name.str(), File = File.str(),
+               Contents = It->second->Contents,
+               Command = Worker->getCurrentCompileCommand(),
+               Ctx = Context::current().derive(kFileBeingProcessed,
+                                               std::string(File)),
+               Action = std::move(Action), this]() mutable {
+    std::shared_ptr<const PreambleData> Preamble;
+    if (Consistency == PreambleConsistency::Stale) {
+      // Wait until the preamble is built for the first time, if preamble
+      // is required. This avoids extra work of processing the preamble
+      // headers in parallel multiple times.
+      Worker->waitForFirstPreamble();
+    }
+    std::shared_ptr<const ASTSignals> Signals;
+    Preamble = Worker->getPossiblyStalePreamble(&Signals);
 
-        std::lock_guard<Semaphore> BarrierLock(Barrier);
-        WithContext Guard(std::move(Ctx));
-        trace::Span Tracer(Name);
-        SPAN_ATTACH(Tracer, "file", File);
-        WithContext WithProvidedContext(Opts.ContextProvider(File));
-        Action(InputsAndPreamble{Contents, Command, Preamble.get()});
-      };
+    std::lock_guard<Semaphore> BarrierLock(Barrier);
+    WithContext Guard(std::move(Ctx));
+    trace::Span Tracer(Name);
+    SPAN_ATTACH(Tracer, "file", File);
+    WithContext WithProvidedContext(Opts.ContextProvider(File));
+    Action(InputsAndPreamble{Contents, Command, Preamble.get(), Signals.get()});
+  };
 
   PreambleTasks->runAsync("task:" + llvm::sys::path::filename(File),
                           std::move(Task));

clang-tools-extra/clangd/TUScheduler.h

@@ -9,6 +9,7 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 
+#include "ASTSignals.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
@@ -43,6 +44,8 @@ struct InputsAndPreamble {
   const tooling::CompileCommand &Command;
   // This can be nullptr if no preamble is available.
   const PreambleData *Preamble;
+  // This can be nullptr if no ASTSignals are available.
+  const ASTSignals *Signals;
 };
 
 /// Determines whether diagnostics should be generated for a file snapshot.

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

@@ -0,0 +1,75 @@
+//===-- ASTSignalsTests.cpp -------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#include "AST.h"
+
+#include "ParsedAST.h"
+#include "TestIndex.h"
+#include "TestTU.h"
+#include "llvm/ADT/StringRef.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+
+using ::testing::_;
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
+TEST(ASTSignals, Derive) {
+  TestTU TU = TestTU::withCode(R"cpp(
+  namespace ns1 {
+  namespace ns2 {
+  namespace {
+  int func() {
+    tar::X a;
+    a.Y = 1;
+    return ADD(tar::kConst, a.Y, tar::foo()) + fooInNS2() + tar::foo();
+  }
+  } // namespace
+  } // namespace ns2
+  } // namespace ns1
+  )cpp");
+
+  TU.HeaderCode = R"cpp(
+  #define ADD(x, y, z) (x + y + z)
+  namespace tar {  // A related namespace.
+  int kConst = 5;
+  int foo();
+  void bar();  // Unused symbols are not recorded.
+  class X {
+    public: int Y;
+  };
+  } // namespace tar
+  namespace ns1::ns2 { int fooInNS2(); }}
+  )cpp";
+  ASTSignals Signals = ASTSignals::derive(TU.build());
+  std::vector<std::pair<StringRef, int>> NS;
+  for (const auto &P : Signals.RelatedNamespaces)
+    NS.emplace_back(P.getKey(), P.getValue());
+  EXPECT_THAT(NS, UnorderedElementsAre(Pair("ns1::", 1), Pair("ns1::ns2::", 1),
+                                       Pair("tar::", /*foo, kConst, X*/ 3)));
+
+  std::vector<std::pair<SymbolID, int>> Sym;
+  for (const auto &P : Signals.ReferencedSymbols)
+    Sym.emplace_back(P.getFirst(), P.getSecond());
+  EXPECT_THAT(
+      Sym,
+      UnorderedElementsAre(
+          Pair(ns("tar").ID, 4), Pair(ns("ns1").ID, 1),
+          Pair(ns("ns1::ns2").ID, 1), Pair(_ /*int func();*/, 1),
+          Pair(cls("tar::X").ID, 1), Pair(var("tar::kConst").ID, 1),
+          Pair(func("tar::foo").ID, 2), Pair(func("ns1::ns2::fooInNS2").ID, 1),
+          Pair(sym("Y", index::SymbolKind::Variable, "@N@tar@S@X@FI@\\0").ID,
+               2),
+          Pair(_ /*a*/, 3)));
+}
+} // namespace
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/unittests/CMakeLists.txt

@@ -35,6 +35,7 @@ add_custom_target(ClangdUnitTests)
 add_unittest(ClangdUnitTests ClangdTests
   Annotations.cpp
   ASTTests.cpp
+  ASTSignalsTests.cpp
   BackgroundIndexTests.cpp
   CallHierarchyTests.cpp
   CanonicalIncludesTests.cpp

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

@@ -14,6 +14,7 @@
 #include "Preamble.h"
 #include "TUScheduler.h"
 #include "TestFS.h"
+#include "TestIndex.h"
 #include "support/Cancellation.h"
 #include "support/Context.h"
 #include "support/Path.h"
@@ -48,6 +49,7 @@ using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::Field;
 using ::testing::IsEmpty;
+using ::testing::Pair;
 using ::testing::Pointee;
 using ::testing::SizeIs;
 using ::testing::UnorderedElementsAre;
@@ -679,12 +681,12 @@ TEST_F(TUSchedulerTests, EmptyPreamble) {
             cantFail(std::move(Preamble)).Preamble->Preamble.getBounds().Size,
             0u);
       });
-  // Wait for the preamble is being built.
+  // Wait while the preamble is being built.
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
 
   // Update the file which results in an empty preamble.
   S.update(Foo, getInputs(Foo, WithEmptyPreamble), WantDiagnostics::Auto);
-  // Wait for the preamble is being built.
+  // Wait while the preamble is being built.
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   S.runWithPreamble(
       "getEmptyPreamble", Foo, TUScheduler::Stale,
@@ -696,6 +698,47 @@ TEST_F(TUSchedulerTests, EmptyPreamble) {
       });
 }
 
+TEST_F(TUSchedulerTests, ASTSignalsSmokeTests) {
+  TUScheduler S(CDB, optsForTest());
+  auto Foo = testPath("foo.cpp");
+  auto Header = testPath("foo.h");
+
+  FS.Files[Header] = "namespace tar { int foo(); }";
+  const char *Contents = R"cpp(
+  #include "foo.h"
+  namespace ns {
+  int func() {
+    return tar::foo());
+  }
+  } // namespace ns
+  )cpp";
+  // Update the file which results in an empty preamble.
+  S.update(Foo, getInputs(Foo, Contents), WantDiagnostics::Yes);
+  // Wait while the preamble is being built.
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  Notification TaskRun;
+  S.runWithPreamble(
+      "ASTSignals", Foo, TUScheduler::Stale,
+      [&](Expected<InputsAndPreamble> IP) {
+        ASSERT_FALSE(!IP);
+        std::vector<std::pair<StringRef, int>> NS;
+        for (const auto &P : IP->Signals->RelatedNamespaces)
+          NS.emplace_back(P.getKey(), P.getValue());
+        EXPECT_THAT(NS,
+                    UnorderedElementsAre(Pair("ns::", 1), Pair("tar::", 1)));
+
+        std::vector<std::pair<SymbolID, int>> Sym;
+        for (const auto &P : IP->Signals->ReferencedSymbols)
+          Sym.emplace_back(P.getFirst(), P.getSecond());
+        EXPECT_THAT(Sym, UnorderedElementsAre(Pair(ns("tar").ID, 1),
+                                              Pair(ns("ns").ID, 1),
+                                              Pair(func("tar::foo").ID, 1),
+                                              Pair(func("ns::func").ID, 1)));
+        TaskRun.notify();
+      });
+  TaskRun.wait();
+}
+
 TEST_F(TUSchedulerTests, RunWaitsForPreamble) {
   // Testing strategy: we update the file and schedule a few preamble reads at
   // the same time. All reads should get the same non-null preamble.