Index: clangd/CodeComplete.cpp
===================================================================
--- clangd/CodeComplete.cpp
+++ clangd/CodeComplete.cpp
@@ -430,12 +430,13 @@
 // It doesn't do scoring or conversion to CompletionItem yet, as we want to
 // merge with index results first.
 struct CompletionRecorder : public CodeCompleteConsumer {
-  CompletionRecorder(const CodeCompleteOptions &Opts)
+  CompletionRecorder(const CodeCompleteOptions &Opts,
+                     UniqueFunction<void()> ResultsCallback)
       : CodeCompleteConsumer(Opts.getClangCompleteOpts(),
                              /*OutputIsBinary=*/false),
         CCContext(CodeCompletionContext::CCC_Other), Opts(Opts),
         CCAllocator(std::make_shared<GlobalCodeCompletionAllocator>()),
-        CCTUInfo(CCAllocator) {}
+        CCTUInfo(CCAllocator), ResultsCallback(std::move(ResultsCallback)) {}
   std::vector<CodeCompletionResult> Results;
   CodeCompletionContext CCContext;
   Sema *CCSema = nullptr; // Sema that created the results.
@@ -466,6 +467,8 @@
         continue;
       Results.push_back(Result);
     }
+    if (ResultsCallback)
+      ResultsCallback();
   }
 
   CodeCompletionAllocator &getAllocator() override { return *CCAllocator; }
@@ -503,6 +506,7 @@
   CodeCompleteOptions Opts;
   std::shared_ptr<GlobalCodeCompletionAllocator> CCAllocator;
   CodeCompletionTUInfo CCTUInfo;
+  UniqueFunction<void()> ResultsCallback;
 };
 
 // Tracks a bounded number of candidates with the best scores.
@@ -660,9 +664,8 @@
 // Callback will be invoked once completion is done, but before cleaning up.
 bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
                       const clang::CodeCompleteOptions &Options,
-                      const SemaCompleteInput &Input,
-                      llvm::function_ref<void()> Callback = nullptr) {
-  auto Tracer = llvm::make_unique<trace::Span>("Sema completion");
+                      const SemaCompleteInput &Input) {
+  trace::Span Tracer("Sema completion");
   std::vector<const char *> ArgStrs;
   for (const auto &S : Input.Command.CommandLine)
     ArgStrs.push_back(S.c_str());
@@ -729,12 +732,6 @@
     log("Execute() failed when running codeComplete for " + Input.FileName);
     return false;
   }
-  Tracer.reset();
-
-  if (Callback)
-    Callback();
-
-  Tracer = llvm::make_unique<trace::Span>("Sema completion cleanup");
   Action.EndSourceFile();
 
   return true;
@@ -834,8 +831,7 @@
   PathRef FileName;
   const CodeCompleteOptions &Opts;
   // Sema takes ownership of Recorder. Recorder is valid until Sema cleanup.
-  std::unique_ptr<CompletionRecorder> RecorderOwner;
-  CompletionRecorder &Recorder;
+  CompletionRecorder *Recorder = nullptr;
   int NSema = 0, NIndex = 0, NBoth = 0; // Counters for logging.
   bool Incomplete = false; // Would more be available with a higher limit?
   llvm::Optional<FuzzyMatcher> Filter; // Initialized once Sema runs.
@@ -843,8 +839,7 @@
 public:
   // A CodeCompleteFlow object is only useful for calling run() exactly once.
   CodeCompleteFlow(PathRef FileName, const CodeCompleteOptions &Opts)
-      : FileName(FileName), Opts(Opts),
-        RecorderOwner(new CompletionRecorder(Opts)), Recorder(*RecorderOwner) {}
+      : FileName(FileName), Opts(Opts) {}
 
   CompletionList run(const SemaCompleteInput &SemaCCInput) && {
     trace::Span Tracer("CodeCompleteFlow");
@@ -852,16 +847,16 @@
     //   - completion results based on the AST. These are saved for merging.
     //   - partial identifier and context. We need these for the index query.
     CompletionList Output;
+    auto RecorderOwner = llvm::make_unique<CompletionRecorder>(Opts, [&]() {
+      assert(Recorder && "Recorder is not set");
+      Output = runWithSema();
+      SPAN_ATTACH(Tracer, "sema_completion_kind",
+                  getCompletionKindString(Recorder->CCContext.getKind()));
+    });
+
+    Recorder = RecorderOwner.get();
     semaCodeComplete(std::move(RecorderOwner), Opts.getClangCompleteOpts(),
-                     SemaCCInput, [&] {
-                       if (Recorder.CCSema) {
-                         Output = runWithSema();
-                         SPAN_ATTACH(
-                             Tracer, "sema_completion_kind",
-                             getCompletionKindString(Recorder.CCContext.getKind()));
-                       } else
-                         log("Code complete: no Sema callback, 0 results");
-                     });
+                     SemaCCInput);
 
     SPAN_ATTACH(Tracer, "sema_results", NSema);
     SPAN_ATTACH(Tracer, "index_results", NIndex);
@@ -883,7 +878,7 @@
   // Sema data structures are torn down. It does all the real work.
   CompletionList runWithSema() {
     Filter = FuzzyMatcher(
-        Recorder.CCSema->getPreprocessor().getCodeCompletionFilter());
+        Recorder->CCSema->getPreprocessor().getCodeCompletionFilter());
     // Sema provides the needed context to query the index.
     // FIXME: in addition to querying for extra/overlapping symbols, we should
     //        explicitly request symbols corresponding to Sema results.
@@ -891,7 +886,7 @@
     // We must copy index results to preserve them, but there are at most Limit.
     auto IndexResults = queryIndex();
     // Merge Sema and Index results, score them, and pick the winners.
-    auto Top = mergeResults(Recorder.Results, IndexResults);
+    auto Top = mergeResults(Recorder->Results, IndexResults);
     // Convert the results to the desired LSP structs.
     CompletionList Output;
     for (auto &C : Top)
@@ -901,7 +896,7 @@
   }
 
   SymbolSlab queryIndex() {
-    if (!Opts.Index || !allowIndex(Recorder.CCContext.getKind()))
+    if (!Opts.Index || !allowIndex(Recorder->CCContext.getKind()))
       return SymbolSlab();
     trace::Span Tracer("Query index");
     SPAN_ATTACH(Tracer, "limit", Opts.Limit);
@@ -912,8 +907,8 @@
     if (Opts.Limit)
       Req.MaxCandidateCount = Opts.Limit;
     Req.Query = Filter->pattern();
-    Req.Scopes =
-        getQueryScopes(Recorder.CCContext, Recorder.CCSema->getSourceManager());
+    Req.Scopes = getQueryScopes(Recorder->CCContext,
+                                Recorder->CCSema->getSourceManager());
     log(llvm::formatv("Code complete: fuzzyFind(\"{0}\", scopes=[{1}])",
                       Req.Query,
                       llvm::join(Req.Scopes.begin(), Req.Scopes.end(), ",")));
@@ -945,7 +940,7 @@
       return nullptr;
     };
     // Emit all Sema results, merging them with Index results if possible.
-    for (auto &SemaResult : Recorder.Results)
+    for (auto &SemaResult : Recorder->Results)
       addCandidate(Top, &SemaResult, CorrespondingIndexResult(SemaResult));
     // Now emit any Index-only results.
     for (const auto &IndexResult : IndexResults) {
@@ -962,7 +957,7 @@
     CompletionCandidate C;
     C.SemaResult = SemaResult;
     C.IndexResult = IndexResult;
-    C.Name = IndexResult ? IndexResult->Name : Recorder.getName(*SemaResult);
+    C.Name = IndexResult ? IndexResult->Name : Recorder->getName(*SemaResult);
 
     CompletionItemScores Scores;
     if (auto FuzzyScore = Filter->match(C.Name))
@@ -986,7 +981,7 @@
                                   const CompletionItemScores &Scores) {
     CodeCompletionString *SemaCCS = nullptr;
     if (auto *SR = Candidate.SemaResult)
-      SemaCCS = Recorder.codeCompletionString(*SR, Opts.IncludeBriefComments);
+      SemaCCS = Recorder->codeCompletionString(*SR, Opts.IncludeBriefComments);
     return Candidate.build(FileName, Scores, Opts, SemaCCS);
   }
 };