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

[usaxena95]

Summary

Changes for new AST signals for TUScheduler.

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.

[usaxena95]

Documentation change.

[adamcz]

clang-tools-extra/clangd/TUScheduler.cpp

860

nit: I think you can just use OS.str() here and below instead of RelatedNS to avoid extra string copy.

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

716

bar or tar?

720

Can you add anonymous namespaces as well?

733

s/for/while?

[usaxena95]

Addressed comments.

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

720

Good point. Avoided capturing anonymous namespaces.

[sammccall]

Thanks, this looks pretty solid!

I want to be pretty careful about APIs/naming/code organization here (even moreso than usual!) because TUScheduler is a big complex beast as it is, and it's easy to get lost.

clang-tools-extra/clangd/TUScheduler.cpp

421

In practice, this is always called together with getPossiblyStalePreamble() for runWithPreamble (2 callsites). The two feel fairly related, and they each lock the same mutex.

The other two callsites of getPossiblyStalePreamble() are measuring memory (which in principle should also get the ASTSignals, though in practice it's likely small enough not to matter), and building ASTs, which can ignore it.

I'd suggest combining them into the somewhat awkward signature

shared_ptr<const PreambleData> getPossiblyStalePreamble(shared_ptr<const ASTSignals> *PossiblyStaleASTSignals = nullptr)

423

why is this public?

843

This implementation doesn't belong in TUScheduler, which is all about managing lifecycles and threading, rather than inspecting ASTs.

Best suggestion I have is to create a separate header ASTSignals.h, which defines the ASTSignals struct and a static ASTSignals ASTSignals::derive(const ParsedAST&) function to extract them.

(It's pretty awkward to even directly depend on it, since this feels more of a "feature" like go-to-definition that's independent of to TUScheduler and injected by ClangdServer, rather than core things like ParsedAST that are layered below TUScheduler. But fundamentally we can't avoid this without somehow obfuscating the data structure we're storing).

853

why not?

855

as mentioned above, you may want to do this only if the per-symbol-count was incremented from 0 to 1.
You'd have to ignore namespaces for symbols with no symbolID, but that doesn't matter.

856

why not NSD->isAnonymous()?

(and generally prefer early continue rather than indenting the whole loop body)

857

rather than checking/printing the namespace every time it's encountered, collect a DenseMap<NamespaceDecl*, unsigned> and then run through them once at the end?

859

use printNamespaceScope() from AST.h, which correctly handles inline namespace, adds trailing "::" etc.

868

nit: to minimize the length of the locked section:

• move the new signals into the shared_ptr first (here I think you can just allocate it with make_shared in the first place). This avoids locking the allocation + move constructor.
• introduce a new scope for the locked section, and std::swap the shared_ptrs rather than moving. Together, this avoids locking the destructor of the old signals (it runs when the pointer you swapped it into goes out of scope, instead).

clang-tools-extra/clangd/TUScheduler.h

38

Neither the comment or the name really explain what this is or what's unusual about it.
I can't think of a great name, but the comment should mention:

• the purpose (provide information that can only be extracted from the AST to actions that can't access an AST)
• the limitations (may be absent, always stale)
• example usage (information about what declarations are used in a file affects code-completion ranking)

41

ReferencedSymbols?

42

I have a hard time parsing this comment (what does it mean for a namespace to be present in the file?) Maybe:

/// Namespaces whose symbols are used in the file, and the number of such symbols.

Looking at the implementation, it looks like it's rather the number of symbol *occurrences*. I suspect distinct symbols as your comment here suggests would be better though. If one symbol from a namespace is used 100 times, that should boost that symbol a lot but not necessarily the namespace a lot.

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

738

(untangling ASTSignals from TUScheduler should make this clearer to test directly, though you may still want to keep a simple smoke test here)

[usaxena95]

Addressed comments.

[usaxena95]

clang-tools-extra/clangd/TUScheduler.cpp

843

Moved struct and calculation to a separate file ASTSignals.h/cpp

853

My mental model suggested these must be namespaces from outside the file which are relevant to this file.
I suppose we can take care of this while using this (in CC signals calculation) and let these represent all namespaces including the primary one.

On a second thought: If we do not explicitly exclude primary namespace, the results from primary namespace will be boosted much more than the results from index which might be acceptable too.

855

Ignored NS without a valid SymbolID.
Counted distinct number of Symbols used per namespace.

856

I guess isAnonymous would cover most cases. I thought of excluding named NS in an anonymous NS but those are super rare I suppose and including those is not harmful too.

[usaxena95]

Removed unused headers.

[sammccall]

Thanks!

clang-tools-extra/clangd/ASTSignals.cpp

9 ↗

(On Diff #316669)

this requires storing all the SymbolIDs again.

You could just have a counter instead. The trick is that you're already counting the number of times a symbol has been seen...

unsigned &SymbolCount = Signals.ReferencedSymbols[ID];
++SymbolCount;
if (SymbolCount == 1) { // first time seeing the symbol
  // increment counter for its namespace
}

[usaxena95]

Addressed comments. Ready for landing.

[usaxena95]

Added newline at end of new file.

clang-tools-extra/clangd/TUScheduler.h

 //===--- TUScheduler.h -------------------------------------------*-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_TUSCHEDULER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
 #include "index/CanonicalIncludes.h"
+#include "index/SymbolID.h"
 #include "support/Function.h"
 #include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/Threading.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include <chrono>
 
 namespace clang {
 namespace clangd {
 class ParsedAST;
 struct PreambleData;
 
 /// Returns a number of a default async threads to use for TUScheduler.
 /// Returned value is always >= 1 (i.e. will not cause requests to be processed
 /// synchronously).
 unsigned getDefaultAsyncThreadsCount();
 
+/// Signals derived from a valid AST of a file.
+struct ASTSignals {
+  /// Number of occurrences of each Symbol present in the file.
+  llvm::DenseMap<SymbolID, unsigned> Symbols;
+  /// Number of Symbols belonging to each namespace present in the file.
+  llvm::StringMap<unsigned> RelatedNamespaces;
+};
+
 struct InputsAndAST {
   const ParseInputs &Inputs;
   ParsedAST &AST;
 };
 
 struct InputsAndPreamble {
   llvm::StringRef Contents;
   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.
 enum class WantDiagnostics {
   Yes,  /// Diagnostics must be generated for this snapshot.
   No,   /// Diagnostics must not be generated for this snapshot.
   Auto, /// Diagnostics must be generated for this snapshot or a subsequent one,
         /// within a bounded amount of time.

clang-tools-extra/clangd/TUScheduler.cpp

 // version of the file. This ensures diagnostics get updated even if the queue
 // is full.
 //
 // Some read requests might just need preamble. Since preambles can be read
 // concurrently, ASTWorker runs these requests on their own thread. These
 // requests will receive latest build preamble, which might possibly be stale.
 
 #include "TUScheduler.h"
+#include "AST.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
+#include "FindTarget.h"
 #include "GlobalCompilationDatabase.h"
 #include "ParsedAST.h"
 #include "Preamble.h"
 #include "index/CanonicalIncludes.h"
 #include "support/Cancellation.h"
 #include "support/Context.h"
 #include "support/Logger.h"
 #include "support/MemoryTree.h"
 #include "support/Path.h"
 #include "support/Threading.h"
 #include "support/Trace.h"
+#include "clang/AST/Decl.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallVector.h"
public:
   tooling::CompileCommand getCurrentCompileCommand() const;
 
   /// Wait for the first build of preamble to finish. Preamble itself can be
   /// accessed via getPossiblyStalePreamble(). Note that this function will
   /// return after an unsuccessful build of the preamble too, i.e. result of
   /// getPossiblyStalePreamble() can be null even after this function returns.
   void waitForFirstPreamble() const;
 
+  std::shared_ptr<const ASTSignals> getASTSignals() const;
+
+  /// Used to calculates new ASTSignals from the AST. It is called when
+  /// Diagnostics are generated.
+  void updateASTSignals(ParsedAST &AST);
+
   TUScheduler::FileStats stats() const;
   bool isASTCached() const;
 
 private:
   // Details of an update request that are relevant to scheduling.
   struct UpdateType {
     // Do we want diagnostics from this version?
     // If Yes, we must always build this version.
llvm::SmallVector<DebouncePolicy::clock::duration>
       RebuildTimes; /* GUARDED_BY(Mutex) */
   /// Set to true to signal run() to finish processing.
   bool Done;                              /* GUARDED_BY(Mutex) */
   std::deque<Request> Requests;           /* GUARDED_BY(Mutex) */
   llvm::Optional<Request> CurrentRequest; /* GUARDED_BY(Mutex) */
   /// 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.
   llvm::Optional<std::shared_ptr<const PreambleData>> LatestPreamble;
   std::queue<Request> PreambleRequests; /* GUARDED_BY(Mutex) */
   /// Signaled whenever LatestPreamble changes state or there's a new
   /// PreambleRequest.
   mutable std::condition_variable PreambleCV;
if (RunSync) {
     PreambleRequests.push({std::move(Task), std::move(TaskName),
                            steady_clock::now(), Context::current().clone(),
                            llvm::None, TUScheduler::NoInvalidation, nullptr});
   }
   PreambleCV.notify_all();
   RequestsCV.notify_all();
 }
 
+void ASTWorker::updateASTSignals(ParsedAST &AST) {
+  ASTSignals Signals;
+  const SourceManager &SM = AST.getSourceManager();
+  findExplicitReferences(
+      AST.getASTContext(), [&Signals, &SM](ReferenceLoc Ref) {
+        for (const NamedDecl *ND : Ref.Targets) {
+          if (!isInsideMainFile(Ref.NameLoc, SM))
+            continue;
+          if (auto ID = getSymbolID(ND))
+            Signals.Symbols[ID] += 1;
+          // FIXME: Do not count the primary namespace as a related namespace.
+          // Eg. clang::clangd:: for this file.
+          if (const auto *NSD = dyn_cast<NamespaceDecl>(ND->getDeclContext())) {
+            std::string NS;
+            llvm::raw_string_ostream OS(NS);
+            NSD->printQualifiedName(OS);
+            std::string RelatedNS = OS.str();
+            if (!RelatedNS.empty()) {
+              Signals.RelatedNamespaces[RelatedNS + "::"] += 1;
+            }
+          }
+        }
+      });
+  // 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);
+  LatestASTSignals = std::make_shared<ASTSignals>(std::move(Signals));
+}
+
 void ASTWorker::generateDiagnostics(
     std::unique_ptr<CompilerInvocation> Invocation, ParseInputs Inputs,
     std::vector<Diag> CIDiags) {
   // Tracks ast cache accesses for publishing diags.
   static constexpr trace::Metric ASTAccessForDiag(
       "ast_access_diag", trace::Metric::Counter, "result");
   assert(Invocation);
   assert(LatestPreamble);
auto RunPublish = [&](llvm::function_ref<void()> Publish) {
     // removed, making sure there are not race conditions.
     std::lock_guard<std::mutex> Lock(PublishMu);
     if (CanPublishResults)
       Publish();
   };
   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.
     // FIXME: we might have got more errors while trying to build the
     // AST, surface them too.
     Callbacks.onFailedAST(FileName, Inputs.Version, CIDiags, RunPublish);
   }
 
   // AST might've been built for an older version of the source, as ASTWorker
   // queue raced ahead while we were waiting on the preamble. In that case the
   // queue can't reuse the AST.
   if (InputsAreLatest) {
     RanASTCallback = *AST != nullptr;
     IdleASTs.put(this, std::move(*AST));
   }
 }
 
 std::shared_ptr<const PreambleData>
 ASTWorker::getPossiblyStalePreamble() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   return LatestPreamble ? *LatestPreamble : nullptr;
 }
 
+std::shared_ptr<const ASTSignals> ASTWorker::getASTSignals() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return LatestASTSignals;
+}
+
 void ASTWorker::waitForFirstPreamble() const {
   std::unique_lock<std::mutex> Lock(Mutex);
   PreambleCV.wait(Lock, [this] { return LatestPreamble.hasValue() || Done; });
 }
 
 tooling::CompileCommand ASTWorker::getCurrentCompileCommand() const {
   std::unique_lock<std::mutex> Lock(Mutex);
   return FileInputs.CompileCommand;
if (It == Files.end()) {
     return;
   }
 
   if (!PreambleTasks) {
     trace::Span Tracer(Name);
     SPAN_ATTACH(Tracer, "file", File);
     std::shared_ptr<const PreambleData> Preamble =
         It->second->Worker->getPossiblyStalePreamble();
+    std::shared_ptr<const ASTSignals> Signals =
+        It->second->Worker->getASTSignals();
     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(),
+  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)),
+               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();
+    std::shared_ptr<const ASTSignals> Signals = Worker->getASTSignals();
 
-        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()});
-      };
+    Action(InputsAndPreamble{Contents, Command, Preamble.get(), Signals.get()});
 
   PreambleTasks->runAsync("task:" + llvm::sys::path::filename(File),
                           std::move(Task));
 }
 
 llvm::StringMap<TUScheduler::FileStats> TUScheduler::fileStats() const {
   llvm::StringMap<TUScheduler::FileStats> Result;
   for (const auto &PathAndFile : Files)

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

 //===-- TUSchedulerTests.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 "Annotations.h"
 #include "ClangdServer.h"
 #include "Diagnostics.h"
 #include "Matchers.h"
 #include "ParsedAST.h"
 #include "Preamble.h"
 #include "TUScheduler.h"
 #include "TestFS.h"
+#include "TestIndex.h"
 #include "support/Cancellation.h"
 #include "support/Context.h"
 #include "support/Path.h"
 #include "support/TestTracer.h"
 #include "support/Threading.h"
 #include "support/ThreadsafeFS.h"
 #include "clang/Basic/DiagnosticDriver.h"
 #include "llvm/ADT/ArrayRef.h"

 #include <memory>
 #include <string>
 #include <utility>
 
 namespace clang {
 namespace clangd {
 namespace {
 
+using ::testing::_;
 using ::testing::AnyOf;
 using ::testing::Each;
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::Field;
 using ::testing::IsEmpty;
+using ::testing::Pair;
 using ::testing::Pointee;
 using ::testing::SizeIs;
 using ::testing::UnorderedElementsAre;
 
 MATCHER_P2(TUState, PreambleActivity, ASTActivity, "") {
   if (arg.PreambleActivity != PreambleActivity) {
     *result_listener << "preamblestate is "
                      << static_cast<uint8_t>(arg.PreambleActivity);
S.runWithPreamble(
       [&](Expected<InputsAndPreamble> Preamble) {
         // We expect to get an empty preamble.
         EXPECT_EQ(
             cantFail(std::move(Preamble)).Preamble->Preamble.getBounds().Size,
             0u);
       });
 }
 
+TEST_F(TUSchedulerTests, ASTSignals) {
+  TUScheduler S(CDB, optsForTest());
+  auto Foo = testPath("foo.cpp");
+  auto Header = testPath("foo.h");
+
+  FS.Files[Header] = R"cpp(
+  #define ADD(x, y, z) (x + y + z)
+  namespace tar {  // A related namespace.
+  int kConst = 5;
+  int foo();
+  void bar();
+  class X {
+    public: int Y;
+  };
+  } // namespace bar
+  )cpp";
+  const char *Contents = R"cpp(
+  #include "foo.h"
+  namespace ns1 {
+  namespace ns2 {
+    void func() {
+      tar::X a;
+      tar::X b;
+      a.Y = 1;
+      b.Y = ADD(tar::kConst, a.Y, tar::foo());
+    }
+  } // namespace ns2
+  } // namespace ns1
+  )cpp";
+  // Update the file which results in an empty preamble.
+  S.update(Foo, getInputs(Foo, Contents), WantDiagnostics::Yes);
+  // Wait for 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("ns1::", 1),
+                                             Pair("ns1::ns2::", 1),
+                                             Pair("tar::", 4)));
+
+        std::vector<std::pair<SymbolID, int>> Sym;
+        for (const auto &P : IP->Signals->Symbols)
+          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(func("ns1::ns2::func").ID, 1),
+                Pair(cls("tar::X").ID, 2), Pair(var("tar::kConst").ID, 1),
+                Pair(func("tar::foo").ID, 1),
+                Pair(sym("Y", index::SymbolKind::Variable, "@N@tar@S@X@FI@\\0")
+                         .ID,
+                     3),
+                Pair(_ /*a*/, 3), Pair(_ /*b*/, 2)));
+        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.
   TUScheduler S(CDB, optsForTest());
   auto Foo = testPath("foo.cpp");
   auto NonEmptyPreamble = R"cpp(
     #define FOO 1
     #define BAR 2