SIL: Store captured environments in SILFunction

Author: slavapestov

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/SIL/IR/SILFunction.cpp

@@ -245,7 +245,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 +281,7 @@ void SILFunction::init(
   assert(!Transparent || !IsDynamicReplaceable);
   validateSubclassScope(classSubclassScope, isThunk, nullptr);
   setDebugScope(DebugScope);
+  setGenericEnvironment(genericEnv);
 }
 
 SILFunction::~SILFunction() {
@@ -1033,14 +1033,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/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

lib/Serialization/SerializeSIL.cpp

@@ -460,6 +460,9 @@ void SILSerializer::writeSILFunction(const SILFunction &F, bool DeclOnly) {
     Linkage = addExternalToLinkage(Linkage);
   }
 
+  assert(F.getCapturedEnvironments().empty() &&
+         "Captured local environments should not survive past SILGen");
+
   // If we have a body, we might have a generic environment.
   GenericSignatureID genericSigID = 0;
   if (!NoBody)