Pack expansion closures, part 5

include/swift/AST/ASTMangler.h

@@ -656,7 +656,8 @@ class ASTMangler : public Mangler {
   void appendClosureEntity(const AbstractClosureExpr *closure);
 
   void appendClosureComponents(Type Ty, unsigned discriminator, bool isImplicit,
-                               const DeclContext *parentContext);
+                               const DeclContext *parentContext,
+                               ArrayRef<GenericEnvironment *> capturedEnvs);
 
   void appendDefaultArgumentEntity(const DeclContext *ctx, unsigned index);
 

include/swift/AST/LocalArchetypeRequirementCollector.h

@@ -53,6 +53,17 @@ struct MapLocalArchetypesOutOfContext {
   Type operator()(SubstitutableType *type) const;
 };
 
+struct MapIntoLocalArchetypeContext {
+  GenericEnvironment *baseGenericEnv;
+  ArrayRef<GenericEnvironment *> capturedEnvs;
+
+  MapIntoLocalArchetypeContext(GenericEnvironment *baseGenericEnv,
+                               ArrayRef<GenericEnvironment *> capturedEnvs)
+      : baseGenericEnv(baseGenericEnv), capturedEnvs(capturedEnvs) {}
+
+  Type operator()(SubstitutableType *type) const;
+};
+
 GenericSignature buildGenericSignatureWithCapturedEnvironments(
     ASTContext &ctx,
     GenericSignature sig,

include/swift/SIL/SILFunction.h

@@ -237,6 +237,9 @@ class SILFunction
   /// The context archetypes of the function.
   GenericEnvironment *GenericEnv = nullptr;
 
+  /// Captured local generic environments.
+  ArrayRef<GenericEnvironment *> CapturedEnvs;
+
   /// The information about specialization.
   /// Only set if this function is a specialization of another function.
   const GenericSpecializationInformation *SpecializationInfo = nullptr;
@@ -1292,9 +1295,21 @@ class SILFunction
   GenericEnvironment *getGenericEnvironment() const {
     return GenericEnv;
   }
+
+  /// Return any captured local generic environments, currently used for pack
+  /// element environments only. After SILGen, these are rewritten into
+  /// primary archetypes.
+  ArrayRef<GenericEnvironment *> getCapturedEnvironments() const {
+    return CapturedEnvs;
+  }
+
+  void setGenericEnvironment(GenericEnvironment *env);
+
   void setGenericEnvironment(GenericEnvironment *env,
-                             SubstitutionMap forwardingSubs=SubstitutionMap()) {
+                             ArrayRef<GenericEnvironment *> capturedEnvs,
+                             SubstitutionMap forwardingSubs) {
     GenericEnv = env;
+    CapturedEnvs = capturedEnvs;
     ForwardingSubMap = forwardingSubs;
   }
 
@@ -1324,9 +1339,30 @@ class SILFunction
   /// responsibility of the caller.
   void eraseAllBlocks();
 
-  /// Return the identity substitutions necessary to forward this call if it is
-  /// generic.
-  SubstitutionMap getForwardingSubstitutionMap();
+  /// A substitution map that sends the generic parameters of the invocation
+  /// generic signature to some combination of primar and local archetypes.
+  ///
+  /// CAUTION: If this is a SILFunction that captures pack element environments,
+  /// then at SILGen time, this is not actually the forwarding substitution map
+  /// of the SILFunction's generic environment. This is because:
+  ///
+  /// 1) The SILFunction's generic signature includes extra generic parameters,
+  ///    to model captured pack elements;
+  /// 2) The SILFunction's generic environment is the AST generic environment,
+  ///    so it's based on the original generic signature;
+  /// 3) SILGen uses this AST generic environment together with local archetypes
+  ///    for lowering SIL instructions.
+  ///
+  /// Therefore, the SILFunction's forwarding substitution map takes the extended
+  /// generic signature (1). It maps the original generic parameters to the
+  /// archetypes of (2), and the extended generic parameters to the local archetypes
+  /// of (3).
+  ///
+  /// After SILGen, all archetypes are re-instantiated inside the SIL function,
+  /// and the forwarding substitution map and generic environment then align.
+  SubstitutionMap getForwardingSubstitutionMap() const {
+    return ForwardingSubMap;
+  }
 
   /// Returns true if this SILFunction must be a defer statement.
   ///

include/swift/SIL/TypeLowering.h

@@ -1080,7 +1080,7 @@ class TypeConverter {
   /// sends the generic parameters of the function's interface type into
   /// archetypes, which will either be primary archetypes from this
   /// environment, or local archetypes captured by this function.
-  std::pair<GenericEnvironment *, SubstitutionMap>
+  std::tuple<GenericEnvironment *, ArrayRef<GenericEnvironment *>, SubstitutionMap>
   getForwardingSubstitutionsForLowering(SILDeclRef constant);
 
   /// Returns the SIL type of a constant reference.

lib/AST/ASTMangler.cpp

@@ -25,6 +25,7 @@
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/LazyResolver.h"
+#include "swift/AST/LocalArchetypeRequirementCollector.h"
 #include "swift/AST/MacroDiscriminatorContext.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/Ownership.h"
@@ -3733,29 +3734,48 @@ void ASTMangler::appendAssociatedTypeName(DependentMemberType *dmt,
 
 void ASTMangler::appendClosureEntity(
                               const SerializedAbstractClosureExpr *closure) {
+  assert(!closure->getType()->hasLocalArchetype() &&
+         "Not enough information here to handle this case");
+
   appendClosureComponents(closure->getType(), closure->getDiscriminator(),
-                          closure->isImplicit(), closure->getParent());
+                          closure->isImplicit(), closure->getParent(),
+                          ArrayRef<GenericEnvironment *>());
 }
 
 void ASTMangler::appendClosureEntity(const AbstractClosureExpr *closure) {
-  appendClosureComponents(closure->getType(), closure->getDiscriminator(),
-                          isa<AutoClosureExpr>(closure), closure->getParent());
+  ArrayRef<GenericEnvironment *> capturedEnvs;
+
+  auto type = closure->getType();
+
+  // FIXME: CodeCompletionResultBuilder calls printValueDeclUSR() but the
+  // closure hasn't been type checked yet.
+  if (!type)
+    type = ErrorType::get(closure->getASTContext());
+
+  if (type->hasLocalArchetype())
+    capturedEnvs = closure->getCaptureInfo().getGenericEnvironments();
+
+  appendClosureComponents(type, closure->getDiscriminator(),
+                          isa<AutoClosureExpr>(closure), closure->getParent(),
+                          capturedEnvs);
 }
 
 void ASTMangler::appendClosureComponents(Type Ty, unsigned discriminator,
                                          bool isImplicit,
-                                         const DeclContext *parentContext) {
+                                         const DeclContext *parentContext,
+                                         ArrayRef<GenericEnvironment *> capturedEnvs) {
   assert(discriminator != AbstractClosureExpr::InvalidDiscriminator
          && "closure must be marked correctly with discriminator");
 
   BaseEntitySignature base(parentContext->getInnermostDeclarationDeclContext());
   appendContext(parentContext, base, StringRef());
 
-  if (!Ty)
-    Ty = ErrorType::get(parentContext->getASTContext());
-
   auto Sig = parentContext->getGenericSignatureOfContext();
-  Ty = Ty->mapTypeOutOfContext();
+
+  Ty = Ty.subst(MapLocalArchetypesOutOfContext(Sig, capturedEnvs),
+                MakeAbstractConformanceForGenericType(),
+                SubstFlags::PreservePackExpansionLevel);
+
   appendType(Ty->getCanonicalType(), Sig);
   appendOperator(isImplicit ? "fu" : "fU", Index(discriminator));
 }

lib/AST/LocalArchetypeRequirementCollector.cpp

@@ -194,6 +194,28 @@ Type MapLocalArchetypesOutOfContext::operator()(SubstitutableType *type) const {
   llvm_unreachable("Fell off the end");
 }
 
+static Type mapIntoLocalContext(GenericTypeParamType *param, unsigned baseDepth,
+                                ArrayRef<GenericEnvironment *> capturedEnvs) {
+  assert(!param->isParameterPack());
+  unsigned envIndex = param->getDepth() - baseDepth;
+  assert(envIndex < capturedEnvs.size());
+  auto *capturedEnv = capturedEnvs[envIndex];
+  auto localInterfaceType = capturedEnv->getGenericSignature()
+      .getInnermostGenericParams()[param->getIndex()];
+  assert(localInterfaceType->getIndex() == param->getIndex());
+  return capturedEnvs[envIndex]->mapTypeIntoContext(localInterfaceType);
+}
+
+Type MapIntoLocalArchetypeContext::operator()(SubstitutableType *type) const {
+  unsigned baseDepth = baseGenericEnv->getGenericSignature().getNextDepth();
+
+  auto param = cast<GenericTypeParamType>(type);
+  if (param->getDepth() >= baseDepth)
+    return mapIntoLocalContext(param, baseDepth, capturedEnvs);
+
+  return baseGenericEnv->mapTypeIntoContext(param);
+}
+
 /// Given a substitution map for a call to a local function or closure, extend
 /// it to include all captured element archetypes; they become primary archetypes
 /// inside the body of the function.
@@ -215,15 +237,8 @@ swift::buildSubstitutionMapWithCapturedEnvironments(
     genericSigWithCaptures,
     [&](SubstitutableType *type) -> Type {
       auto param = cast<GenericTypeParamType>(type);
-      if (param->getDepth() >= baseDepth) {
-        assert(!param->isParameterPack());
-        unsigned envIndex = param->getDepth() - baseDepth;
-        assert(envIndex < capturedEnvs.size());
-        auto *capturedEnv = capturedEnvs[envIndex];
-        auto localInterfaceType = capturedEnv->getGenericSignature()
-            .getInnermostGenericParams()[param->getIndex()];
-        return capturedEnvs[envIndex]->mapTypeIntoContext(localInterfaceType);
-      }
+      if (param->getDepth() >= baseDepth)
+        return mapIntoLocalContext(param, baseDepth, capturedEnvs);
       return Type(type).subst(baseSubMap);
     },
     [&](CanType origType, Type substType,

lib/SIL/IR/SILFunction.cpp

@@ -17,6 +17,7 @@
 #include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/IRGenOptions.h"
+#include "swift/AST/LocalArchetypeRequirementCollector.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/Stmt.h"
 #include "swift/Basic/OptimizationMode.h"
@@ -245,7 +246,6 @@ void SILFunction::init(
          "function type has open type parameters");
 
   this->LoweredType = LoweredType;
-  this->GenericEnv = genericEnv;
   this->SpecializationInfo = nullptr;
   this->EntryCount = entryCount;
   this->Availability = AvailabilityContext::alwaysAvailable();
@@ -282,6 +282,7 @@ void SILFunction::init(
   assert(!Transparent || !IsDynamicReplaceable);
   validateSubclassScope(classSubclassScope, isThunk, nullptr);
   setDebugScope(DebugScope);
+  setGenericEnvironment(genericEnv);
 }
 
 SILFunction::~SILFunction() {
@@ -492,24 +493,51 @@ bool SILFunction::shouldOptimize() const {
 }
 
 Type SILFunction::mapTypeIntoContext(Type type) const {
-  return GenericEnvironment::mapTypeIntoContext(
-      getGenericEnvironment(), type);
+  assert(!type->hasPrimaryArchetype());
+
+  if (GenericEnv) {
+    // The complication here is that we sometimes call this with an AST interface
+    // type, which might contain element archetypes, if it was the interface type
+    // of a closure or local variable.
+    if (type->hasElementArchetype())
+      return GenericEnv->mapTypeIntoContext(type);
+
+    // Otherwise, assume we have an interface type for the "combined" captured
+    // environment.
+    return type.subst(MapIntoLocalArchetypeContext(GenericEnv, CapturedEnvs),
+                      LookUpConformanceInModule(Module.getSwiftModule()),
+                      SubstFlags::AllowLoweredTypes |
+                      SubstFlags::PreservePackExpansionLevel);
+  }
+
+  assert(!type->hasTypeParameter());
+  return type;
 }
 
 SILType SILFunction::mapTypeIntoContext(SILType type) const {
-  if (auto *genericEnv = getGenericEnvironment())
-    return genericEnv->mapTypeIntoContext(getModule(), type);
+  assert(!type.hasPrimaryArchetype());
+
+  if (GenericEnv) {
+    auto genericSig = GenericEnv->getGenericSignature().getCanonicalSignature();
+    return type.subst(Module,
+                      MapIntoLocalArchetypeContext(GenericEnv, CapturedEnvs),
+                      LookUpConformanceInModule(Module.getSwiftModule()),
+                      genericSig,
+                      SubstFlags::PreservePackExpansionLevel);
+  }
+
+  assert(!type.hasTypeParameter());
   return type;
 }
 
 SILType GenericEnvironment::mapTypeIntoContext(SILModule &M,
                                                SILType type) const {
-  assert(!type.hasArchetype());
+  assert(!type.hasPrimaryArchetype());
 
   auto genericSig = getGenericSignature().getCanonicalSignature();
   return type.subst(M,
                     QueryInterfaceTypeSubstitutions(this),
-                    LookUpConformanceInSignature(genericSig.getPointer()),
+                    LookUpConformanceInModule(M.getSwiftModule()),
                     genericSig,
                     SubstFlags::PreservePackExpansionLevel);
 }
@@ -1033,14 +1061,10 @@ void SILFunction::eraseAllBlocks() {
   BlockList.clear();
 }
 
-SubstitutionMap SILFunction::getForwardingSubstitutionMap() {
-  if (ForwardingSubMap)
-    return ForwardingSubMap;
-
-  if (auto *env = getGenericEnvironment())
-    ForwardingSubMap = env->getForwardingSubstitutionMap();
-
-  return ForwardingSubMap;
+void SILFunction::setGenericEnvironment(GenericEnvironment *env) {
+  setGenericEnvironment(env, ArrayRef<GenericEnvironment *>(),
+                        env ? env->getForwardingSubstitutionMap()
+                            : SubstitutionMap());
 }
 
 bool SILFunction::shouldVerifyOwnership() const {

lib/SIL/IR/SILModule.cpp

@@ -688,6 +688,23 @@ SILValue SILModule::getRootLocalArchetypeDef(CanLocalArchetypeType archetype,
 }
 
 bool SILModule::hasUnresolvedLocalArchetypeDefinitions() {
+  // Garbage collect dead placeholders first.
+  llvm::DenseMap<LocalArchetypeKey, SILValue> newLocalArchetypeDefs;
+
+  for (auto pair : RootLocalArchetypeDefs) {
+    if (auto *placeholder = dyn_cast<PlaceholderValue>(pair.second)) {
+      if (placeholder->use_empty()) {
+        --numUnresolvedLocalArchetypes;
+        ::delete placeholder;
+        continue;
+      }
+    }
+
+    newLocalArchetypeDefs.insert(pair);
+  }
+
+  std::swap(newLocalArchetypeDefs, RootLocalArchetypeDefs);
+
   return numUnresolvedLocalArchetypes != 0;
 }
 
@@ -753,6 +770,8 @@ void SILModule::notifyAddedInstruction(SILInstruction *inst) {
       auto *placeholder = cast<PlaceholderValue>(val);
       placeholder->replaceAllUsesWith(dependency);
       ::delete placeholder;
+
+      assert(numUnresolvedLocalArchetypes > 0);
       numUnresolvedLocalArchetypes--;
     }
     val = dependency;

lib/SIL/IR/TypeLowering.cpp

@@ -4116,17 +4116,15 @@ TypeConverter::getConstantGenericEnvironment(SILDeclRef c) {
       .genericSig.getGenericEnvironment();
 }
 
-std::pair<GenericEnvironment *, SubstitutionMap>
+std::tuple<GenericEnvironment *, ArrayRef<GenericEnvironment *>, SubstitutionMap>
 TypeConverter::getForwardingSubstitutionsForLowering(SILDeclRef constant) {
   auto sig = getGenericSignatureWithCapturedEnvironments(constant);
 
-  GenericEnvironment *genericEnv = nullptr;
+  auto *genericEnv = sig.baseGenericSig.getGenericEnvironment();
   SubstitutionMap forwardingSubs;
 
-  if (sig.baseGenericSig) {
-    genericEnv = sig.baseGenericSig.getGenericEnvironment();
+  if (sig.baseGenericSig)
     forwardingSubs = genericEnv->getForwardingSubstitutionMap();
-  }
 
   if (!sig.capturedEnvs.empty()) {
     assert(sig.genericSig && !sig.genericSig->isEqual(sig.baseGenericSig));
@@ -4135,7 +4133,7 @@ TypeConverter::getForwardingSubstitutionsForLowering(SILDeclRef constant) {
         forwardingSubs, sig.genericSig, sig.capturedEnvs);
   }
 
-  return std::make_pair(genericEnv, forwardingSubs);
+  return std::make_tuple(genericEnv, sig.capturedEnvs, forwardingSubs);
 }
 
 SILType TypeConverter::getSubstitutedStorageType(TypeExpansionContext context,

lib/SILGen/SILGen.cpp

@@ -1224,8 +1224,9 @@ void SILGenModule::preEmitFunction(SILDeclRef constant, SILFunction *F,
   assert(F->empty() && "already emitted function?!");
 
   if (F->getLoweredFunctionType()->isPolymorphic()) {
-    auto pair = Types.getForwardingSubstitutionsForLowering(constant);
-    F->setGenericEnvironment(pair.first, pair.second);
+    auto [genericEnv, capturedEnvs, forwardingSubs]
+        = Types.getForwardingSubstitutionsForLowering(constant);
+    F->setGenericEnvironment(genericEnv, capturedEnvs, forwardingSubs);
   }
 
   // If we have global actor isolation for our constant, put the isolation onto