[CodeCompletion] Split result delivery into a result colleciton and consumer phase

This will allow us to run two different completion kinds and deliver results from both of them.

Also: Compute a unified type context for global lookup. Previously, we always used the expected type context of the last lookup. But really, we should be considering all possible types from all constraint system solutions when computing code completion results from the cache.

include/swift/IDE/AfterPoundExprCompletion.h

@@ -47,8 +47,7 @@ public:
       : CompletionExpr(CompletionExpr), DC(DC), ParentStmtKind(ParentStmtKind) {
   }
 
-  void deliverResults(ide::CodeCompletionContext &CompletionCtx,
+  void collectResults(ide::CodeCompletionContext &CompletionCtx);
-                      CodeCompletionConsumer &Consumer);
 };
 
 } // end namespace ide

include/swift/IDE/ArgumentCompletion.h

@@ -93,9 +93,8 @@ public:
   /// function signature instead of suggesting individual labels. Used when
   /// completing after the opening '(' of a function call \param Loc The
   /// location of the code completion token
-  void deliverResults(bool IncludeSignature, SourceLoc Loc, DeclContext *DC,
+  void collectResults(bool IncludeSignature, SourceLoc Loc, DeclContext *DC,
-                      CodeCompletionContext &CompletionCtx,
+                      CodeCompletionContext &CompletionCtx);
-                      CodeCompletionConsumer &Consumer);
 };
 
 } // end namespace ide

include/swift/IDE/CodeCompletion.h

@@ -65,9 +65,10 @@ void postProcessCompletionResults(
     MutableArrayRef<CodeCompletionResult *> results, CompletionKind Kind,
     const DeclContext *DC, CodeCompletionResultSink *Sink);
 
-void deliverCompletionResults(CodeCompletionContext &CompletionContext,
+void collectCompletionResults(CodeCompletionContext &CompletionContext,
                               CompletionLookup &Lookup, DeclContext *DC,
-                              CodeCompletionConsumer &Consumer);
+                              const ExpectedTypeContext &TypeContext,
+                              bool CanCurrDeclContextHandleAsync);
 
 /// Create a factory for code completion callbacks.
 IDEInspectionCallbacksFactory *

include/swift/IDE/CodeCompletionConsumer.h

@@ -22,30 +22,9 @@ namespace ide {
 struct RequestedCachedModule;
 
 /// An abstract base class for consumers of code completion results.
-/// \see \c SimpleCachingCodeCompletionConsumer.
 class CodeCompletionConsumer {
 public:
   virtual ~CodeCompletionConsumer() {}
-  virtual void
-  handleResultsAndModules(CodeCompletionContext &context,
-                          ArrayRef<RequestedCachedModule> requestedModules,
-                          const ExpectedTypeContext *TypeContext,
-                          const DeclContext *DC,
-                          bool CanCurrDeclContextHandleAsync) = 0;
-};
-
-/// A simplified code completion consumer interface that clients can use to get
-/// CodeCompletionResults with automatic caching of top-level completions from
-/// imported modules.
-struct SimpleCachingCodeCompletionConsumer : public CodeCompletionConsumer {
-
-  // Implement the CodeCompletionConsumer interface.
-  void handleResultsAndModules(CodeCompletionContext &context,
-                               ArrayRef<RequestedCachedModule> requestedModules,
-                               const ExpectedTypeContext *TypeContext,
-                               const DeclContext *DCForModules,
-                               bool CanCurrDeclContextHandleAsync) override;
-
   /// Clients should override this method to receive \p Results.
   virtual void handleResults(CodeCompletionContext &context) = 0;
 };

include/swift/IDE/CodeCompletionContext.h

@@ -20,6 +20,7 @@ namespace swift {
 namespace ide {
 
 class CodeCompletionCache;
+struct RequestedCachedModule;
 
 class CodeCompletionContext {
   friend class CodeCompletionResultBuilder;
@@ -103,6 +104,13 @@ public:
   sortCompletionResults(ArrayRef<CodeCompletionResult *> Results);
 
   CodeCompletionResultSink &getResultSink() { return CurrentResults; }
+
+  /// Add code completion results from the given requested modules to this
+  /// context.
+  void addResultsFromModules(ArrayRef<RequestedCachedModule> RequestedModules,
+                             const ExpectedTypeContext &TypeContext,
+                             const DeclContext *DC,
+                             bool CanCurrDeclContextHandleAsync);
 };
 
 } // end namespace ide

include/swift/IDE/CodeCompletionResultType.h

@@ -74,6 +74,26 @@ public:
     }
   }
 
+  /// Form a union of this expected type context with \p Other.
+  ///
+  /// Any possible type from either type context will be considered a possible
+  /// type in the merged type context.
+  void merge(const ExpectedTypeContext &Other) {
+    PossibleTypes.append(Other.PossibleTypes);
+
+    // We can't merge ideal types. Setting to a null type is the best thing we
+    // can do if they differ.
+    if (IdealType.isNull() != Other.IdealType.isNull()) {
+      IdealType = Type();
+    } else if (IdealType && Other.IdealType &&
+               !IdealType->isEqual(Other.IdealType)) {
+      IdealType = Type();
+    }
+    IsImplicitSingleExpressionReturn |= Other.IsImplicitSingleExpressionReturn;
+    PreferNonVoid &= Other.PreferNonVoid;
+    ExpectedCustomAttributeKinds |= Other.ExpectedCustomAttributeKinds;
+  }
+
   Type getIdealType() const { return IdealType; }
 
   void setIdealType(Type IdealType) { this->IdealType = IdealType; }

include/swift/IDE/ExprCompletion.h

@@ -90,9 +90,8 @@ public:
         AddUnresolvedMemberCompletions(AddUnresolvedMemberCompletions) {}
 
   /// \param CCLoc The location of the code completion token.
-  void deliverResults(SourceLoc CCLoc,
+  void collectResults(SourceLoc CCLoc,
-                      ide::CodeCompletionContext &CompletionCtx,
+                      ide::CodeCompletionContext &CompletionCtx);
-                      CodeCompletionConsumer &Consumer);
 };
 
 } // end namespace ide

include/swift/IDE/KeyPathCompletion.h

@@ -37,9 +37,8 @@ class KeyPathTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
 public:
   KeyPathTypeCheckCompletionCallback(KeyPathExpr *KeyPath) : KeyPath(KeyPath) {}
 
-  void deliverResults(DeclContext *DC, SourceLoc DotLoc,
+  void collectResults(DeclContext *DC, SourceLoc DotLoc,
-                      ide::CodeCompletionContext &CompletionCtx,
+                      ide::CodeCompletionContext &CompletionCtx);
-                      CodeCompletionConsumer &Consumer);
 };
 
 } // end namespace ide

include/swift/IDE/PostfixCompletion.h

@@ -103,10 +103,9 @@ public:
   ///                         \p DotLoc is invalid
   /// \param HasLeadingSpace Whether there is a space separating the exiting
   ///                        expression and the code completion token.
-  void deliverResults(SourceLoc DotLoc, bool IsInSelector,
+  void collectResults(SourceLoc DotLoc, bool IsInSelector,
                       bool IncludeOperators, bool HasLeadingSpace,
-                      CodeCompletionContext &CompletionCtx,
+                      CodeCompletionContext &CompletionCtx);
-                      CodeCompletionConsumer &Consumer);
 };
 
 } // end namespace ide

include/swift/IDE/UnresolvedMemberCompletion.h

@@ -59,9 +59,8 @@ public:
       CodeCompletionExpr *CompletionExpr, DeclContext *DC)
       : CompletionExpr(CompletionExpr), DC(DC) {}
 
-  void deliverResults(DeclContext *DC, SourceLoc DotLoc,
+  void collectResults(DeclContext *DC, SourceLoc DotLoc,
-                      ide::CodeCompletionContext &CompletionCtx,
+                      ide::CodeCompletionContext &CompletionCtx);
-                      CodeCompletionConsumer &Consumer);
 };
 
 } // end namespace ide

lib/IDE/AfterPoundExprCompletion.cpp

@@ -37,15 +37,18 @@ void AfterPoundExprCompletion::sawSolutionImpl(const constraints::Solution &S) {
   }
 }
 
-void AfterPoundExprCompletion::deliverResults(
+void AfterPoundExprCompletion::collectResults(
-    ide::CodeCompletionContext &CompletionCtx,
+    ide::CodeCompletionContext &CompletionCtx) {
-    CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
 
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
 
+  // The type context that is being used for global results.
+  ExpectedTypeContext UnifiedTypeContext;
+  UnifiedTypeContext.setPreferNonVoid(true);
+
   for (auto &Result : Results) {
     Lookup.setExpectedTypes({Result.ExpectedTy},
                             Result.IsImplicitSingleExpressionReturn,
@@ -53,7 +56,10 @@ void AfterPoundExprCompletion::deliverResults(
     Lookup.addPoundAvailable(ParentStmtKind);
     Lookup.addObjCPoundKeywordCompletions(/*needPound=*/false);
     Lookup.getMacroCompletions(CodeCompletionMacroRole::Expression);
+
+    UnifiedTypeContext.merge(*Lookup.getExpectedTypeContext());
   }
 
-  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+  collectCompletionResults(CompletionCtx, Lookup, DC, UnifiedTypeContext,
+                           /*CanCurrDeclContextHandleAsync=*/false);
 }

lib/IDE/ArgumentCompletion.cpp

@@ -295,10 +295,9 @@ void ArgumentTypeCheckCompletionCallback::computeShadowedDecls(
   }
 }
 
-void ArgumentTypeCheckCompletionCallback::deliverResults(
+void ArgumentTypeCheckCompletionCallback::collectResults(
     bool IncludeSignature, SourceLoc Loc, DeclContext *DC,
-    ide::CodeCompletionContext &CompletionCtx,
+    ide::CodeCompletionContext &CompletionCtx) {
-    CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
@@ -402,5 +401,7 @@ void ArgumentTypeCheckCompletionCallback::deliverResults(
     addExprKeywords(CompletionCtx.getResultSink(), DC);
   }
 
-  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+  collectCompletionResults(CompletionCtx, Lookup, DC,
+                           *Lookup.getExpectedTypeContext(),
+                           Lookup.canCurrDeclContextHandleAsync());
 }

lib/IDE/CMakeLists.txt

@@ -5,7 +5,6 @@ add_swift_host_library(swiftIDE STATIC
   ArgumentCompletion.cpp
   CodeCompletion.cpp
   CodeCompletionCache.cpp
-  CodeCompletionConsumer.cpp
   CodeCompletionContext.cpp
   CodeCompletionDiagnostics.cpp
   CodeCompletionResult.cpp

lib/IDE/CodeCompletion.cpp

@@ -1298,9 +1298,10 @@ void swift::ide::postProcessCompletionResults(
   }
 }
 
-void swift::ide::deliverCompletionResults(
+void swift::ide::collectCompletionResults(
     CodeCompletionContext &CompletionContext, CompletionLookup &Lookup,
-    DeclContext *DC, CodeCompletionConsumer &Consumer) {
+    DeclContext *DC, const ExpectedTypeContext &TypeContext,
+    bool CanCurrDeclContextHandleAsync) {
   auto &SF = *DC->getParentSourceFile();
   llvm::SmallPtrSet<Identifier, 8> seenModuleNames;
   std::vector<RequestedCachedModule> RequestedModules;
@@ -1420,9 +1421,8 @@ void swift::ide::deliverCompletionResults(
                                CompletionContext.CodeCompletionKind, DC,
                                /*Sink=*/nullptr);
 
-  Consumer.handleResultsAndModules(CompletionContext, RequestedModules,
+  CompletionContext.addResultsFromModules(RequestedModules, TypeContext, DC,
-                                   Lookup.getExpectedTypeContext(), DC,
+                                          CanCurrDeclContextHandleAsync);
-                                   Lookup.canCurrDeclContextHandleAsync());
 }
 
 bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
@@ -1489,8 +1489,9 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
 
     bool IncludeOperators = (Kind == CompletionKind::PostfixExpr);
 
-    Lookup.deliverResults(DotLoc, isInsideObjCSelector(), IncludeOperators,
+    Lookup.collectResults(DotLoc, isInsideObjCSelector(), IncludeOperators,
-                          HasSpace, CompletionContext, Consumer);
+                          HasSpace, CompletionContext);
+    Consumer.handleResults(CompletionContext);
     return true;
   }
   case CompletionKind::UnresolvedMember: {
@@ -1502,7 +1503,8 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     typeCheckWithLookup(Lookup);
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
-    Lookup.deliverResults(CurDeclContext, DotLoc, CompletionContext, Consumer);
+    Lookup.collectResults(CurDeclContext, DotLoc, CompletionContext);
+    Consumer.handleResults(CompletionContext);
     return true;
   }
   case CompletionKind::KeyPathExprSwift: {
@@ -1514,7 +1516,8 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     KeyPathTypeCheckCompletionCallback Lookup(KeyPath);
     typeCheckWithLookup(Lookup);
 
-    Lookup.deliverResults(CurDeclContext, DotLoc, CompletionContext, Consumer);
+    Lookup.collectResults(CurDeclContext, DotLoc, CompletionContext);
+    Consumer.handleResults(CompletionContext);
     return true;
   }
   case CompletionKind::PostfixExprParen:
@@ -1525,8 +1528,9 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
                                                CurDeclContext);
     typeCheckWithLookup(Lookup);
 
-    Lookup.deliverResults(ShouldCompleteCallPatternAfterParen, CompletionLoc,
+    Lookup.collectResults(ShouldCompleteCallPatternAfterParen, CompletionLoc,
-                          CurDeclContext, CompletionContext, Consumer);
+                          CurDeclContext, CompletionContext);
+    Consumer.handleResults(CompletionContext);
     return true;
   }
   case CompletionKind::AccessorBeginning:
@@ -1556,7 +1560,8 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
 
     SourceLoc CCLoc = P.Context.SourceMgr.getIDEInspectionTargetLoc();
-    Lookup.deliverResults(CCLoc, CompletionContext, Consumer);
+    Lookup.collectResults(CCLoc, CompletionContext);
+    Consumer.handleResults(CompletionContext);
     return true;
   }
   case CompletionKind::AfterPoundExpr: {
@@ -1569,7 +1574,8 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
 
-    Lookup.deliverResults(CompletionContext, Consumer);
+    Lookup.collectResults(CompletionContext);
+    Consumer.handleResults(CompletionContext);
     return true;
   }
   default:
@@ -2041,7 +2047,10 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
     break;
   }
 
-  deliverCompletionResults(CompletionContext, Lookup, CurDeclContext, Consumer);
+  collectCompletionResults(CompletionContext, Lookup, CurDeclContext,
+                           *Lookup.getExpectedTypeContext(),
+                           Lookup.canCurrDeclContextHandleAsync());
+  Consumer.handleResults(CompletionContext);
 }
 
 namespace {

lib/IDE/CodeCompletionConsumer.cpp

@@ -1,160 +0,0 @@
-//===--- CodeCompletionConsumer.cpp ---------------------------------------===//
-//
-// This source file is part of the Swift.org open source project
-//
-// Copyright (c) 2022 Apple Inc. and the Swift project authors
-// Licensed under Apache License v2.0 with Runtime Library Exception
-//
-// See https://swift.org/LICENSE.txt for license information
-// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
-//
-//===----------------------------------------------------------------------===//
-
-#include "swift/IDE/CodeCompletionConsumer.h"
-#include "swift/IDE/CodeCompletionCache.h"
-
-using namespace swift;
-using namespace swift::ide;
-
-static MutableArrayRef<CodeCompletionResult *> copyCodeCompletionResults(
-    CodeCompletionResultSink &targetSink, CodeCompletionCache::Value &source,
-    CodeCompletionFilter filter, const ExpectedTypeContext *TypeContext,
-    const DeclContext *DC, bool CanCurrDeclContextHandleAsync) {
-  assert(filter && "Should never have an empty filter");
-
-  // We will be adding foreign results (from another sink) into TargetSink.
-  // TargetSink should have an owning pointer to the allocator that keeps the
-  // results alive.
-  targetSink.ForeignAllocators.push_back(source.Allocator);
-  auto startSize = targetSink.Results.size();
-
-  CodeCompletionMacroRoles expectedMacroRoles = getCompletionMacroRoles(filter);
-  std::function<bool(const ContextFreeCodeCompletionResult *)>
-      shouldIncludeResult =
-          [filter, expectedMacroRoles](
-              const ContextFreeCodeCompletionResult *R) -> bool {
-    if (R->getKind() != CodeCompletionResultKind::Declaration)
-      return false;
-
-    switch (R->getAssociatedDeclKind()) {
-    case CodeCompletionDeclKind::EnumElement:
-    case CodeCompletionDeclKind::Constructor:
-    case CodeCompletionDeclKind::Destructor:
-    case CodeCompletionDeclKind::Subscript:
-    case CodeCompletionDeclKind::StaticMethod:
-    case CodeCompletionDeclKind::InstanceMethod:
-    case CodeCompletionDeclKind::PrefixOperatorFunction:
-    case CodeCompletionDeclKind::PostfixOperatorFunction:
-    case CodeCompletionDeclKind::InfixOperatorFunction:
-    case CodeCompletionDeclKind::FreeFunction:
-    case CodeCompletionDeclKind::StaticVar:
-    case CodeCompletionDeclKind::InstanceVar:
-    case CodeCompletionDeclKind::LocalVar:
-    case CodeCompletionDeclKind::GlobalVar:
-      return filter.contains(CodeCompletionFilterFlag::Expr);
-
-    case CodeCompletionDeclKind::Module:
-    case CodeCompletionDeclKind::Class:
-    case CodeCompletionDeclKind::Actor:
-    case CodeCompletionDeclKind::Struct:
-    case CodeCompletionDeclKind::Enum:
-    case CodeCompletionDeclKind::Protocol:
-    case CodeCompletionDeclKind::TypeAlias:
-    case CodeCompletionDeclKind::AssociatedType:
-    case CodeCompletionDeclKind::GenericTypeParam:
-      return filter.contains(CodeCompletionFilterFlag::Type);
-
-    case CodeCompletionDeclKind::PrecedenceGroup:
-      return filter.contains(CodeCompletionFilterFlag::PrecedenceGroup);
-
-    case CodeCompletionDeclKind::Macro:
-      return (bool)(R->getMacroRoles() & expectedMacroRoles);
-    }
-
-    llvm_unreachable("Unhandled associated decl kind");
-  };
-
-  USRBasedTypeContext USRTypeContext(TypeContext, source.USRTypeArena);
-
-  for (auto contextFreeResult : source.Results) {
-    if (!shouldIncludeResult(contextFreeResult)) {
-      continue;
-    }
-
-    CodeCompletionResultTypeRelation typeRelation =
-        contextFreeResult->calculateContextualTypeRelation(DC, TypeContext,
-                                                           &USRTypeContext);
-    ContextualNotRecommendedReason notRecommendedReason =
-        contextFreeResult->calculateContextualNotRecommendedReason(
-            ContextualNotRecommendedReason::None,
-            CanCurrDeclContextHandleAsync);
-
-    auto contextualResult = new (*targetSink.Allocator) CodeCompletionResult(
-        *contextFreeResult, SemanticContextKind::OtherModule,
-        CodeCompletionFlair(),
-        /*numBytesToErase=*/0, typeRelation, notRecommendedReason);
-    targetSink.Results.push_back(contextualResult);
-  }
-
-  return llvm::makeMutableArrayRef(targetSink.Results.data() + startSize,
-                                   targetSink.Results.size() - startSize);
-}
-
-void SimpleCachingCodeCompletionConsumer::handleResultsAndModules(
-    CodeCompletionContext &context,
-    ArrayRef<RequestedCachedModule> requestedModules,
-    const ExpectedTypeContext *TypeContext, const DeclContext *DC,
-    bool CanCurrDeclContextHandleAsync) {
-
-  // Use the current SourceFile as the DeclContext so that we can use it to
-  // perform qualified lookup, and to get the correct visibility for
-  // @testable imports. Also it cannot use 'DC' since it would apply decl
-  // context changes to cached results.
-  const SourceFile *SF = DC->getParentSourceFile();
-
-  for (auto &R : requestedModules) {
-    // FIXME(thread-safety): lock the whole AST context.  We might load a
-    // module.
-    llvm::Optional<CodeCompletionCache::ValueRefCntPtr> V =
-        context.Cache.get(R.Key);
-    if (!V.has_value()) {
-      // No cached results found. Fill the cache.
-      V = context.Cache.createValue();
-      // Temporary sink in which we gather the result. The cache value retains
-      // the sink's allocator.
-      CodeCompletionResultSink Sink;
-      Sink.annotateResult = context.getAnnotateResult();
-      Sink.addInitsToTopLevel = context.getAddInitsToTopLevel();
-      Sink.enableCallPatternHeuristics = context.getCallPatternHeuristics();
-      Sink.includeObjectLiterals = context.includeObjectLiterals();
-      Sink.addCallWithNoDefaultArgs = context.addCallWithNoDefaultArgs();
-      Sink.setProduceContextFreeResults((*V)->USRTypeArena);
-      lookupCodeCompletionResultsFromModule(Sink, R.TheModule, R.Key.AccessPath,
-                                            R.Key.ResultsHaveLeadingDot, SF);
-      (*V)->Allocator = Sink.Allocator;
-      auto &CachedResults = (*V)->Results;
-      CachedResults.reserve(Sink.Results.size());
-      // Instead of copying the context free results out of the sink's allocator
-      // retain the sink's entire allocator (which also includes the contextual
-      // properities) and simply store pointers to the context free results that
-      // back the contextual results.
-      for (auto Result : Sink.Results) {
-        assert(
-            Result->getContextFreeResult().getResultType().isBackedByUSRs() &&
-            "Results stored in the cache should have their result types backed "
-            "by a USR because the cache might outlive the ASTContext the "
-            "results were created from.");
-        CachedResults.push_back(&Result->getContextFreeResult());
-      }
-      context.Cache.set(R.Key, *V);
-    }
-    assert(V.has_value());
-    auto newItems = copyCodeCompletionResults(context.getResultSink(), **V,
-                                              R.Filter, TypeContext, DC,
-                                              CanCurrDeclContextHandleAsync);
-    postProcessCompletionResults(newItems, context.CodeCompletionKind, DC,
-                                 &context.getResultSink());
-  }
-
-  handleResults(context);
-}

lib/IDE/CodeCompletionContext.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/IDE/CodeCompletionContext.h"
+#include "swift/IDE/CodeCompletionCache.h"
 
 using namespace swift;
 using namespace swift::ide;
@@ -34,3 +35,141 @@ CodeCompletionContext::sortCompletionResults(
 
   return SortedResults;
 }
+
+static MutableArrayRef<CodeCompletionResult *> copyCodeCompletionResults(
+    CodeCompletionResultSink &targetSink, CodeCompletionCache::Value &source,
+    CodeCompletionFilter filter, const ExpectedTypeContext *TypeContext,
+    const DeclContext *DC, bool CanCurrDeclContextHandleAsync) {
+  assert(filter && "Should never have an empty filter");
+
+  // We will be adding foreign results (from another sink) into TargetSink.
+  // TargetSink should have an owning pointer to the allocator that keeps the
+  // results alive.
+  targetSink.ForeignAllocators.push_back(source.Allocator);
+  auto startSize = targetSink.Results.size();
+
+  CodeCompletionMacroRoles expectedMacroRoles = getCompletionMacroRoles(filter);
+  std::function<bool(const ContextFreeCodeCompletionResult *)>
+      shouldIncludeResult =
+          [filter, expectedMacroRoles](
+              const ContextFreeCodeCompletionResult *R) -> bool {
+    if (R->getKind() != CodeCompletionResultKind::Declaration)
+      return false;
+
+    switch (R->getAssociatedDeclKind()) {
+    case CodeCompletionDeclKind::EnumElement:
+    case CodeCompletionDeclKind::Constructor:
+    case CodeCompletionDeclKind::Destructor:
+    case CodeCompletionDeclKind::Subscript:
+    case CodeCompletionDeclKind::StaticMethod:
+    case CodeCompletionDeclKind::InstanceMethod:
+    case CodeCompletionDeclKind::PrefixOperatorFunction:
+    case CodeCompletionDeclKind::PostfixOperatorFunction:
+    case CodeCompletionDeclKind::InfixOperatorFunction:
+    case CodeCompletionDeclKind::FreeFunction:
+    case CodeCompletionDeclKind::StaticVar:
+    case CodeCompletionDeclKind::InstanceVar:
+    case CodeCompletionDeclKind::LocalVar:
+    case CodeCompletionDeclKind::GlobalVar:
+      return filter.contains(CodeCompletionFilterFlag::Expr);
+
+    case CodeCompletionDeclKind::Module:
+    case CodeCompletionDeclKind::Class:
+    case CodeCompletionDeclKind::Actor:
+    case CodeCompletionDeclKind::Struct:
+    case CodeCompletionDeclKind::Enum:
+    case CodeCompletionDeclKind::Protocol:
+    case CodeCompletionDeclKind::TypeAlias:
+    case CodeCompletionDeclKind::AssociatedType:
+    case CodeCompletionDeclKind::GenericTypeParam:
+      return filter.contains(CodeCompletionFilterFlag::Type);
+
+    case CodeCompletionDeclKind::PrecedenceGroup:
+      return filter.contains(CodeCompletionFilterFlag::PrecedenceGroup);
+
+    case CodeCompletionDeclKind::Macro:
+      return (bool)(R->getMacroRoles() & expectedMacroRoles);
+    }
+
+    llvm_unreachable("Unhandled associated decl kind");
+  };
+
+  USRBasedTypeContext USRTypeContext(TypeContext, source.USRTypeArena);
+
+  for (auto contextFreeResult : source.Results) {
+    if (!shouldIncludeResult(contextFreeResult)) {
+      continue;
+    }
+
+    CodeCompletionResultTypeRelation typeRelation =
+        contextFreeResult->calculateContextualTypeRelation(DC, TypeContext,
+                                                           &USRTypeContext);
+    ContextualNotRecommendedReason notRecommendedReason =
+        contextFreeResult->calculateContextualNotRecommendedReason(
+            ContextualNotRecommendedReason::None,
+            CanCurrDeclContextHandleAsync);
+
+    auto contextualResult = new (*targetSink.Allocator) CodeCompletionResult(
+        *contextFreeResult, SemanticContextKind::OtherModule,
+        CodeCompletionFlair(),
+        /*numBytesToErase=*/0, typeRelation, notRecommendedReason);
+    targetSink.Results.push_back(contextualResult);
+  }
+
+  return llvm::makeMutableArrayRef(targetSink.Results.data() + startSize,
+                                   targetSink.Results.size() - startSize);
+}
+
+void CodeCompletionContext::addResultsFromModules(
+    ArrayRef<RequestedCachedModule> RequestedModules,
+    const ExpectedTypeContext &TypeContext, const DeclContext *DC,
+    bool CanCurrDeclContextHandleAsync) {
+  // Use the current SourceFile as the DeclContext so that we can use it to
+  // perform qualified lookup, and to get the correct visibility for
+  // @testable imports. Also it cannot use 'DC' since it would apply decl
+  // context changes to cached results.
+  const SourceFile *SF = DC->getParentSourceFile();
+
+  for (auto &R : RequestedModules) {
+    // FIXME(thread-safety): lock the whole AST context.  We might load a
+    // module.
+    llvm::Optional<CodeCompletionCache::ValueRefCntPtr> V = Cache.get(R.Key);
+    if (!V.has_value()) {
+      // No cached results found. Fill the cache.
+      V = Cache.createValue();
+      // Temporary sink in which we gather the result. The cache value retains
+      // the sink's allocator.
+      CodeCompletionResultSink Sink;
+      Sink.annotateResult = getAnnotateResult();
+      Sink.addInitsToTopLevel = getAddInitsToTopLevel();
+      Sink.enableCallPatternHeuristics = getCallPatternHeuristics();
+      Sink.includeObjectLiterals = includeObjectLiterals();
+      Sink.addCallWithNoDefaultArgs = addCallWithNoDefaultArgs();
+      Sink.setProduceContextFreeResults((*V)->USRTypeArena);
+      lookupCodeCompletionResultsFromModule(Sink, R.TheModule, R.Key.AccessPath,
+                                            R.Key.ResultsHaveLeadingDot, SF);
+      (*V)->Allocator = Sink.Allocator;
+      auto &CachedResults = (*V)->Results;
+      CachedResults.reserve(Sink.Results.size());
+      // Instead of copying the context free results out of the sink's allocator
+      // retain the sink's entire allocator (which also includes the contextual
+      // properities) and simply store pointers to the context free results that
+      // back the contextual results.
+      for (auto Result : Sink.Results) {
+        assert(
+            Result->getContextFreeResult().getResultType().isBackedByUSRs() &&
+            "Results stored in the cache should have their result types backed "
+            "by a USR because the cache might outlive the ASTContext the "
+            "results were created from.");
+        CachedResults.push_back(&Result->getContextFreeResult());
+      }
+      Cache.set(R.Key, *V);
+    }
+    assert(V.has_value());
+    auto newItems =
+        copyCodeCompletionResults(getResultSink(), **V, R.Filter, &TypeContext,
+                                  DC, CanCurrDeclContextHandleAsync);
+    postProcessCompletionResults(newItems, CodeCompletionKind, DC,
+                                 &getResultSink());
+  }
+}

lib/IDE/ExprCompletion.cpp

@@ -96,14 +96,18 @@ void ExprTypeCheckCompletionCallback::sawSolutionImpl(
   }
 }
 
-void ExprTypeCheckCompletionCallback::deliverResults(
+void ExprTypeCheckCompletionCallback::collectResults(
-    SourceLoc CCLoc, ide::CodeCompletionContext &CompletionCtx,
+    SourceLoc CCLoc, ide::CodeCompletionContext &CompletionCtx) {
-    CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
 
+  // The type context that is being used for global results.
+  ExpectedTypeContext UnifiedTypeContext;
+  UnifiedTypeContext.setPreferNonVoid(true);
+  bool UnifiedCanHandleAsync = false;
+
   for (auto &Result : Results) {
     WithSolutionSpecificVarTypesRAII VarTypes(Result.SolutionSpecificVarTypes);
 
@@ -117,7 +121,11 @@ void ExprTypeCheckCompletionCallback::deliverResults(
     if (Result.UnresolvedMemberBaseType) {
       Lookup.getUnresolvedMemberCompletions(Result.UnresolvedMemberBaseType);
     }
+
+    UnifiedTypeContext.merge(*Lookup.getExpectedTypeContext());
+    UnifiedCanHandleAsync |= Result.IsInAsyncContext;
   }
 
-  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+  collectCompletionResults(CompletionCtx, Lookup, DC, UnifiedTypeContext,
+                           UnifiedCanHandleAsync);
 }

lib/IDE/KeyPathCompletion.cpp

@@ -73,10 +73,9 @@ void KeyPathTypeCheckCompletionCallback::sawSolutionImpl(
   Results.push_back({BaseType, /*OnRoot=*/(ComponentIndex == 0)});
 }
 
-void KeyPathTypeCheckCompletionCallback::deliverResults(
+void KeyPathTypeCheckCompletionCallback::collectResults(
     DeclContext *DC, SourceLoc DotLoc,
-    ide::CodeCompletionContext &CompletionCtx,
+    ide::CodeCompletionContext &CompletionCtx) {
-    CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
@@ -91,5 +90,6 @@ void KeyPathTypeCheckCompletionCallback::deliverResults(
     Lookup.getValueExprCompletions(Result.BaseType);
   }
 
-  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+  collectCompletionResults(CompletionCtx, Lookup, DC, ExpectedTypeContext(),
+                           /*CanCurrDeclContextHandleAsync=*/false);
 }

lib/IDE/PostfixCompletion.cpp

@@ -396,10 +396,9 @@ static void addOperatorResults(Type LHSType, ArrayRef<OperatorDecl *> Operators,
   }
 }
 
-void PostfixCompletionCallback::deliverResults(
+void PostfixCompletionCallback::collectResults(
     SourceLoc DotLoc, bool IsInSelector, bool IncludeOperators,
-    bool HasLeadingSpace, CodeCompletionContext &CompletionCtx,
+    bool HasLeadingSpace, CodeCompletionContext &CompletionCtx) {
-    CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
@@ -429,6 +428,11 @@ void PostfixCompletionCallback::deliverResults(
     Lookup.collectOperators(Operators);
   }
 
+  // The type context that is being used for global results.
+  ExpectedTypeContext UnifiedTypeContext;
+  UnifiedTypeContext.setPreferNonVoid(true);
+  bool UnifiedCanHandleAsync = false;
+
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
   for (auto &Result : Results) {
     Lookup.setCanCurrDeclContextHandleAsync(Result.IsInAsyncContext);
@@ -446,7 +450,11 @@ void PostfixCompletionCallback::deliverResults(
     if (IncludeOperators && !Result.BaseIsStaticMetaType) {
       addOperatorResults(Result.BaseTy, Operators, DC, Lookup);
     }
+
+    UnifiedTypeContext.merge(*Lookup.getExpectedTypeContext());
+    UnifiedCanHandleAsync |= Result.IsInAsyncContext;
   }
 
-  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+  collectCompletionResults(CompletionCtx, Lookup, DC, UnifiedTypeContext,
+                           UnifiedCanHandleAsync);
 }

lib/IDE/REPLCodeCompletion.cpp

@@ -174,7 +174,7 @@ static void toDisplayString(CodeCompletionResult *Result,
 }
 
 namespace swift {
-class REPLCodeCompletionConsumer : public SimpleCachingCodeCompletionConsumer {
+class REPLCodeCompletionConsumer : public CodeCompletionConsumer {
   REPLCompletions &Completions;
 
 public:

lib/IDE/UnresolvedMemberCompletion.cpp

@@ -88,10 +88,9 @@ void UnresolvedMemberTypeCheckCompletionCallback::sawSolutionImpl(
   }
 }
 
-void UnresolvedMemberTypeCheckCompletionCallback::deliverResults(
+void UnresolvedMemberTypeCheckCompletionCallback::collectResults(
     DeclContext *DC, SourceLoc DotLoc,
-    ide::CodeCompletionContext &CompletionCtx,
+    ide::CodeCompletionContext &CompletionCtx) {
-    CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
   CompletionLookup Lookup(CompletionCtx.getResultSink(), Ctx, DC,
                           &CompletionCtx);
@@ -125,6 +124,11 @@ void UnresolvedMemberTypeCheckCompletionCallback::deliverResults(
     Lookup.getUnresolvedMemberCompletions(Result.ExpectedTy);
   }
 
+  // The type context that is being used for global results.
+  ExpectedTypeContext UnifiedTypeContext;
+  UnifiedTypeContext.setPreferNonVoid(true);
+  bool UnifiedCanHandleAsync = false;
+
   // Offer completions when interpreting the pattern match as an
   // EnumElementPattern.
   for (auto &Result : EnumPatternTypes) {
@@ -141,7 +145,11 @@ void UnresolvedMemberTypeCheckCompletionCallback::deliverResults(
     }
 
     Lookup.getEnumElementPatternCompletions(Ty);
+
+    UnifiedTypeContext.merge(*Lookup.getExpectedTypeContext());
+    UnifiedCanHandleAsync |= Result.IsInAsyncContext;
   }
 
-  deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
+  collectCompletionResults(CompletionCtx, Lookup, DC, UnifiedTypeContext,
+                           UnifiedCanHandleAsync);
 }

lib/IDETool/IDEInspectionInstance.cpp

@@ -594,8 +594,7 @@ void swift::ide::IDEInspectionInstance::codeComplete(
     llvm::function_ref<void(CancellableResult<CodeCompleteResult>)> Callback) {
   using ResultType = CancellableResult<CodeCompleteResult>;
 
-  struct ConsumerToCallbackAdapter
+  struct ConsumerToCallbackAdapter : public CodeCompletionConsumer {
-      : public SimpleCachingCodeCompletionConsumer {
     SwiftCompletionInfo SwiftContext;
     ImportDepth ImportDep;
     std::shared_ptr<std::atomic<bool>> CancellationFlag;