[Code Completion] Split code completion results into context free part that can be cached and contextual part displayed to the user

This allows makes the distinction between cachable and non-cachable properties cleaner and allows us to more easily compute contextual information (like type relations) for cached items later.

Author: ahoppen

include/swift/IDE/CodeCompletion.h

@@ -61,8 +61,14 @@ class CodeCompletionCache {
   };
 
   struct Value : public llvm::ThreadSafeRefCountedBase<Value> {
+    /// The allocator used to allocate the results stored in this cache.
+    std::shared_ptr<llvm::BumpPtrAllocator> Allocator;
+
     llvm::sys::TimePoint<> ModuleModificationTime;
-    CodeCompletionResultSink Sink;
+
+    std::vector<const ContextFreeCodeCompletionResult *> Results;
+
+    Value() : Allocator(std::make_shared<llvm::BumpPtrAllocator>()) {}
   };
   using ValueRefCntPtr = llvm::IntrusiveRefCntPtr<Value>;
 

include/swift/IDE/CodeCompletionCache.h

@@ -41,7 +41,7 @@ namespace swift {
     struct CacheValueCostInfo<swift::ide::CodeCompletionCacheImpl::Value> {
       static size_t
       getCost(const swift::ide::CodeCompletionCacheImpl::Value &V) {
-        return V.Sink.Allocator->getTotalMemory();
+        return V.Allocator->getTotalMemory();
       }
     };
   } // namespace sys
@@ -102,7 +102,7 @@ CodeCompletionCache::~CodeCompletionCache() {}
 ///
 /// This should be incremented any time we commit a change to the format of the
 /// cached results. This isn't expected to change very often.
-static constexpr uint32_t onDiskCompletionCacheVersion = 3; // Removed "source file path".
+static constexpr uint32_t onDiskCompletionCacheVersion = 4; // Store ContextFreeCodeCompletionResults in cache
 
 /// Deserializes CodeCompletionResults from \p in and stores them in \p V.
 /// \see writeCacheModule.
@@ -166,7 +166,7 @@ static bool readCachedModule(llvm::MemoryBuffer *in,
 
     const char *p = strings + index;
     auto size = read32le(p);
-    auto str = copyString(*V.Sink.Allocator, StringRef(p, size));
+    auto str = copyString(*V.Allocator, StringRef(p, size));
     knownStrings[index] = str;
     return str;
   };
@@ -192,20 +192,19 @@ static bool readCachedModule(llvm::MemoryBuffer *in,
       }
     }
 
-    return CodeCompletionString::create(*V.Sink.Allocator, chunkList);
+    return CodeCompletionString::create(*V.Allocator, chunkList);
   };
 
   // RESULTS
   while (cursor != resultEnd) {
     auto kind = static_cast<CodeCompletionResultKind>(*cursor++);
-    auto declKind = static_cast<CodeCompletionDeclKind>(*cursor++);
+    auto associatedKind = static_cast<uint8_t>(*cursor++);
     auto opKind = static_cast<CodeCompletionOperatorKind>(*cursor++);
-    auto context = static_cast<SemanticContextKind>(*cursor++);
-    auto notRecommended = static_cast<NotRecommendedReason>(*cursor++);
+    auto notRecommended =
+        static_cast<ContextFreeNotRecommendedReason>(*cursor++);
     auto diagSeverity =
         static_cast<CodeCompletionDiagnosticSeverity>(*cursor++);
     auto isSystem = static_cast<bool>(*cursor++);
-    auto numBytesToErase = static_cast<unsigned>(*cursor++);
     auto chunkIndex = read32le(cursor);
     auto moduleIndex = read32le(cursor);
     auto briefDocIndex = read32le(cursor);
@@ -222,21 +221,14 @@ static bool readCachedModule(llvm::MemoryBuffer *in,
     auto briefDocComment = getString(briefDocIndex);
     auto diagMessage = getString(diagMessageIndex);
 
-    CodeCompletionResult *result = nullptr;
-    if (kind == CodeCompletionResultKind::Declaration) {
-      result = new (*V.Sink.Allocator) CodeCompletionResult(
-          context, CodeCompletionFlair(), numBytesToErase, string, declKind,
-          isSystem, moduleName, notRecommended, diagSeverity, diagMessage,
-          briefDocComment,
-          copyArray(*V.Sink.Allocator, ArrayRef<StringRef>(assocUSRs)),
-          CodeCompletionResult::ExpectedTypeRelation::Unknown, opKind);
-    } else {
-      result = new (*V.Sink.Allocator) CodeCompletionResult(
-          kind, context, CodeCompletionFlair(), numBytesToErase, string,
-          CodeCompletionResult::ExpectedTypeRelation::NotApplicable, opKind);
-    }
+    ContextFreeCodeCompletionResult *result =
+        new (*V.Allocator) ContextFreeCodeCompletionResult(
+            kind, associatedKind, opKind, isSystem, string, moduleName,
+            briefDocComment,
+            copyArray(*V.Allocator, ArrayRef<StringRef>(assocUSRs)),
+            notRecommended, diagSeverity, diagMessage);
 
-    V.Sink.Results.push_back(result);
+    V.Results.push_back(result);
   }
 
   return true;
@@ -342,27 +334,18 @@ static void writeCachedModule(llvm::raw_ostream &out,
   // RESULTS
   {
     endian::Writer LE(results, little);
-    for (CodeCompletionResult *R : V.Sink.Results) {
-      assert(!R->getFlair().toRaw() && "Any flairs should not be cached");
-      assert(R->getNotRecommendedReason() !=
-                 NotRecommendedReason::InvalidAsyncContext &&
-             "InvalidAsyncContext is decl context specific, cannot be cached");
-
+    for (const ContextFreeCodeCompletionResult *R : V.Results) {
       // FIXME: compress bitfield
       LE.write(static_cast<uint8_t>(R->getKind()));
-      if (R->getKind() == CodeCompletionResultKind::Declaration)
-        LE.write(static_cast<uint8_t>(R->getAssociatedDeclKind()));
-      else
-        LE.write(static_cast<uint8_t>(~0u));
-      if (R->isOperator())
-        LE.write(static_cast<uint8_t>(R->getOperatorKind()));
-      else
+      LE.write(static_cast<uint8_t>(R->getOpaqueAssociatedKind()));
+      if (R->isOperator()) {
+        LE.write(static_cast<uint8_t>(R->getKnownOperatorKind()));
+      } else {
         LE.write(static_cast<uint8_t>(CodeCompletionOperatorKind::None));
-      LE.write(static_cast<uint8_t>(R->getSemanticContext()));
+      }
       LE.write(static_cast<uint8_t>(R->getNotRecommendedReason()));
       LE.write(static_cast<uint8_t>(R->getDiagnosticSeverity()));
       LE.write(static_cast<uint8_t>(R->isSystem()));
-      LE.write(static_cast<uint8_t>(R->getNumBytesToErase()));
       LE.write(
           static_cast<uint32_t>(addCompletionString(R->getCompletionString())));
       LE.write(addString(R->getModuleName()));      // index into strings

lib/IDE/CodeCompletion.cpp

@@ -144,33 +144,43 @@ bool CodeCompletionDiagnostics::getDiagnosticForDeprecated(
 
 } // namespace
 
-bool swift::ide::getCompletionDiagnostics(
-    NotRecommendedReason reason, const ValueDecl *D,
-    CodeCompletionDiagnosticSeverity &severity, llvm::raw_ostream &Out) {
-  ASTContext &ctx = D->getASTContext();
-
-  CodeCompletionDiagnostics Diag(ctx);
-  switch (reason) {
-  case NotRecommendedReason::Deprecated:
-  case NotRecommendedReason::SoftDeprecated:
-    return Diag.getDiagnosticForDeprecated(D, severity, Out);
-  case NotRecommendedReason::InvalidAsyncContext:
+bool swift::ide::getContextFreeCompletionDiagnostics(
+    ContextFreeNotRecommendedReason Reason, const ValueDecl *D,
+    CodeCompletionDiagnosticSeverity &Severity, llvm::raw_ostream &Out) {
+  CodeCompletionDiagnostics Diag(D->getASTContext());
+  switch (Reason) {
+  case ContextFreeNotRecommendedReason::Deprecated:
+  case ContextFreeNotRecommendedReason::SoftDeprecated:
+    return Diag.getDiagnosticForDeprecated(D, Severity, Out);
+  case ContextFreeNotRecommendedReason::None:
+    llvm_unreachable("invalid not recommended reason");
+  }
+  return true;
+}
+
+bool swift::ide::getContextualCompletionDiagnostics(
+    ContextualNotRecommendedReason Reason, const ValueDecl *D,
+    CodeCompletionDiagnosticSeverity &Severity, llvm::raw_ostream &Out) {
+  CodeCompletionDiagnostics Diag(D->getASTContext());
+  switch (Reason) {
+  case ContextualNotRecommendedReason::InvalidAsyncContext:
     // FIXME: Could we use 'diag::async_in_nonasync_function'?
-    return Diag.getDiagnostics(severity, Out, diag::ide_async_in_nonasync_context,
-                        D->getName());
-  case NotRecommendedReason::CrossActorReference:
-    return Diag.getDiagnostics(severity, Out, diag::ide_cross_actor_reference_swift5,
-                        D->getName());
-  case NotRecommendedReason::RedundantImport:
-    return Diag.getDiagnostics(severity, Out, diag::ide_redundant_import,
-                        D->getName());
-  case NotRecommendedReason::RedundantImportIndirect:
-    return Diag.getDiagnostics(severity, Out, diag::ide_redundant_import_indirect,
-                        D->getName());
-  case NotRecommendedReason::VariableUsedInOwnDefinition:
-    return Diag.getDiagnostics(severity, Out, diag::recursive_accessor_reference,
-                        D->getName().getBaseIdentifier(), /*"getter"*/ 0);
-  case NotRecommendedReason::None:
+    return Diag.getDiagnostics(
+        Severity, Out, diag::ide_async_in_nonasync_context, D->getName());
+  case ContextualNotRecommendedReason::CrossActorReference:
+    return Diag.getDiagnostics(
+        Severity, Out, diag::ide_cross_actor_reference_swift5, D->getName());
+  case ContextualNotRecommendedReason::RedundantImport:
+    return Diag.getDiagnostics(Severity, Out, diag::ide_redundant_import,
+                               D->getName());
+  case ContextualNotRecommendedReason::RedundantImportIndirect:
+    return Diag.getDiagnostics(
+        Severity, Out, diag::ide_redundant_import_indirect, D->getName());
+  case ContextualNotRecommendedReason::VariableUsedInOwnDefinition:
+    return Diag.getDiagnostics(
+        Severity, Out, diag::recursive_accessor_reference,
+        D->getName().getBaseIdentifier(), /*"getter"*/ 0);
+  case ContextualNotRecommendedReason::None:
     llvm_unreachable("invalid not recommended reason");
   }
   return true;

lib/IDE/CodeCompletionCache.cpp

@@ -21,11 +21,21 @@ class ValueDecl;
 
 namespace ide {
 
-/// Populate \p severity and \p Out with the diagnostics for \p D.
+/// Populate \p severity and \p Out with the context-free diagnostics for \p D.
+/// See \c NotRecommendedReason for an explaination of context-free vs.
+/// contextual diagnostics.
 /// Returns \c true if it fails to generate the diagnostics.
-bool getCompletionDiagnostics(NotRecommendedReason reason, const ValueDecl *D,
-                              CodeCompletionDiagnosticSeverity &severity,
-                              llvm::raw_ostream &Out);
+bool getContextFreeCompletionDiagnostics(
+    ContextFreeNotRecommendedReason reason, const ValueDecl *D,
+    CodeCompletionDiagnosticSeverity &severity, llvm::raw_ostream &Out);
+
+/// Populate \p severity and \p Out with the contextual diagnostics for \p D.
+/// See \c NotRecommendedReason for an explaination of context-free vs.
+/// contextual diagnostics.
+/// Returns \c true if it fails to generate the diagnostics.
+bool getContextualCompletionDiagnostics(
+    ContextualNotRecommendedReason reason, const ValueDecl *D,
+    CodeCompletionDiagnosticSeverity &severity, llvm::raw_ostream &Out);
 
 } // namespace ide
 } // namespace swift

lib/IDE/CodeCompletionDiagnostics.cpp

@@ -89,7 +89,10 @@ class CodeCompletionResultBuilder {
   CodeCompletionResult::ExpectedTypeRelation ExpectedTypeRelation =
       CodeCompletionResult::ExpectedTypeRelation::Unknown;
   bool Cancelled = false;
-  NotRecommendedReason NotRecReason = NotRecommendedReason::None;
+  ContextFreeNotRecommendedReason ContextFreeNotRecReason =
+      ContextFreeNotRecommendedReason::None;
+  ContextualNotRecommendedReason ContextualNotRecReason =
+      ContextualNotRecommendedReason::None;
   StringRef BriefDocComment;
 
   void addChunkWithText(CodeCompletionString::Chunk::ChunkKind Kind,
@@ -145,10 +148,11 @@ class CodeCompletionResultBuilder {
 
   void setLiteralKind(CodeCompletionLiteralKind kind) { LiteralKind = kind; }
   void setKeywordKind(CodeCompletionKeywordKind kind) { KeywordKind = kind; }
-  void setNotRecommended(NotRecommendedReason Reason) { NotRecReason = Reason; }
-
-  void setSemanticContext(SemanticContextKind Kind) {
-    SemanticContext = Kind;
+  void setContextFreeNotRecommended(ContextFreeNotRecommendedReason Reason) {
+    ContextFreeNotRecReason = Reason;
+  }
+  void setContextualNotRecommended(ContextualNotRecommendedReason Reason) {
+    ContextualNotRecReason = Reason;
   }
 
   void addFlair(CodeCompletionFlair Options) {

lib/IDE/CodeCompletionDiagnostics.h

@@ -154,11 +154,19 @@ bool SourceKit::CodeCompletion::addCustomCompletions(
     auto chunk = Chunk::createWithText(Chunk::ChunkKind::Text, 0, nameCopy);
     auto *completionString =
         CodeCompletionString::create(sink.allocator, chunk);
+    auto *contextFreeResult =
+        new (sink.allocator) ContextFreeCodeCompletionResult(
+            CodeCompletionResultKind::Pattern, completionString,
+            CodeCompletionOperatorKind::None, /*BriefDocComment=*/"",
+            ContextFreeNotRecommendedReason::None,
+            CodeCompletionDiagnosticSeverity::None, /*DiagnosticMessage=*/"");
     auto *swiftResult = new (sink.allocator) CodeCompletion::SwiftResult(
-        CodeCompletionResultKind::Pattern, SemanticContextKind::Local,
+        *contextFreeResult, SemanticContextKind::Local,
         CodeCompletionFlairBit::ExpressionSpecific,
-        /*NumBytesToErase=*/0, completionString,
-        CodeCompletionResult::ExpectedTypeRelation::Unknown);
+        /*NumBytesToErase=*/0,
+        CodeCompletionResult::ExpectedTypeRelation::Unknown,
+        ContextualNotRecommendedReason::None,
+        CodeCompletionDiagnosticSeverity::None, /*DiagnosticMessage=*/"");
 
     CompletionBuilder builder(sink, *swiftResult);
     builder.setCustomKind(customCompletion.Kind);
@@ -1176,18 +1184,26 @@ Completion *CompletionBuilder::finish() {
     if (base.isOperator())
       opKind = base.getOperatorKind();
 
-    if (base.getKind() == CodeCompletionResultKind::Declaration) {
-      newBase = new (sink.allocator) SwiftResult(
-          semanticContext, flair, base.getNumBytesToErase(), completionString,
-          base.getAssociatedDeclKind(), base.isSystem(), base.getModuleName(),
-          base.getNotRecommendedReason(), base.getDiagnosticSeverity(),
-          base.getDiagnosticMessage(), base.getBriefDocComment(),
-          base.getAssociatedUSRs(), typeRelation, opKind);
-    } else {
-      newBase = new (sink.allocator) SwiftResult(
-          base.getKind(), semanticContext, flair, base.getNumBytesToErase(),
-          completionString, typeRelation, opKind);
-    }
+    const ContextFreeCodeCompletionResult &contextFreeBase =
+        base.getContextFreeResult();
+
+    ContextFreeCodeCompletionResult *contextFreeResult =
+        new (sink.allocator) ContextFreeCodeCompletionResult(
+            contextFreeBase.getKind(),
+            contextFreeBase.getOpaqueAssociatedKind(), opKind,
+            contextFreeBase.isSystem(), completionString,
+            contextFreeBase.getModuleName(),
+            contextFreeBase.getBriefDocComment(),
+            contextFreeBase.getAssociatedUSRs(),
+            contextFreeBase.getNotRecommendedReason(),
+            contextFreeBase.getDiagnosticSeverity(),
+            contextFreeBase.getDiagnosticMessage());
+
+    newBase = new (sink.allocator) SwiftResult(
+        *contextFreeResult, semanticContext, flair, base.getNumBytesToErase(),
+        typeRelation, base.getContextualNotRecommendedReason(),
+        base.getContextualDiagnosticSeverity(),
+        base.getContextualDiagnosticMessage());
 
     llvm::raw_svector_ostream OSS(nameStorage);
     ide::printCodeCompletionResultFilterName(*newBase, OSS);

lib/IDE/CodeCompletionResultBuilder.h

@@ -919,12 +919,20 @@ static void transformAndForwardResults(
   auto buildInnerResult = [&](ArrayRef<CodeCompletionString::Chunk> chunks) {
     auto *completionString =
         CodeCompletionString::create(innerSink.allocator, chunks);
+    ContextFreeCodeCompletionResult *contextFreeResult =
+        new (innerSink.allocator) ContextFreeCodeCompletionResult(
+            CodeCompletionResultKind::BuiltinOperator, completionString,
+            CodeCompletionOperatorKind::None,
+            /*BriefDocComment=*/"", ContextFreeNotRecommendedReason::None,
+            CodeCompletionDiagnosticSeverity::None,
+            /*DiagnosticMessage=*/"");
     auto *paren = new (innerSink.allocator) CodeCompletion::SwiftResult(
-        CodeCompletionResultKind::BuiltinOperator,
-        SemanticContextKind::CurrentNominal,
+        *contextFreeResult, SemanticContextKind::CurrentNominal,
         CodeCompletionFlairBit::ExpressionSpecific,
-        exactMatch ? exactMatch->getNumBytesToErase() : 0, completionString,
-        CodeCompletionResult::ExpectedTypeRelation::NotApplicable);
+        exactMatch ? exactMatch->getNumBytesToErase() : 0,
+        CodeCompletionResult::ExpectedTypeRelation::NotApplicable,
+        ContextualNotRecommendedReason::None,
+        CodeCompletionDiagnosticSeverity::None, /*DiagnosticMessage=*/"");
 
     SwiftCompletionInfo info;
     std::vector<Completion *> extended = extendCompletions(

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp

@@ -983,7 +983,7 @@ class SemanticAnnotator : public SourceEntityWalker {
 
     unsigned ByteOffset = SM.getLocOffsetInBuffer(Range.getStart(), BufferID);
     unsigned Length = Range.getByteLength();
-    auto Kind = CodeCompletionResult::getCodeCompletionDeclKind(D);
+    auto Kind = ContextFreeCodeCompletionResult::getCodeCompletionDeclKind(D);
     bool IsSystem = D->getModuleContext()->isSystemModule();
     SemaToks.emplace_back(Kind, ByteOffset, Length, IsRef, IsSystem);
   }

tools/SourceKit/lib/SwiftLang/SwiftCompletion.cpp

@@ -4179,10 +4179,25 @@ int main(int argc, char *argv[]) {
         // FIXME: error?
         continue;
       }
+      // Make contextual CodeCompletionResults from the
+      // ContextFreeCodeCompletionResults so we can print them.
+      std::vector<CodeCompletionResult *> contextualResults;
+      contextualResults.reserve(resultsOpt->get()->Results.size());
+      for (auto contextFreeResult : resultsOpt->get()->Results) {
+        // We are leaking these results but it doesn't matter since the process
+        // just terminates afterwards anyway.
+        auto contextualResult = new CodeCompletionResult(
+            *contextFreeResult, SemanticContextKind::OtherModule,
+            CodeCompletionFlair(),
+            /*numBytesToErase=*/0,
+            CodeCompletionResult::ExpectedTypeRelation::Unknown,
+            ContextualNotRecommendedReason::None,
+            CodeCompletionDiagnosticSeverity::None, /*DiagnosticMessage=*/"");
+        contextualResults.push_back(contextualResult);
+      }
       printCodeCompletionResultsImpl(
-          resultsOpt->get()->Sink.Results, llvm::outs(),
-          options::CodeCompletionKeywords, options::CodeCompletionComments,
-          options::CodeCompletionSourceText,
+          contextualResults, llvm::outs(), options::CodeCompletionKeywords,
+          options::CodeCompletionComments, options::CodeCompletionSourceText,
           options::CodeCompletionAnnotateResults);
     }