xref: /llvm-project/clang-tools-extra/clangd/ClangdServer.cpp (revision 61fe67a4017375fd675f75652e857e837f77fa51) 
1 //===--- ClangdServer.cpp - Main clangd server code --------------*- C++-*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===-------------------------------------------------------------------===//
8 
9 #include "ClangdServer.h"
10 #include "CodeComplete.h"
11 #include "Config.h"
12 #include "Diagnostics.h"
13 #include "DumpAST.h"
14 #include "FindSymbols.h"
15 #include "Format.h"
16 #include "HeaderSourceSwitch.h"
17 #include "InlayHints.h"
18 #include "ParsedAST.h"
19 #include "Preamble.h"
20 #include "Protocol.h"
21 #include "SemanticHighlighting.h"
22 #include "SemanticSelection.h"
23 #include "SourceCode.h"
24 #include "TUScheduler.h"
25 #include "XRefs.h"
26 #include "clang-include-cleaner/Record.h"
27 #include "index/FileIndex.h"
28 #include "index/Merge.h"
29 #include "index/StdLib.h"
30 #include "refactor/Rename.h"
31 #include "refactor/Tweak.h"
32 #include "support/Cancellation.h"
33 #include "support/Context.h"
34 #include "support/Logger.h"
35 #include "support/MemoryTree.h"
36 #include "support/ThreadsafeFS.h"
37 #include "support/Trace.h"
38 #include "clang/Basic/Stack.h"
39 #include "clang/Format/Format.h"
40 #include "clang/Lex/Preprocessor.h"
41 #include "clang/Tooling/CompilationDatabase.h"
42 #include "clang/Tooling/Core/Replacement.h"
43 #include "llvm/ADT/ArrayRef.h"
44 #include "llvm/ADT/STLExtras.h"
45 #include "llvm/ADT/StringRef.h"
46 #include "llvm/Support/Error.h"
47 #include "llvm/Support/Path.h"
48 #include "llvm/Support/raw_ostream.h"
49 #include <algorithm>
50 #include <chrono>
51 #include <future>
52 #include <memory>
53 #include <mutex>
54 #include <optional>
55 #include <string>
56 #include <type_traits>
57 #include <utility>
58 #include <vector>
59 
60 namespace clang {
61 namespace clangd {
62 namespace {
63 
64 // Tracks number of times a tweak has been offered.
65 static constexpr trace::Metric TweakAvailable(
66     "tweak_available", trace::Metric::Counter, "tweak_id");
67 
68 // Update the FileIndex with new ASTs and plumb the diagnostics responses.
69 struct UpdateIndexCallbacks : public ParsingCallbacks {
70   UpdateIndexCallbacks(FileIndex *FIndex,
71                        ClangdServer::Callbacks *ServerCallbacks,
72                        const ThreadsafeFS &TFS, AsyncTaskRunner *Tasks,
73                        bool CollectInactiveRegions)
74       : FIndex(FIndex), ServerCallbacks(ServerCallbacks), TFS(TFS),
75         Stdlib{std::make_shared<StdLibSet>()}, Tasks(Tasks),
76         CollectInactiveRegions(CollectInactiveRegions) {}
77 
78   void onPreambleAST(
79       PathRef Path, llvm::StringRef Version, CapturedASTCtx ASTCtx,
80       std::shared_ptr<const include_cleaner::PragmaIncludes> PI) override {
81 
82     if (!FIndex)
83       return;
84 
85     auto &PP = ASTCtx.getPreprocessor();
86     auto &CI = ASTCtx.getCompilerInvocation();
87     if (auto Loc = Stdlib->add(CI.getLangOpts(), PP.getHeaderSearchInfo()))
88       indexStdlib(CI, std::move(*Loc));
89 
90     // FIndex outlives the UpdateIndexCallbacks.
91     auto Task = [FIndex(FIndex), Path(Path.str()), Version(Version.str()),
92                  ASTCtx(std::move(ASTCtx)), PI(std::move(PI))]() mutable {
93       trace::Span Tracer("PreambleIndexing");
94       FIndex->updatePreamble(Path, Version, ASTCtx.getASTContext(),
95                              ASTCtx.getPreprocessor(), *PI);
96     };
97 
98     if (Tasks) {
99       Tasks->runAsync("Preamble indexing for:" + Path + Version,
100                       std::move(Task));
101     } else
102       Task();
103   }
104 
105   void indexStdlib(const CompilerInvocation &CI, StdLibLocation Loc) {
106     // This task is owned by Tasks, which outlives the TUScheduler and
107     // therefore the UpdateIndexCallbacks.
108     // We must be careful that the references we capture outlive TUScheduler.
109     auto Task = [LO(CI.getLangOpts()), Loc(std::move(Loc)),
110                  CI(std::make_unique<CompilerInvocation>(CI)),
111                  // External values that outlive ClangdServer
112                  TFS(&TFS),
113                  // Index outlives TUScheduler (declared first)
114                  FIndex(FIndex),
115                  // shared_ptr extends lifetime
116                  Stdlib(Stdlib),
117                  // We have some FS implementations that rely on information in
118                  // the context.
119                  Ctx(Context::current().clone())]() mutable {
120       // Make sure we install the context into current thread.
121       WithContext C(std::move(Ctx));
122       clang::noteBottomOfStack();
123       IndexFileIn IF;
124       IF.Symbols = indexStandardLibrary(std::move(CI), Loc, *TFS);
125       if (Stdlib->isBest(LO))
126         FIndex->updatePreamble(std::move(IF));
127     };
128     if (Tasks)
129       // This doesn't have a semaphore to enforce -j, but it's rare.
130       Tasks->runAsync("IndexStdlib", std::move(Task));
131     else
132       Task();
133   }
134 
135   void onMainAST(PathRef Path, ParsedAST &AST, PublishFn Publish) override {
136     if (FIndex)
137       FIndex->updateMain(Path, AST);
138 
139     if (ServerCallbacks)
140       Publish([&]() {
141         ServerCallbacks->onDiagnosticsReady(Path, AST.version(),
142                                             AST.getDiagnostics());
143         if (CollectInactiveRegions) {
144           ServerCallbacks->onInactiveRegionsReady(Path,
145                                                   getInactiveRegions(AST));
146         }
147       });
148   }
149 
150   void onFailedAST(PathRef Path, llvm::StringRef Version,
151                    std::vector<Diag> Diags, PublishFn Publish) override {
152     if (ServerCallbacks)
153       Publish(
154           [&]() { ServerCallbacks->onDiagnosticsReady(Path, Version, Diags); });
155   }
156 
157   void onFileUpdated(PathRef File, const TUStatus &Status) override {
158     if (ServerCallbacks)
159       ServerCallbacks->onFileUpdated(File, Status);
160   }
161 
162   void onPreamblePublished(PathRef File) override {
163     if (ServerCallbacks)
164       ServerCallbacks->onSemanticsMaybeChanged(File);
165   }
166 
167 private:
168   FileIndex *FIndex;
169   ClangdServer::Callbacks *ServerCallbacks;
170   const ThreadsafeFS &TFS;
171   std::shared_ptr<StdLibSet> Stdlib;
172   AsyncTaskRunner *Tasks;
173   bool CollectInactiveRegions;
174 };
175 
176 class DraftStoreFS : public ThreadsafeFS {
177 public:
178   DraftStoreFS(const ThreadsafeFS &Base, const DraftStore &Drafts)
179       : Base(Base), DirtyFiles(Drafts) {}
180 
181 private:
182   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> viewImpl() const override {
183     auto OFS = llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(
184         Base.view(std::nullopt));
185     OFS->pushOverlay(DirtyFiles.asVFS());
186     return OFS;
187   }
188 
189   const ThreadsafeFS &Base;
190   const DraftStore &DirtyFiles;
191 };
192 
193 } // namespace
194 
195 ClangdServer::Options ClangdServer::optsForTest() {
196   ClangdServer::Options Opts;
197   Opts.UpdateDebounce = DebouncePolicy::fixed(/*zero*/ {});
198   Opts.StorePreamblesInMemory = true;
199   Opts.AsyncThreadsCount = 4; // Consistent!
200   return Opts;
201 }
202 
203 ClangdServer::Options::operator TUScheduler::Options() const {
204   TUScheduler::Options Opts;
205   Opts.AsyncThreadsCount = AsyncThreadsCount;
206   Opts.RetentionPolicy = RetentionPolicy;
207   Opts.StorePreamblesInMemory = StorePreamblesInMemory;
208   Opts.UpdateDebounce = UpdateDebounce;
209   Opts.ContextProvider = ContextProvider;
210   Opts.PreambleThrottler = PreambleThrottler;
211   return Opts;
212 }
213 
214 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
215                            const ThreadsafeFS &TFS, const Options &Opts,
216                            Callbacks *Callbacks)
217     : FeatureModules(Opts.FeatureModules), CDB(CDB), TFS(TFS),
218       DynamicIdx(Opts.BuildDynamicSymbolIndex
219                      ? new FileIndex(Opts.EnableOutgoingCalls)
220                      : nullptr),
221       ModulesManager(Opts.ModulesManager),
222       ClangTidyProvider(Opts.ClangTidyProvider),
223       UseDirtyHeaders(Opts.UseDirtyHeaders),
224       LineFoldingOnly(Opts.LineFoldingOnly),
225       PreambleParseForwardingFunctions(Opts.PreambleParseForwardingFunctions),
226       ImportInsertions(Opts.ImportInsertions),
227       PublishInactiveRegions(Opts.PublishInactiveRegions),
228       WorkspaceRoot(Opts.WorkspaceRoot),
229       Transient(Opts.ImplicitCancellation ? TUScheduler::InvalidateOnUpdate
230                                           : TUScheduler::NoInvalidation),
231       DirtyFS(std::make_unique<DraftStoreFS>(TFS, DraftMgr)) {
232   if (Opts.AsyncThreadsCount != 0)
233     IndexTasks.emplace();
234   // Pass a callback into `WorkScheduler` to extract symbols from a newly
235   // parsed file and rebuild the file index synchronously each time an AST
236   // is parsed.
237   WorkScheduler.emplace(CDB, TUScheduler::Options(Opts),
238                         std::make_unique<UpdateIndexCallbacks>(
239                             DynamicIdx.get(), Callbacks, TFS,
240                             IndexTasks ? &*IndexTasks : nullptr,
241                             PublishInactiveRegions));
242   // Adds an index to the stack, at higher priority than existing indexes.
243   auto AddIndex = [&](SymbolIndex *Idx) {
244     if (this->Index != nullptr) {
245       MergedIdx.push_back(std::make_unique<MergedIndex>(Idx, this->Index));
246       this->Index = MergedIdx.back().get();
247     } else {
248       this->Index = Idx;
249     }
250   };
251   if (Opts.StaticIndex)
252     AddIndex(Opts.StaticIndex);
253   if (Opts.BackgroundIndex) {
254     BackgroundIndex::Options BGOpts;
255     BGOpts.ThreadPoolSize = std::max(Opts.AsyncThreadsCount, 1u);
256     BGOpts.OnProgress = [Callbacks](BackgroundQueue::Stats S) {
257       if (Callbacks)
258         Callbacks->onBackgroundIndexProgress(S);
259     };
260     BGOpts.ContextProvider = Opts.ContextProvider;
261     BGOpts.SupportContainedRefs = Opts.EnableOutgoingCalls;
262     BackgroundIdx = std::make_unique<BackgroundIndex>(
263         TFS, CDB,
264         BackgroundIndexStorage::createDiskBackedStorageFactory(
265             [&CDB](llvm::StringRef File) { return CDB.getProjectInfo(File); }),
266         std::move(BGOpts));
267     AddIndex(BackgroundIdx.get());
268   }
269   if (DynamicIdx)
270     AddIndex(DynamicIdx.get());
271 
272   if (Opts.FeatureModules) {
273     FeatureModule::Facilities F{
274         *this->WorkScheduler,
275         this->Index,
276         this->TFS,
277     };
278     for (auto &Mod : *Opts.FeatureModules)
279       Mod.initialize(F);
280   }
281 }
282 
283 ClangdServer::~ClangdServer() {
284   // Destroying TUScheduler first shuts down request threads that might
285   // otherwise access members concurrently.
286   // (Nobody can be using TUScheduler because we're on the main thread).
287   WorkScheduler.reset();
288   // Now requests have stopped, we can shut down feature modules.
289   if (FeatureModules) {
290     for (auto &Mod : *FeatureModules)
291       Mod.stop();
292     for (auto &Mod : *FeatureModules)
293       Mod.blockUntilIdle(Deadline::infinity());
294   }
295 }
296 
297 void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
298                                llvm::StringRef Version,
299                                WantDiagnostics WantDiags, bool ForceRebuild) {
300   std::string ActualVersion = DraftMgr.addDraft(File, Version, Contents);
301   ParseOptions Opts;
302   Opts.PreambleParseForwardingFunctions = PreambleParseForwardingFunctions;
303   Opts.ImportInsertions = ImportInsertions;
304 
305   // Compile command is set asynchronously during update, as it can be slow.
306   ParseInputs Inputs;
307   Inputs.TFS = &getHeaderFS();
308   Inputs.Contents = std::string(Contents);
309   Inputs.Version = std::move(ActualVersion);
310   Inputs.ForceRebuild = ForceRebuild;
311   Inputs.Opts = std::move(Opts);
312   Inputs.Index = Index;
313   Inputs.ClangTidyProvider = ClangTidyProvider;
314   Inputs.FeatureModules = FeatureModules;
315   Inputs.ModulesManager = ModulesManager;
316   bool NewFile = WorkScheduler->update(File, Inputs, WantDiags);
317   // If we loaded Foo.h, we want to make sure Foo.cpp is indexed.
318   if (NewFile && BackgroundIdx)
319     BackgroundIdx->boostRelated(File);
320 }
321 
322 void ClangdServer::reparseOpenFilesIfNeeded(
323     llvm::function_ref<bool(llvm::StringRef File)> Filter) {
324   // Reparse only opened files that were modified.
325   for (const Path &FilePath : DraftMgr.getActiveFiles())
326     if (Filter(FilePath))
327       if (auto Draft = DraftMgr.getDraft(FilePath)) // else disappeared in race?
328         addDocument(FilePath, *Draft->Contents, Draft->Version,
329                     WantDiagnostics::Auto);
330 }
331 
332 std::shared_ptr<const std::string> ClangdServer::getDraft(PathRef File) const {
333   auto Draft = DraftMgr.getDraft(File);
334   if (!Draft)
335     return nullptr;
336   return std::move(Draft->Contents);
337 }
338 
339 std::function<Context(PathRef)>
340 ClangdServer::createConfiguredContextProvider(const config::Provider *Provider,
341                                               Callbacks *Publish) {
342   if (!Provider)
343     return [](llvm::StringRef) { return Context::current().clone(); };
344 
345   struct Impl {
346     const config::Provider *Provider;
347     ClangdServer::Callbacks *Publish;
348     std::mutex PublishMu;
349 
350     Impl(const config::Provider *Provider, ClangdServer::Callbacks *Publish)
351         : Provider(Provider), Publish(Publish) {}
352 
353     Context operator()(llvm::StringRef File) {
354       config::Params Params;
355       // Don't reread config files excessively often.
356       // FIXME: when we see a config file change event, use the event timestamp?
357       Params.FreshTime =
358           std::chrono::steady_clock::now() - std::chrono::seconds(5);
359       llvm::SmallString<256> PosixPath;
360       if (!File.empty()) {
361         assert(llvm::sys::path::is_absolute(File));
362         llvm::sys::path::native(File, PosixPath, llvm::sys::path::Style::posix);
363         Params.Path = PosixPath.str();
364       }
365 
366       llvm::StringMap<std::vector<Diag>> ReportableDiagnostics;
367       Config C = Provider->getConfig(Params, [&](const llvm::SMDiagnostic &D) {
368         // Create the map entry even for note diagnostics we don't report.
369         // This means that when the file is parsed with no warnings, we
370         // publish an empty set of diagnostics, clearing any the client has.
371         handleDiagnostic(D, !Publish || D.getFilename().empty()
372                                 ? nullptr
373                                 : &ReportableDiagnostics[D.getFilename()]);
374       });
375       // Blindly publish diagnostics for the (unopened) parsed config files.
376       // We must avoid reporting diagnostics for *the same file* concurrently.
377       // Source diags are published elsewhere, but those are different files.
378       if (!ReportableDiagnostics.empty()) {
379         std::lock_guard<std::mutex> Lock(PublishMu);
380         for (auto &Entry : ReportableDiagnostics)
381           Publish->onDiagnosticsReady(Entry.first(), /*Version=*/"",
382                                       Entry.second);
383       }
384       return Context::current().derive(Config::Key, std::move(C));
385     }
386 
387     void handleDiagnostic(const llvm::SMDiagnostic &D,
388                           std::vector<Diag> *ClientDiagnostics) {
389       switch (D.getKind()) {
390       case llvm::SourceMgr::DK_Error:
391         elog("config error at {0}:{1}:{2}: {3}", D.getFilename(), D.getLineNo(),
392              D.getColumnNo(), D.getMessage());
393         break;
394       case llvm::SourceMgr::DK_Warning:
395         log("config warning at {0}:{1}:{2}: {3}", D.getFilename(),
396             D.getLineNo(), D.getColumnNo(), D.getMessage());
397         break;
398       case llvm::SourceMgr::DK_Note:
399       case llvm::SourceMgr::DK_Remark:
400         vlog("config note at {0}:{1}:{2}: {3}", D.getFilename(), D.getLineNo(),
401              D.getColumnNo(), D.getMessage());
402         ClientDiagnostics = nullptr; // Don't emit notes as LSP diagnostics.
403         break;
404       }
405       if (ClientDiagnostics)
406         ClientDiagnostics->push_back(toDiag(D, Diag::ClangdConfig));
407     }
408   };
409 
410   // Copyable wrapper.
411   return [I(std::make_shared<Impl>(Provider, Publish))](llvm::StringRef Path) {
412     return (*I)(Path);
413   };
414 }
415 
416 void ClangdServer::removeDocument(PathRef File) {
417   DraftMgr.removeDraft(File);
418   WorkScheduler->remove(File);
419 }
420 
421 void ClangdServer::codeComplete(PathRef File, Position Pos,
422                                 const clangd::CodeCompleteOptions &Opts,
423                                 Callback<CodeCompleteResult> CB) {
424   // Copy completion options for passing them to async task handler.
425   auto CodeCompleteOpts = Opts;
426   if (!CodeCompleteOpts.Index) // Respect overridden index.
427     CodeCompleteOpts.Index = Index;
428 
429   auto Task = [Pos, CodeCompleteOpts, File = File.str(), CB = std::move(CB),
430                this](llvm::Expected<InputsAndPreamble> IP) mutable {
431     if (!IP)
432       return CB(IP.takeError());
433     if (auto Reason = isCancelled())
434       return CB(llvm::make_error<CancelledError>(Reason));
435 
436     std::optional<SpeculativeFuzzyFind> SpecFuzzyFind;
437     if (!IP->Preamble) {
438       // No speculation in Fallback mode, as it's supposed to be much faster
439       // without compiling.
440       vlog("Build for file {0} is not ready. Enter fallback mode.", File);
441     } else if (CodeCompleteOpts.Index) {
442       SpecFuzzyFind.emplace();
443       {
444         std::lock_guard<std::mutex> Lock(CachedCompletionFuzzyFindRequestMutex);
445         SpecFuzzyFind->CachedReq = CachedCompletionFuzzyFindRequestByFile[File];
446       }
447     }
448     ParseInputs ParseInput{IP->Command, &getHeaderFS(), IP->Contents.str()};
449     // FIXME: Add traling new line if there is none at eof, workaround a crash,
450     // see https://github.com/clangd/clangd/issues/332
451     if (!IP->Contents.ends_with("\n"))
452       ParseInput.Contents.append("\n");
453     ParseInput.Index = Index;
454 
455     CodeCompleteOpts.MainFileSignals = IP->Signals;
456     CodeCompleteOpts.AllScopes = Config::current().Completion.AllScopes;
457     CodeCompleteOpts.ArgumentLists = Config::current().Completion.ArgumentLists;
458     // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
459     // both the old and the new version in case only one of them matches.
460     CodeCompleteResult Result = clangd::codeComplete(
461         File, Pos, IP->Preamble, ParseInput, CodeCompleteOpts,
462         SpecFuzzyFind ? &*SpecFuzzyFind : nullptr);
463     {
464       clang::clangd::trace::Span Tracer("Completion results callback");
465       CB(std::move(Result));
466     }
467     if (SpecFuzzyFind && SpecFuzzyFind->NewReq) {
468       std::lock_guard<std::mutex> Lock(CachedCompletionFuzzyFindRequestMutex);
469       CachedCompletionFuzzyFindRequestByFile[File] = *SpecFuzzyFind->NewReq;
470     }
471     // SpecFuzzyFind is only destroyed after speculative fuzzy find finishes.
472     // We don't want `codeComplete` to wait for the async call if it doesn't use
473     // the result (e.g. non-index completion, speculation fails), so that `CB`
474     // is called as soon as results are available.
475   };
476 
477   // We use a potentially-stale preamble because latency is critical here.
478   WorkScheduler->runWithPreamble(
479       "CodeComplete", File,
480       (Opts.RunParser == CodeCompleteOptions::AlwaysParse)
481           ? TUScheduler::Stale
482           : TUScheduler::StaleOrAbsent,
483       std::move(Task));
484 }
485 
486 void ClangdServer::signatureHelp(PathRef File, Position Pos,
487                                  MarkupKind DocumentationFormat,
488                                  Callback<SignatureHelp> CB) {
489 
490   auto Action = [Pos, File = File.str(), CB = std::move(CB),
491                  DocumentationFormat,
492                  this](llvm::Expected<InputsAndPreamble> IP) mutable {
493     if (!IP)
494       return CB(IP.takeError());
495 
496     const auto *PreambleData = IP->Preamble;
497     if (!PreambleData)
498       return CB(error("Failed to parse includes"));
499 
500     ParseInputs ParseInput{IP->Command, &getHeaderFS(), IP->Contents.str()};
501     // FIXME: Add traling new line if there is none at eof, workaround a crash,
502     // see https://github.com/clangd/clangd/issues/332
503     if (!IP->Contents.ends_with("\n"))
504       ParseInput.Contents.append("\n");
505     ParseInput.Index = Index;
506     CB(clangd::signatureHelp(File, Pos, *PreambleData, ParseInput,
507                              DocumentationFormat));
508   };
509 
510   // Unlike code completion, we wait for a preamble here.
511   WorkScheduler->runWithPreamble("SignatureHelp", File, TUScheduler::Stale,
512                                  std::move(Action));
513 }
514 
515 void ClangdServer::formatFile(PathRef File, std::optional<Range> Rng,
516                               Callback<tooling::Replacements> CB) {
517   auto Code = getDraft(File);
518   if (!Code)
519     return CB(llvm::make_error<LSPError>("trying to format non-added document",
520                                          ErrorCode::InvalidParams));
521   tooling::Range RequestedRange;
522   if (Rng) {
523     llvm::Expected<size_t> Begin = positionToOffset(*Code, Rng->start);
524     if (!Begin)
525       return CB(Begin.takeError());
526     llvm::Expected<size_t> End = positionToOffset(*Code, Rng->end);
527     if (!End)
528       return CB(End.takeError());
529     RequestedRange = tooling::Range(*Begin, *End - *Begin);
530   } else {
531     RequestedRange = tooling::Range(0, Code->size());
532   }
533 
534   // Call clang-format.
535   auto Action = [File = File.str(), Code = std::move(*Code),
536                  Ranges = std::vector<tooling::Range>{RequestedRange},
537                  CB = std::move(CB), this]() mutable {
538     format::FormatStyle Style = getFormatStyleForFile(File, Code, TFS, true);
539     tooling::Replacements IncludeReplaces =
540         format::sortIncludes(Style, Code, Ranges, File);
541     auto Changed = tooling::applyAllReplacements(Code, IncludeReplaces);
542     if (!Changed)
543       return CB(Changed.takeError());
544 
545     CB(IncludeReplaces.merge(format::reformat(
546         Style, *Changed,
547         tooling::calculateRangesAfterReplacements(IncludeReplaces, Ranges),
548         File)));
549   };
550   WorkScheduler->runQuick("Format", File, std::move(Action));
551 }
552 
553 void ClangdServer::formatOnType(PathRef File, Position Pos,
554                                 StringRef TriggerText,
555                                 Callback<std::vector<TextEdit>> CB) {
556   auto Code = getDraft(File);
557   if (!Code)
558     return CB(llvm::make_error<LSPError>("trying to format non-added document",
559                                          ErrorCode::InvalidParams));
560   llvm::Expected<size_t> CursorPos = positionToOffset(*Code, Pos);
561   if (!CursorPos)
562     return CB(CursorPos.takeError());
563   auto Action = [File = File.str(), Code = std::move(*Code),
564                  TriggerText = TriggerText.str(), CursorPos = *CursorPos,
565                  CB = std::move(CB), this]() mutable {
566     auto Style = getFormatStyleForFile(File, Code, TFS, false);
567     std::vector<TextEdit> Result;
568     for (const tooling::Replacement &R :
569          formatIncremental(Code, CursorPos, TriggerText, Style))
570       Result.push_back(replacementToEdit(Code, R));
571     return CB(Result);
572   };
573   WorkScheduler->runQuick("FormatOnType", File, std::move(Action));
574 }
575 
576 void ClangdServer::prepareRename(PathRef File, Position Pos,
577                                  std::optional<std::string> NewName,
578                                  const RenameOptions &RenameOpts,
579                                  Callback<RenameResult> CB) {
580   auto Action = [Pos, File = File.str(), CB = std::move(CB),
581                  NewName = std::move(NewName),
582                  RenameOpts](llvm::Expected<InputsAndAST> InpAST) mutable {
583     if (!InpAST)
584       return CB(InpAST.takeError());
585     // prepareRename is latency-sensitive: we don't query the index, as we
586     // only need main-file references
587     auto Results =
588         clangd::rename({Pos, NewName.value_or("__clangd_rename_placeholder"),
589                         InpAST->AST, File, /*FS=*/nullptr,
590                         /*Index=*/nullptr, RenameOpts});
591     if (!Results) {
592       // LSP says to return null on failure, but that will result in a generic
593       // failure message. If we send an LSP error response, clients can surface
594       // the message to users (VSCode does).
595       return CB(Results.takeError());
596     }
597     return CB(*Results);
598   };
599   WorkScheduler->runWithAST("PrepareRename", File, std::move(Action));
600 }
601 
602 void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
603                           const RenameOptions &Opts,
604                           Callback<RenameResult> CB) {
605   auto Action = [File = File.str(), NewName = NewName.str(), Pos, Opts,
606                  CB = std::move(CB),
607                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
608     // Tracks number of files edited per invocation.
609     static constexpr trace::Metric RenameFiles("rename_files",
610                                                trace::Metric::Distribution);
611     if (!InpAST)
612       return CB(InpAST.takeError());
613     auto R = clangd::rename({Pos, NewName, InpAST->AST, File,
614                              DirtyFS->view(std::nullopt), Index, Opts});
615     if (!R)
616       return CB(R.takeError());
617 
618     if (Opts.WantFormat) {
619       auto Style = getFormatStyleForFile(File, InpAST->Inputs.Contents,
620                                          *InpAST->Inputs.TFS, false);
621       llvm::Error Err = llvm::Error::success();
622       for (auto &E : R->GlobalChanges)
623         Err =
624             llvm::joinErrors(reformatEdit(E.getValue(), Style), std::move(Err));
625 
626       if (Err)
627         return CB(std::move(Err));
628     }
629     RenameFiles.record(R->GlobalChanges.size());
630     return CB(*R);
631   };
632   WorkScheduler->runWithAST("Rename", File, std::move(Action));
633 }
634 
635 namespace {
636 // May generate several candidate selections, due to SelectionTree ambiguity.
637 // vector of pointers because GCC doesn't like non-copyable Selection.
638 llvm::Expected<std::vector<std::unique_ptr<Tweak::Selection>>>
639 tweakSelection(const Range &Sel, const InputsAndAST &AST,
640                llvm::vfs::FileSystem *FS) {
641   auto Begin = positionToOffset(AST.Inputs.Contents, Sel.start);
642   if (!Begin)
643     return Begin.takeError();
644   auto End = positionToOffset(AST.Inputs.Contents, Sel.end);
645   if (!End)
646     return End.takeError();
647   std::vector<std::unique_ptr<Tweak::Selection>> Result;
648   SelectionTree::createEach(
649       AST.AST.getASTContext(), AST.AST.getTokens(), *Begin, *End,
650       [&](SelectionTree T) {
651         Result.push_back(std::make_unique<Tweak::Selection>(
652             AST.Inputs.Index, AST.AST, *Begin, *End, std::move(T), FS));
653         return false;
654       });
655   assert(!Result.empty() && "Expected at least one SelectionTree");
656   return std::move(Result);
657 }
658 
659 // Some fixes may perform local renaming, we want to convert those to clangd
660 // rename commands, such that we can leverage the index for more accurate
661 // results.
662 std::optional<ClangdServer::CodeActionResult::Rename>
663 tryConvertToRename(const Diag *Diag, const Fix &Fix) {
664   bool IsClangTidyRename = Diag->Source == Diag::ClangTidy &&
665                            Diag->Name == "readability-identifier-naming" &&
666                            !Fix.Edits.empty();
667   if (IsClangTidyRename && Diag->InsideMainFile) {
668     ClangdServer::CodeActionResult::Rename R;
669     R.NewName = Fix.Edits.front().newText;
670     R.FixMessage = Fix.Message;
671     R.Diag = {Diag->Range, Diag->Message};
672     return R;
673   }
674 
675   return std::nullopt;
676 }
677 
678 } // namespace
679 
680 void ClangdServer::codeAction(const CodeActionInputs &Params,
681                               Callback<CodeActionResult> CB) {
682   auto Action = [Params, CB = std::move(CB),
683                  FeatureModules(this->FeatureModules)](
684                     Expected<InputsAndAST> InpAST) mutable {
685     if (!InpAST)
686       return CB(InpAST.takeError());
687     auto KindAllowed =
688         [Only(Params.RequestedActionKinds)](llvm::StringRef Kind) {
689           if (Only.empty())
690             return true;
691           return llvm::any_of(Only, [&](llvm::StringRef Base) {
692             return Kind.consume_front(Base) &&
693                    (Kind.empty() || Kind.starts_with("."));
694           });
695         };
696 
697     CodeActionResult Result;
698     Result.Version = InpAST->AST.version().str();
699     if (KindAllowed(CodeAction::QUICKFIX_KIND)) {
700       auto FindMatchedDiag = [&InpAST](const DiagRef &DR) -> const Diag * {
701         for (const auto &Diag : InpAST->AST.getDiagnostics())
702           if (Diag.Range == DR.Range && Diag.Message == DR.Message)
703             return &Diag;
704         return nullptr;
705       };
706       for (const auto &DiagRef : Params.Diagnostics) {
707         if (const auto *Diag = FindMatchedDiag(DiagRef))
708           for (const auto &Fix : Diag->Fixes) {
709             if (auto Rename = tryConvertToRename(Diag, Fix)) {
710               Result.Renames.emplace_back(std::move(*Rename));
711             } else {
712               Result.QuickFixes.push_back({DiagRef, Fix});
713             }
714           }
715       }
716     }
717 
718     // Collect Tweaks
719     auto Selections = tweakSelection(Params.Selection, *InpAST, /*FS=*/nullptr);
720     if (!Selections)
721       return CB(Selections.takeError());
722     // Don't allow a tweak to fire more than once across ambiguous selections.
723     llvm::DenseSet<llvm::StringRef> PreparedTweaks;
724     auto DeduplicatingFilter = [&](const Tweak &T) {
725       return KindAllowed(T.kind()) && Params.TweakFilter(T) &&
726              !PreparedTweaks.count(T.id());
727     };
728     for (const auto &Sel : *Selections) {
729       for (auto &T : prepareTweaks(*Sel, DeduplicatingFilter, FeatureModules)) {
730         Result.TweakRefs.push_back(TweakRef{T->id(), T->title(), T->kind()});
731         PreparedTweaks.insert(T->id());
732         TweakAvailable.record(1, T->id());
733       }
734     }
735     CB(std::move(Result));
736   };
737 
738   WorkScheduler->runWithAST("codeAction", Params.File, std::move(Action),
739                             Transient);
740 }
741 
742 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
743                               Callback<Tweak::Effect> CB) {
744   // Tracks number of times a tweak has been attempted.
745   static constexpr trace::Metric TweakAttempt(
746       "tweak_attempt", trace::Metric::Counter, "tweak_id");
747   // Tracks number of times a tweak has failed to produce edits.
748   static constexpr trace::Metric TweakFailed(
749       "tweak_failed", trace::Metric::Counter, "tweak_id");
750   TweakAttempt.record(1, TweakID);
751   auto Action = [File = File.str(), Sel, TweakID = TweakID.str(),
752                  CB = std::move(CB),
753                  this](Expected<InputsAndAST> InpAST) mutable {
754     if (!InpAST)
755       return CB(InpAST.takeError());
756     auto FS = DirtyFS->view(std::nullopt);
757     auto Selections = tweakSelection(Sel, *InpAST, FS.get());
758     if (!Selections)
759       return CB(Selections.takeError());
760     std::optional<llvm::Expected<Tweak::Effect>> Effect;
761     // Try each selection, take the first one that prepare()s.
762     // If they all fail, Effect will hold get the last error.
763     for (const auto &Selection : *Selections) {
764       auto T = prepareTweak(TweakID, *Selection, FeatureModules);
765       if (T) {
766         Effect = (*T)->apply(*Selection);
767         break;
768       }
769       Effect = T.takeError();
770     }
771     assert(Effect && "Expected at least one selection");
772     if (*Effect && (*Effect)->FormatEdits) {
773       // Format tweaks that require it centrally here.
774       for (auto &It : (*Effect)->ApplyEdits) {
775         Edit &E = It.second;
776         format::FormatStyle Style =
777             getFormatStyleForFile(File, E.InitialCode, TFS, false);
778         if (llvm::Error Err = reformatEdit(E, Style))
779           elog("Failed to format {0}: {1}", It.first(), std::move(Err));
780       }
781     } else {
782       TweakFailed.record(1, TweakID);
783     }
784     return CB(std::move(*Effect));
785   };
786   WorkScheduler->runWithAST("ApplyTweak", File, std::move(Action));
787 }
788 
789 void ClangdServer::locateSymbolAt(PathRef File, Position Pos,
790                                   Callback<std::vector<LocatedSymbol>> CB) {
791   auto Action = [Pos, CB = std::move(CB),
792                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
793     if (!InpAST)
794       return CB(InpAST.takeError());
795     CB(clangd::locateSymbolAt(InpAST->AST, Pos, Index));
796   };
797 
798   WorkScheduler->runWithAST("Definitions", File, std::move(Action));
799 }
800 
801 void ClangdServer::switchSourceHeader(
802     PathRef Path, Callback<std::optional<clangd::Path>> CB) {
803   // We want to return the result as fast as possible, strategy is:
804   //  1) use the file-only heuristic, it requires some IO but it is much
805   //     faster than building AST, but it only works when .h/.cc files are in
806   //     the same directory.
807   //  2) if 1) fails, we use the AST&Index approach, it is slower but supports
808   //     different code layout.
809   if (auto CorrespondingFile =
810           getCorrespondingHeaderOrSource(Path, TFS.view(std::nullopt)))
811     return CB(std::move(CorrespondingFile));
812   auto Action = [Path = Path.str(), CB = std::move(CB),
813                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
814     if (!InpAST)
815       return CB(InpAST.takeError());
816     CB(getCorrespondingHeaderOrSource(Path, InpAST->AST, Index));
817   };
818   WorkScheduler->runWithAST("SwitchHeaderSource", Path, std::move(Action));
819 }
820 
821 void ClangdServer::findDocumentHighlights(
822     PathRef File, Position Pos, Callback<std::vector<DocumentHighlight>> CB) {
823   auto Action =
824       [Pos, CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
825         if (!InpAST)
826           return CB(InpAST.takeError());
827         CB(clangd::findDocumentHighlights(InpAST->AST, Pos));
828       };
829 
830   WorkScheduler->runWithAST("Highlights", File, std::move(Action), Transient);
831 }
832 
833 void ClangdServer::findHover(PathRef File, Position Pos,
834                              Callback<std::optional<HoverInfo>> CB) {
835   auto Action = [File = File.str(), Pos, CB = std::move(CB),
836                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
837     if (!InpAST)
838       return CB(InpAST.takeError());
839     format::FormatStyle Style = getFormatStyleForFile(
840         File, InpAST->Inputs.Contents, *InpAST->Inputs.TFS, false);
841     CB(clangd::getHover(InpAST->AST, Pos, std::move(Style), Index));
842   };
843 
844   WorkScheduler->runWithAST("Hover", File, std::move(Action), Transient);
845 }
846 
847 void ClangdServer::typeHierarchy(PathRef File, Position Pos, int Resolve,
848                                  TypeHierarchyDirection Direction,
849                                  Callback<std::vector<TypeHierarchyItem>> CB) {
850   auto Action = [File = File.str(), Pos, Resolve, Direction, CB = std::move(CB),
851                  this](Expected<InputsAndAST> InpAST) mutable {
852     if (!InpAST)
853       return CB(InpAST.takeError());
854     CB(clangd::getTypeHierarchy(InpAST->AST, Pos, Resolve, Direction, Index,
855                                 File));
856   };
857 
858   WorkScheduler->runWithAST("TypeHierarchy", File, std::move(Action));
859 }
860 
861 void ClangdServer::superTypes(
862     const TypeHierarchyItem &Item,
863     Callback<std::optional<std::vector<TypeHierarchyItem>>> CB) {
864   WorkScheduler->run("typeHierarchy/superTypes", /*Path=*/"",
865                      [=, CB = std::move(CB)]() mutable {
866                        CB(clangd::superTypes(Item, Index));
867                      });
868 }
869 
870 void ClangdServer::subTypes(const TypeHierarchyItem &Item,
871                             Callback<std::vector<TypeHierarchyItem>> CB) {
872   WorkScheduler->run(
873       "typeHierarchy/subTypes", /*Path=*/"",
874       [=, CB = std::move(CB)]() mutable { CB(clangd::subTypes(Item, Index)); });
875 }
876 
877 void ClangdServer::resolveTypeHierarchy(
878     TypeHierarchyItem Item, int Resolve, TypeHierarchyDirection Direction,
879     Callback<std::optional<TypeHierarchyItem>> CB) {
880   WorkScheduler->run(
881       "Resolve Type Hierarchy", "", [=, CB = std::move(CB)]() mutable {
882         clangd::resolveTypeHierarchy(Item, Resolve, Direction, Index);
883         CB(Item);
884       });
885 }
886 
887 void ClangdServer::prepareCallHierarchy(
888     PathRef File, Position Pos, Callback<std::vector<CallHierarchyItem>> CB) {
889   auto Action = [File = File.str(), Pos,
890                  CB = std::move(CB)](Expected<InputsAndAST> InpAST) mutable {
891     if (!InpAST)
892       return CB(InpAST.takeError());
893     CB(clangd::prepareCallHierarchy(InpAST->AST, Pos, File));
894   };
895   WorkScheduler->runWithAST("CallHierarchy", File, std::move(Action));
896 }
897 
898 void ClangdServer::incomingCalls(
899     const CallHierarchyItem &Item,
900     Callback<std::vector<CallHierarchyIncomingCall>> CB) {
901   WorkScheduler->run("Incoming Calls", "",
902                      [CB = std::move(CB), Item, this]() mutable {
903                        CB(clangd::incomingCalls(Item, Index));
904                      });
905 }
906 
907 void ClangdServer::inlayHints(PathRef File, std::optional<Range> RestrictRange,
908                               Callback<std::vector<InlayHint>> CB) {
909   auto Action = [RestrictRange(std::move(RestrictRange)),
910                  CB = std::move(CB)](Expected<InputsAndAST> InpAST) mutable {
911     if (!InpAST)
912       return CB(InpAST.takeError());
913     CB(clangd::inlayHints(InpAST->AST, std::move(RestrictRange)));
914   };
915   WorkScheduler->runWithAST("InlayHints", File, std::move(Action), Transient);
916 }
917 
918 void ClangdServer::outgoingCalls(
919     const CallHierarchyItem &Item,
920     Callback<std::vector<CallHierarchyOutgoingCall>> CB) {
921   WorkScheduler->run("Outgoing Calls", "",
922                      [CB = std::move(CB), Item, this]() mutable {
923                        CB(clangd::outgoingCalls(Item, Index));
924                      });
925 }
926 
927 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
928   // FIXME: Do nothing for now. This will be used for indexing and potentially
929   // invalidating other caches.
930 }
931 
932 void ClangdServer::workspaceSymbols(
933     llvm::StringRef Query, int Limit,
934     Callback<std::vector<SymbolInformation>> CB) {
935   WorkScheduler->run(
936       "getWorkspaceSymbols", /*Path=*/"",
937       [Query = Query.str(), Limit, CB = std::move(CB), this]() mutable {
938         CB(clangd::getWorkspaceSymbols(Query, Limit, Index,
939                                        WorkspaceRoot.value_or("")));
940       });
941 }
942 
943 void ClangdServer::documentSymbols(llvm::StringRef File,
944                                    Callback<std::vector<DocumentSymbol>> CB) {
945   auto Action =
946       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
947         if (!InpAST)
948           return CB(InpAST.takeError());
949         CB(clangd::getDocumentSymbols(InpAST->AST));
950       };
951   WorkScheduler->runWithAST("DocumentSymbols", File, std::move(Action),
952                             Transient);
953 }
954 
955 void ClangdServer::foldingRanges(llvm::StringRef File,
956                                  Callback<std::vector<FoldingRange>> CB) {
957   auto Code = getDraft(File);
958   if (!Code)
959     return CB(llvm::make_error<LSPError>(
960         "trying to compute folding ranges for non-added document",
961         ErrorCode::InvalidParams));
962   auto Action = [LineFoldingOnly = LineFoldingOnly, CB = std::move(CB),
963                  Code = std::move(*Code)]() mutable {
964     CB(clangd::getFoldingRanges(Code, LineFoldingOnly));
965   };
966   // We want to make sure folding ranges are always available for all the open
967   // files, hence prefer runQuick to not wait for operations on other files.
968   WorkScheduler->runQuick("FoldingRanges", File, std::move(Action));
969 }
970 
971 void ClangdServer::findType(llvm::StringRef File, Position Pos,
972                             Callback<std::vector<LocatedSymbol>> CB) {
973   auto Action = [Pos, CB = std::move(CB),
974                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
975     if (!InpAST)
976       return CB(InpAST.takeError());
977     CB(clangd::findType(InpAST->AST, Pos, Index));
978   };
979   WorkScheduler->runWithAST("FindType", File, std::move(Action));
980 }
981 
982 void ClangdServer::findImplementations(
983     PathRef File, Position Pos, Callback<std::vector<LocatedSymbol>> CB) {
984   auto Action = [Pos, CB = std::move(CB),
985                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
986     if (!InpAST)
987       return CB(InpAST.takeError());
988     CB(clangd::findImplementations(InpAST->AST, Pos, Index));
989   };
990 
991   WorkScheduler->runWithAST("Implementations", File, std::move(Action));
992 }
993 
994 void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit,
995                                   bool AddContainer,
996                                   Callback<ReferencesResult> CB) {
997   auto Action = [Pos, Limit, AddContainer, CB = std::move(CB),
998                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
999     if (!InpAST)
1000       return CB(InpAST.takeError());
1001     CB(clangd::findReferences(InpAST->AST, Pos, Limit, Index, AddContainer));
1002   };
1003 
1004   WorkScheduler->runWithAST("References", File, std::move(Action));
1005 }
1006 
1007 void ClangdServer::symbolInfo(PathRef File, Position Pos,
1008                               Callback<std::vector<SymbolDetails>> CB) {
1009   auto Action =
1010       [Pos, CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
1011         if (!InpAST)
1012           return CB(InpAST.takeError());
1013         CB(clangd::getSymbolInfo(InpAST->AST, Pos));
1014       };
1015 
1016   WorkScheduler->runWithAST("SymbolInfo", File, std::move(Action));
1017 }
1018 
1019 void ClangdServer::semanticRanges(PathRef File,
1020                                   const std::vector<Position> &Positions,
1021                                   Callback<std::vector<SelectionRange>> CB) {
1022   auto Action = [Positions, CB = std::move(CB)](
1023                     llvm::Expected<InputsAndAST> InpAST) mutable {
1024     if (!InpAST)
1025       return CB(InpAST.takeError());
1026     std::vector<SelectionRange> Result;
1027     for (const auto &Pos : Positions) {
1028       if (auto Range = clangd::getSemanticRanges(InpAST->AST, Pos))
1029         Result.push_back(std::move(*Range));
1030       else
1031         return CB(Range.takeError());
1032     }
1033     CB(std::move(Result));
1034   };
1035   WorkScheduler->runWithAST("SemanticRanges", File, std::move(Action));
1036 }
1037 
1038 void ClangdServer::documentLinks(PathRef File,
1039                                  Callback<std::vector<DocumentLink>> CB) {
1040   auto Action =
1041       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
1042         if (!InpAST)
1043           return CB(InpAST.takeError());
1044         CB(clangd::getDocumentLinks(InpAST->AST));
1045       };
1046   WorkScheduler->runWithAST("DocumentLinks", File, std::move(Action),
1047                             Transient);
1048 }
1049 
1050 void ClangdServer::semanticHighlights(
1051     PathRef File, Callback<std::vector<HighlightingToken>> CB) {
1052 
1053   auto Action = [CB = std::move(CB),
1054                  PublishInactiveRegions = PublishInactiveRegions](
1055                     llvm::Expected<InputsAndAST> InpAST) mutable {
1056     if (!InpAST)
1057       return CB(InpAST.takeError());
1058     // Include inactive regions in semantic highlighting tokens only if the
1059     // client doesn't support a dedicated protocol for being informed about
1060     // them.
1061     CB(clangd::getSemanticHighlightings(InpAST->AST, !PublishInactiveRegions));
1062   };
1063   WorkScheduler->runWithAST("SemanticHighlights", File, std::move(Action),
1064                             Transient);
1065 }
1066 
1067 void ClangdServer::getAST(PathRef File, std::optional<Range> R,
1068                           Callback<std::optional<ASTNode>> CB) {
1069   auto Action =
1070       [R, CB(std::move(CB))](llvm::Expected<InputsAndAST> Inputs) mutable {
1071         if (!Inputs)
1072           return CB(Inputs.takeError());
1073         if (!R) {
1074           // It's safe to pass in the TU, as dumpAST() does not
1075           // deserialize the preamble.
1076           auto Node = DynTypedNode::create(
1077               *Inputs->AST.getASTContext().getTranslationUnitDecl());
1078           return CB(dumpAST(Node, Inputs->AST.getTokens(),
1079                             Inputs->AST.getASTContext()));
1080         }
1081         unsigned Start, End;
1082         if (auto Offset = positionToOffset(Inputs->Inputs.Contents, R->start))
1083           Start = *Offset;
1084         else
1085           return CB(Offset.takeError());
1086         if (auto Offset = positionToOffset(Inputs->Inputs.Contents, R->end))
1087           End = *Offset;
1088         else
1089           return CB(Offset.takeError());
1090         bool Success = SelectionTree::createEach(
1091             Inputs->AST.getASTContext(), Inputs->AST.getTokens(), Start, End,
1092             [&](SelectionTree T) {
1093               if (const SelectionTree::Node *N = T.commonAncestor()) {
1094                 CB(dumpAST(N->ASTNode, Inputs->AST.getTokens(),
1095                            Inputs->AST.getASTContext()));
1096                 return true;
1097               }
1098               return false;
1099             });
1100         if (!Success)
1101           CB(std::nullopt);
1102       };
1103   WorkScheduler->runWithAST("GetAST", File, std::move(Action));
1104 }
1105 
1106 void ClangdServer::customAction(PathRef File, llvm::StringRef Name,
1107                                 Callback<InputsAndAST> Action) {
1108   WorkScheduler->runWithAST(Name, File, std::move(Action));
1109 }
1110 
1111 void ClangdServer::diagnostics(PathRef File, Callback<std::vector<Diag>> CB) {
1112   auto Action =
1113       [CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
1114         if (!InpAST)
1115           return CB(InpAST.takeError());
1116         return CB(InpAST->AST.getDiagnostics());
1117       };
1118 
1119   WorkScheduler->runWithAST("Diagnostics", File, std::move(Action));
1120 }
1121 
1122 llvm::StringMap<TUScheduler::FileStats> ClangdServer::fileStats() const {
1123   return WorkScheduler->fileStats();
1124 }
1125 
1126 [[nodiscard]] bool
1127 ClangdServer::blockUntilIdleForTest(std::optional<double> TimeoutSeconds) {
1128   // Order is important here: we don't want to block on A and then B,
1129   // if B might schedule work on A.
1130 
1131 #if defined(__has_feature) &&                                                  \
1132     (__has_feature(address_sanitizer) || __has_feature(hwaddress_sanitizer) || \
1133      __has_feature(memory_sanitizer) || __has_feature(thread_sanitizer))
1134   if (TimeoutSeconds.has_value())
1135     (*TimeoutSeconds) *= 10;
1136 #endif
1137 
1138   // Nothing else can schedule work on TUScheduler, because it's not threadsafe
1139   // and we're blocking the main thread.
1140   if (!WorkScheduler->blockUntilIdle(timeoutSeconds(TimeoutSeconds)))
1141     return false;
1142   // TUScheduler is the only thing that starts background indexing work.
1143   if (IndexTasks && !IndexTasks->wait(timeoutSeconds(TimeoutSeconds)))
1144     return false;
1145 
1146   // Unfortunately we don't have strict topological order between the rest of
1147   // the components. E.g. CDB broadcast triggers backrgound indexing.
1148   // This queries the CDB which may discover new work if disk has changed.
1149   //
1150   // So try each one a few times in a loop.
1151   // If there are no tricky interactions then all after the first are no-ops.
1152   // Then on the last iteration, verify they're idle without waiting.
1153   //
1154   // There's a small chance they're juggling work and we didn't catch them :-(
1155   for (std::optional<double> Timeout :
1156        {TimeoutSeconds, TimeoutSeconds, std::optional<double>(0)}) {
1157     if (!CDB.blockUntilIdle(timeoutSeconds(Timeout)))
1158       return false;
1159     if (BackgroundIdx && !BackgroundIdx->blockUntilIdleForTest(Timeout))
1160       return false;
1161     if (FeatureModules && llvm::any_of(*FeatureModules, [&](FeatureModule &M) {
1162           return !M.blockUntilIdle(timeoutSeconds(Timeout));
1163         }))
1164       return false;
1165   }
1166 
1167   assert(WorkScheduler->blockUntilIdle(Deadline::zero()) &&
1168          "Something scheduled work while we're blocking the main thread!");
1169   return true;
1170 }
1171 
1172 void ClangdServer::profile(MemoryTree &MT) const {
1173   if (DynamicIdx)
1174     DynamicIdx->profile(MT.child("dynamic_index"));
1175   if (BackgroundIdx)
1176     BackgroundIdx->profile(MT.child("background_index"));
1177   WorkScheduler->profile(MT.child("tuscheduler"));
1178 }
1179 } // namespace clangd
1180 } // namespace clang
1181