[CodeCompletion] Don't modify CodeCompletionResult in place for cached results

Instead, allocate new `CodeCompletionResult` in the current sink.

(cherry picked from commit f553826)

Author: rintaro

include/swift/IDE/CodeCompletion.h

@@ -40,6 +40,7 @@ namespace ide {
 class CodeCompletionCache;
 class CodeCompletionContext;
 class CodeCompletionResultBuilder;
+struct CodeCompletionResultSink;
 struct RequestedCachedModule;
 
 /// A routine to remove code completion tokens from code completion
@@ -760,6 +761,12 @@ class CodeCompletionResult {
            getOperatorKind() != CodeCompletionOperatorKind::None);
   }
 
+  /// Copy this result to \p Sink with \p newFlair . Note that this does NOT
+  /// copy the value of \c CompletionString , \c AssociatedUSRs etc. it only
+  /// copies the pointers to them.
+  CodeCompletionResult *withFlair(CodeCompletionFlair newFlair,
+                                  CodeCompletionResultSink &Sink);
+
   ResultKind getKind() const { return static_cast<ResultKind>(Kind); }
 
   CodeCompletionDeclKind getAssociatedDeclKind() const {
@@ -815,6 +822,7 @@ class CodeCompletionResult {
     return static_cast<CodeCompletionFlair>(Flair);
   }
 
+  /// Modify "flair" of this result *in place*.
   void setFlair(CodeCompletionFlair flair) {
     Flair = unsigned(flair.toRaw());
   }
@@ -1030,7 +1038,7 @@ void lookupCodeCompletionResultsFromModule(CodeCompletionResultSink &targetSink,
 
 /// Copy code completion results from \p sourceSink to \p targetSink, possibly
 /// restricting by \p onlyTypes. Returns copied results in \p targetSink.
-ArrayRef<CodeCompletionResult *>
+MutableArrayRef<CodeCompletionResult *>
 copyCodeCompletionResults(CodeCompletionResultSink &targetSink,
                           CodeCompletionResultSink &sourceSink,
                           bool onlyTypes,

lib/IDE/CodeCompletion.cpp

@@ -801,6 +801,24 @@ void CodeCompletionResult::dump() const {
   llvm::errs() << "\n";
 }
 
+CodeCompletionResult *
+CodeCompletionResult::withFlair(CodeCompletionFlair newFlair,
+                                CodeCompletionResultSink &Sink) {
+  if (Kind == ResultKind::Declaration) {
+    return new (*Sink.Allocator) CodeCompletionResult(
+        getSemanticContext(), newFlair, getNumBytesToErase(),
+        getCompletionString(), getAssociatedDeclKind(), isSystem(),
+        getModuleName(), getNotRecommendedReason(), getBriefDocComment(),
+        getAssociatedUSRs(), getDeclKeywords(), getExpectedTypeRelation(),
+        isOperator() ? getOperatorKind() : CodeCompletionOperatorKind::None);
+  } else {
+    return new (*Sink.Allocator) CodeCompletionResult(
+        getKind(), getSemanticContext(), newFlair, getNumBytesToErase(),
+        getCompletionString(), getExpectedTypeRelation(),
+        isOperator() ? getOperatorKind() : CodeCompletionOperatorKind::None);
+  }
+}
+
 void CodeCompletionResultBuilder::withNestedGroup(
     CodeCompletionString::Chunk::ChunkKind Kind,
     llvm::function_ref<void()> body) {
@@ -6497,9 +6515,16 @@ static void addConditionalCompilationFlags(ASTContext &Ctx,
   }
 }
 
-static void postProcessResults(ArrayRef<CodeCompletionResult *> results,
-                               CompletionKind Kind, DeclContext *DC) {
-  for (CodeCompletionResult *result : results) {
+/// Add flairs to the each item in \p results .
+///
+/// If \p Sink is passed, the pointer of the each result may be replaced with a
+/// pointer to the new item allocated in \p Sink.
+/// If \p Sink is nullptr, the pointee of each result may be modified in place.
+static void postProcessResults(MutableArrayRef<CodeCompletionResult *> results,
+                               CompletionKind Kind, DeclContext *DC,
+                               CodeCompletionResultSink *Sink) {
+  for (CodeCompletionResult *&result : results) {
+    bool modified = false;
     auto flair = result->getFlair();
 
     // Starting a statement with a protocol name is not common. So protocol
@@ -6512,15 +6537,27 @@ static void postProcessResults(ArrayRef<CodeCompletionResult *> results,
         Kind != CompletionKind::TypeDeclResultBeginning &&
         Kind != CompletionKind::GenericRequirement) {
       flair |= CodeCompletionFlairBit::RareTypeAtCurrentPosition;
+      modified = true;
     }
 
     // Starting a statement at top-level in non-script files is invalid.
     if (Kind == CompletionKind::StmtOrExpr &&
         result->getKind() == CodeCompletionResult::ResultKind::Declaration &&
         isCodeCompletionAtTopLevelOfLibraryFile(DC)) {
       flair |= CodeCompletionFlairBit::ExpressionAtNonScriptOrMainFileScope;
+      modified = true;
+    }
+
+    if (!modified)
+      continue;
+
+    if (Sink) {
+      // Replace the result with a new result with the flair.
+      result = result->withFlair(flair, *Sink);
+    } else {
+      // 'Sink' == nullptr means the result is modifiable in place.
+      result->setFlair(flair);
     }
-    result->setFlair(flair);
   }
 }
 
@@ -6640,7 +6677,8 @@ static void deliverCompletionResults(CodeCompletionContext &CompletionContext,
   CompletionContext.typeContextKind = Lookup.typeContextKind();
 
   postProcessResults(CompletionContext.getResultSink().Results,
-                     CompletionContext.CodeCompletionKind, DC);
+                     CompletionContext.CodeCompletionKind, DC,
+                     /*Sink=*/nullptr);
 
   Consumer.handleResultsAndModules(CompletionContext, RequestedModules, DC);
 }
@@ -7424,7 +7462,8 @@ void swift::ide::lookupCodeCompletionResultsFromModule(
   CompletionLookup Lookup(targetSink, module->getASTContext(), SF);
   Lookup.lookupExternalModuleDecls(module, accessPath, needLeadingDot);
 }
-ArrayRef<CodeCompletionResult *>
+
+MutableArrayRef<CodeCompletionResult *>
 swift::ide::copyCodeCompletionResults(CodeCompletionResultSink &targetSink,
                                       CodeCompletionResultSink &sourceSink,
                                       bool onlyTypes,
@@ -7485,8 +7524,8 @@ swift::ide::copyCodeCompletionResults(CodeCompletionResultSink &targetSink,
                               sourceSink.Results.end());
   }
 
-  return llvm::makeArrayRef(targetSink.Results.data() + startSize,
-                            targetSink.Results.size() - startSize);
+  return llvm::makeMutableArrayRef(targetSink.Results.data() + startSize,
+                                   targetSink.Results.size() - startSize);
 }
 
 void SimpleCachingCodeCompletionConsumer::handleResultsAndModules(
@@ -7515,9 +7554,11 @@ void SimpleCachingCodeCompletionConsumer::handleResultsAndModules(
       context.Cache.set(R.Key, *V);
     }
     assert(V.hasValue());
-    auto newItems = copyCodeCompletionResults(context.getResultSink(), (*V)->Sink,
-                                              R.OnlyTypes, R.OnlyPrecedenceGroups);
-    postProcessResults(newItems, context.CodeCompletionKind, DC);
+    auto newItems =
+        copyCodeCompletionResults(context.getResultSink(), (*V)->Sink,
+                                  R.OnlyTypes, R.OnlyPrecedenceGroups);
+    postProcessResults(newItems, context.CodeCompletionKind, DC,
+                       &context.getResultSink());
   }
 
   handleResults(context.takeResults());

test/IDE/complete_flair_cache.swift

@@ -0,0 +1,46 @@
+// %empty-directory(%t)
+// %empty-directory(%t.ccp)
+
+// Repeat twice to ensure the cache file is not affected by the flair.
+// RUN: %target-swift-ide-test -batch-code-completion -source-filename %s -parse-as-library -filecheck %raw-FileCheck -completion-output-dir %t -completion-cache-path=%t.ccp
+// RUN: %target-swift-ide-test -batch-code-completion -source-filename %s -parse-as-library -filecheck %raw-FileCheck -completion-output-dir %t -completion-cache-path=%t.ccp
+
+struct MyStruct {}
+protocol MyProtocol {}
+
+func testType() {
+    let a: #^TYPE^#
+// TYPE: Begin completions
+// TYPE-NOT: Keyword[import]
+// TYPE-NOT: Keyword[struct]
+// TYPE-NOT: Keyword[defer]
+// TYPE-DAG: Decl[Struct]/CurrModule: MyStruct[#MyStruct#]; name=MyStruct
+// TYPE-DAG: Decl[Protocol]/CurrModule: MyProtocol[#MyProtocol#]; name=MyProtocol
+// TYPE-DAG: Decl[Protocol]/OtherModule[Swift]/IsSystem: IteratorProtocol[#IteratorProtocol#]; name=IteratorProtocol
+// TYPE-DAG: Decl[Struct]/OtherModule[Swift]/IsSystem: Int[#Int#]; name=Int
+// TYPE: End completions
+}
+
+func testType() {
+    #^EXPR^#
+// EXPR: Begin completions
+// EXPR-DAG: Keyword[import]/None/Flair[RareKeyword]: import; name=import
+// EXPR-DAG: Keyword[struct]/None/Flair[RareKeyword]: struct; name=struct
+// EXPR-DAG: Keyword[defer]/None: defer; name=defer
+// EXPR-DAG: Decl[Struct]/CurrModule: MyStruct[#MyStruct#]; name=MyStruct
+// EXPR-DAG: Decl[Protocol]/CurrModule/Flair[RareType]: MyProtocol[#MyProtocol#]; name=MyProtocol
+// EXPR-DAG: Decl[Protocol]/OtherModule[Swift]/Flair[RareType]/IsSystem: IteratorProtocol[#IteratorProtocol#]; name=IteratorProtocol
+// EXPR-DAG: Decl[Struct]/OtherModule[Swift]/IsSystem: Int[#Int#]; name=Int
+// EXPR: End completions
+}
+
+#^TOPLEVEL^#
+// TOPLEVEL: Begin completions
+// TOPLEVEL-DAG: Keyword[import]/None:       import; name=import
+// TOPLEVEL-DAG: Keyword[struct]/None/Flair[CommonKeyword]: struct; name=struct
+// TOPLEVEL-DAG: Keyword[defer]/None/Flair[ExprAtFileScope]: defer; name=defer
+// TOPLEVEL-DAG: Decl[Struct]/CurrModule/Flair[ExprAtFileScope]: MyStruct[#MyStruct#]; name=MyStruct
+// TOPLEVEL-DAG: Decl[Protocol]/CurrModule/Flair[RareType,ExprAtFileScope]: MyProtocol[#MyProtocol#]; name=MyProtocol
+// TOPLEVEL-DAG: Decl[Protocol]/OtherModule[Swift]/Flair[RareType,ExprAtFileScope]/IsSystem: IteratorProtocol[#IteratorProtocol#]; name=IteratorProtocol
+// TOPLEVEL-DAG: Decl[Struct]/OtherModule[Swift]/Flair[ExprAtFileScope]/IsSystem: Int[#Int#]; name=Int
+// TOPLEVEL: End completions