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;
@@ -1276,8 +1279,22 @@ class SILFunction
   GenericEnvironment *getGenericEnvironment() const {
     return GenericEnv;
   }
-  void setGenericEnvironment(GenericEnvironment *env) {
+
+  /// 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,
+                             ArrayRef<GenericEnvironment *> capturedEnvs,
+                             SubstitutionMap forwardingSubs) {
     GenericEnv = env;
+    CapturedEnvs = capturedEnvs;
+    ForwardingSubMap = forwardingSubs;
   }
 
   /// Retrieve the generic signature from the generic environment of this
@@ -1306,9 +1323,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

@@ -1076,6 +1076,13 @@ class TypeConverter {
   /// Get the generic environment for a constant.
   GenericEnvironment *getConstantGenericEnvironment(SILDeclRef constant);
 
+  /// Get the generic environment for SILGen to use. The substitution map
+  /// 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::tuple<GenericEnvironment *, ArrayRef<GenericEnvironment *>, SubstitutionMap>
+  getForwardingSubstitutionsForLowering(SILDeclRef constant);
+
   /// Returns the SIL type of a constant reference.
   SILType getConstantType(TypeExpansionContext context, SILDeclRef constant) {
     return getConstantInfo(context, constant).getSILType();

lib/SIL/IR/SILFunction.cpp

@@ -244,7 +244,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();
@@ -280,6 +279,7 @@ void SILFunction::init(
   assert(!Transparent || !IsDynamicReplaceable);
   validateSubclassScope(classSubclassScope, isThunk, nullptr);
   setDebugScope(DebugScope);
+  setGenericEnvironment(genericEnv);
 }
 
 SILFunction::~SILFunction() {
@@ -993,14 +993,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

@@ -4094,6 +4094,26 @@ TypeConverter::getConstantGenericEnvironment(SILDeclRef c) {
       .genericSig.getGenericEnvironment();
 }
 
+std::tuple<GenericEnvironment *, ArrayRef<GenericEnvironment *>, SubstitutionMap>
+TypeConverter::getForwardingSubstitutionsForLowering(SILDeclRef constant) {
+  auto sig = getGenericSignatureWithCapturedEnvironments(constant);
+
+  auto *genericEnv = sig.baseGenericSig.getGenericEnvironment();
+  SubstitutionMap forwardingSubs;
+
+  if (sig.baseGenericSig)
+    forwardingSubs = genericEnv->getForwardingSubstitutionMap();
+
+  if (!sig.capturedEnvs.empty()) {
+    assert(sig.genericSig && !sig.genericSig->isEqual(sig.baseGenericSig));
+
+    forwardingSubs = buildSubstitutionMapWithCapturedEnvironments(
+        forwardingSubs, sig.genericSig, sig.capturedEnvs);
+  }
+
+  return std::make_tuple(genericEnv, sig.capturedEnvs, forwardingSubs);
+}
+
 SILType TypeConverter::getSubstitutedStorageType(TypeExpansionContext context,
                                                  AbstractStorageDecl *value,
                                                  Type lvalueType) {

lib/SILGen/SILGen.cpp

@@ -1223,8 +1223,11 @@ void SILGenModule::preEmitFunction(SILDeclRef constant, SILFunction *F,
                                    SILLocation Loc) {
   assert(F->empty() && "already emitted function?!");
 
-  if (F->getLoweredFunctionType()->isPolymorphic())
-    F->setGenericEnvironment(Types.getConstantGenericEnvironment(constant));
+  if (F->getLoweredFunctionType()->isPolymorphic()) {
+    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
   // the function.

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)