Merge pull request #40724 from ahoppen/pr/sourcekit-completion-reference-to-swift-result

[CodeComplete] Make SourceKit::CodeCompletion::Completion store a reference to the underlying swift result instead of extending that type

Author: ahoppen

include/swift/IDE/CodeCompletion.h

@@ -1003,8 +1003,7 @@ class ImportDepth {
 class CodeCompletionContext {
   friend class CodeCompletionResultBuilder;
 
-  /// A set of current completion results, not yet delivered to the
-  /// consumer.
+  /// A set of current completion results.
   CodeCompletionResultSink CurrentResults;
 
 public:
@@ -1072,13 +1071,10 @@ class CodeCompletionContext {
   /// Allocate a string owned by the code completion context.
   StringRef copyString(StringRef Str);
 
-  /// Return current code completion results.
-  MutableArrayRef<CodeCompletionResult *> takeResults();
-
   /// Sort code completion results in an implementation-defined order
   /// in place.
-  static void sortCompletionResults(
-      MutableArrayRef<CodeCompletionResult *> Results);
+  static std::vector<CodeCompletionResult *>
+  sortCompletionResults(ArrayRef<CodeCompletionResult *> Results);
 
   CodeCompletionResultSink &getResultSink() {
     return CurrentResults;

include/swift/IDE/CompletionInstance.h

@@ -53,7 +53,7 @@ struct CompletionInstanceResult {
 
 /// The results returned from \c CompletionInstance::codeComplete.
 struct CodeCompleteResult {
-  MutableArrayRef<CodeCompletionResult *> Results;
+  CodeCompletionResultSink &ResultSink;
   SwiftCompletionInfo &Info;
 };
 

lib/IDE/CodeCompletion.cpp

@@ -1354,18 +1354,6 @@ void CodeCompletionResultBuilder::finishResult() {
     Sink.Results.push_back(takeResult());
 }
 
-
-MutableArrayRef<CodeCompletionResult *> CodeCompletionContext::takeResults() {
-  // Copy pointers to the results.
-  const size_t Count = CurrentResults.Results.size();
-  CodeCompletionResult **Results =
-      CurrentResults.Allocator->Allocate<CodeCompletionResult *>(Count);
-  std::copy(CurrentResults.Results.begin(), CurrentResults.Results.end(),
-            Results);
-  CurrentResults.Results.clear();
-  return MutableArrayRef<CodeCompletionResult *>(Results, Count);
-}
-
 Optional<unsigned> CodeCompletionString::getFirstTextChunkIndex(
     bool includeLeadingPunctuation) const {
   for (auto i : indices(getChunks())) {
@@ -1451,25 +1439,29 @@ CodeCompletionString::getFirstTextChunk(bool includeLeadingPunctuation) const {
   return StringRef();
 }
 
-void CodeCompletionContext::sortCompletionResults(
-    MutableArrayRef<CodeCompletionResult *> Results) {
+std::vector<CodeCompletionResult *>
+CodeCompletionContext::sortCompletionResults(
+    ArrayRef<CodeCompletionResult *> Results) {
+  std::vector<CodeCompletionResult *> SortedResults(Results.begin(),
+                                                    Results.end());
   struct ResultAndName {
     CodeCompletionResult *result;
     std::string name;
   };
 
   // Caching the name of each field is important to avoid unnecessary calls to
   // CodeCompletionString::getName().
-  std::vector<ResultAndName> nameCache(Results.size());
-  for (unsigned i = 0, n = Results.size(); i < n; ++i) {
-    auto *result = Results[i];
+  std::vector<ResultAndName> nameCache(SortedResults.size());
+  for (unsigned i = 0, n = SortedResults.size(); i < n; ++i) {
+    auto *result = SortedResults[i];
     nameCache[i].result = result;
     llvm::raw_string_ostream OS(nameCache[i].name);
     printCodeCompletionResultFilterName(*result, OS);
     OS.flush();
   }
 
-  // Sort nameCache, and then transform Results to return the pointers in order.
+  // Sort nameCache, and then transform SortedResults to return the pointers in
+  // order.
   std::sort(nameCache.begin(), nameCache.end(),
             [](const ResultAndName &LHS, const ResultAndName &RHS) {
               int Result = StringRef(LHS.name).compare_insensitive(RHS.name);
@@ -1480,8 +1472,9 @@ void CodeCompletionContext::sortCompletionResults(
               return Result < 0;
             });
 
-  llvm::transform(nameCache, Results.begin(),
+  llvm::transform(nameCache, SortedResults.begin(),
                   [](const ResultAndName &entry) { return entry.result; });
+  return SortedResults;
 }
 
 namespace {

lib/IDE/CompletionInstance.cpp

@@ -709,9 +709,8 @@ void swift::ide::CompletionInstance::codeComplete(
           CancellationFlag->load(std::memory_order_relaxed)) {
         Callback(ResultType::cancelled());
       } else {
-        MutableArrayRef<CodeCompletionResult *> Results = context.takeResults();
         assert(SwiftContext.swiftASTContext);
-        Callback(ResultType::success({Results, SwiftContext}));
+        Callback(ResultType::success({context.getResultSink(), SwiftContext}));
       }
     }
   };
@@ -736,7 +735,8 @@ void swift::ide::CompletionInstance::codeComplete(
                 SwiftCompletionInfo Info{&CI.getASTContext(),
                                          &CI.getInvocation(),
                                          &CompletionContext};
-                DeliverTransformed(ResultType::success({/*Results=*/{}, Info}));
+                CodeCompletionResultSink ResultSink;
+                DeliverTransformed(ResultType::success({ResultSink, Info}));
                 return;
               }
 
@@ -753,7 +753,8 @@ void swift::ide::CompletionInstance::codeComplete(
                 SwiftCompletionInfo Info{&CI.getASTContext(),
                                          &CI.getInvocation(),
                                          &CompletionContext};
-                DeliverTransformed(ResultType::success({/*Results=*/{}, Info}));
+                CodeCompletionResultSink ResultSink;
+                DeliverTransformed(ResultType::success({ResultSink, Info}));
               }
             },
             Callback);

lib/IDE/REPLCodeCompletion.cpp

@@ -178,9 +178,9 @@ class REPLCodeCompletionConsumer : public SimpleCachingCodeCompletionConsumer {
       : Completions(Completions) {}
 
   void handleResults(CodeCompletionContext &context) override {
-    MutableArrayRef<CodeCompletionResult *> Results = context.takeResults();
-    CodeCompletionContext::sortCompletionResults(Results);
-    for (auto Result : Results) {
+    auto SortedResults = CodeCompletionContext::sortCompletionResults(
+        context.getResultSink().Results);
+    for (auto Result : SortedResults) {
       std::string InsertableString = toInsertableString(Result);
       if (StringRef(InsertableString).startswith(Completions.Prefix)) {
         llvm::SmallString<128> PrintedResult;

tools/SourceKit/lib/SwiftLang/CodeCompletion.h

@@ -23,11 +23,12 @@ namespace SourceKit {
 namespace CodeCompletion {
 
 using swift::ide::CodeCompletionDeclKind;
+using swift::ide::CodeCompletionFlair;
 using swift::ide::CodeCompletionKeywordKind;
 using swift::ide::CodeCompletionLiteralKind;
-using swift::ide::SemanticContextKind;
-using swift::ide::CodeCompletionFlair;
+using swift::ide::CodeCompletionOperatorKind;
 using swift::ide::CodeCompletionString;
+using swift::ide::SemanticContextKind;
 using SwiftResult = swift::ide::CodeCompletionResult;
 using swift::ide::CompletionKind;
 
@@ -75,7 +76,8 @@ struct NameStyle {
 ///
 /// Extends a \c swift::ide::CodeCompletionResult with extra fields that are
 /// filled in by SourceKit. Generally stored in an \c CompletionSink.
-class Completion : public SwiftResult {
+class Completion {
+  const SwiftResult &base;
   void *opaqueCustomKind = nullptr;
   Optional<uint8_t> moduleImportDepth;
   PopularityFactor popularityFactor;
@@ -89,9 +91,12 @@ class Completion : public SwiftResult {
 
   /// Wraps \p base with an \c Completion.  The \p name and \p description
   /// should outlive the result, generally by being stored in the same
-  /// \c CompletionSink.
-  Completion(SwiftResult base, StringRef name, StringRef description)
-      : SwiftResult(base), name(name), description(description) {}
+  /// \c CompletionSink or in a sink that was adopted by the sink that this
+  /// \c Compleiton is being stored in.
+  Completion(const SwiftResult &base, StringRef name, StringRef description)
+      : base(base), name(name), description(description) {}
+
+  const SwiftResult &getSwiftResult() const { return base; }
 
   bool hasCustomKind() const { return opaqueCustomKind; }
   void *getCustomKind() const { return opaqueCustomKind; }
@@ -102,6 +107,63 @@ class Completion : public SwiftResult {
   /// A popularity factory in the range [-1, 1]. The higher the value, the more
   /// 'popular' this result is.  0 indicates unknown.
   PopularityFactor getPopularityFactor() const { return popularityFactor; }
+
+  // MARK: Methods that forward to the SwiftResult
+
+  SwiftResult::ResultKind getKind() const { return getSwiftResult().getKind(); }
+
+  CodeCompletionDeclKind getAssociatedDeclKind() const {
+    return getSwiftResult().getAssociatedDeclKind();
+  }
+
+  CodeCompletionLiteralKind getLiteralKind() const {
+    return getSwiftResult().getLiteralKind();
+  }
+
+  CodeCompletionKeywordKind getKeywordKind() const {
+    return getSwiftResult().getKeywordKind();
+  }
+
+  bool isOperator() const { return getSwiftResult().isOperator(); }
+
+  CodeCompletionOperatorKind getOperatorKind() const {
+    return getSwiftResult().getOperatorKind();
+  }
+
+  bool isSystem() const { return getSwiftResult().isSystem(); }
+
+  SwiftResult::ExpectedTypeRelation getExpectedTypeRelation() const {
+    return getSwiftResult().getExpectedTypeRelation();
+  }
+
+  SemanticContextKind getSemanticContext() const {
+    return getSwiftResult().getSemanticContext();
+  }
+
+  CodeCompletionFlair getFlair() const { return getSwiftResult().getFlair(); }
+
+  bool isNotRecommended() const { return getSwiftResult().isNotRecommended(); }
+
+  unsigned getNumBytesToErase() const {
+    return getSwiftResult().getNumBytesToErase();
+  }
+
+  CodeCompletionString *getCompletionString() const {
+    return getSwiftResult().getCompletionString();
+  }
+
+  StringRef getModuleName() const { return getSwiftResult().getModuleName(); }
+
+  StringRef getBriefDocComment() const {
+    return getSwiftResult().getBriefDocComment();
+  }
+
+  ArrayRef<StringRef> getAssociatedUSRs() const {
+    return getSwiftResult().getAssociatedUSRs();
+  }
+
+  /// Allow "upcasting" the completion result to a SwiftResult.
+  operator const SwiftResult &() const { return getSwiftResult(); }
 };
 
 /// Storage sink for \c Completion objects.
@@ -123,9 +185,9 @@ struct CompletionSink {
 
 class CompletionBuilder {
   CompletionSink &sink;
-  SwiftResult &current;
+  const SwiftResult &base;
   bool modified = false;
-  Completion::ExpectedTypeRelation typeRelation;
+  SwiftResult::ExpectedTypeRelation typeRelation;
   SemanticContextKind semanticContext;
   CodeCompletionFlair flair;
   CodeCompletionString *completionString;
@@ -135,7 +197,7 @@ class CompletionBuilder {
   PopularityFactor popularityFactor;
 
 public:
-  CompletionBuilder(CompletionSink &sink, SwiftResult &base);
+  CompletionBuilder(CompletionSink &sink, const SwiftResult &base);
 
   void setCustomKind(void *opaqueCustomKind) { customKind = opaqueCustomKind; }
 
@@ -144,7 +206,7 @@ class CompletionBuilder {
     moduleImportDepth = value;
   }
 
-  void setExpectedTypeRelation(Completion::ExpectedTypeRelation Relation) {
+  void setExpectedTypeRelation(SwiftResult::ExpectedTypeRelation Relation) {
     modified = true;
     typeRelation = Relation;
   }
@@ -233,8 +295,8 @@ struct FilterRules {
   llvm::StringMap<bool> hideByFilterName;
   llvm::StringMap<bool> hideByDescription;
 
-  bool hideCompletion(Completion *completion) const;
-  bool hideCompletion(SwiftResult *completion, StringRef name,
+  bool hideCompletion(const Completion &completion) const;
+  bool hideCompletion(const SwiftResult &completion, StringRef name,
                       StringRef description, void *customKind = nullptr) const;
   bool hideFilterName(StringRef name) const;
 };