SILCloner: Redo local archetype remapping

Author: slavapestov

include/swift/SIL/SILCloner.h

@@ -19,6 +19,7 @@
 
 #include "swift/AST/ConformanceLookup.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/LocalArchetypeRequirementCollector.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/DebugUtils.h"
@@ -32,14 +33,47 @@ namespace swift {
 
 struct SubstitutionMapWithLocalArchetypes {
   std::optional<SubstitutionMap> SubsMap;
-  TypeSubstitutionMap LocalArchetypeSubs;
+  llvm::DenseMap<GenericEnvironment *, GenericEnvironment *> LocalArchetypeSubs;
+  GenericSignature BaseGenericSig;
+  llvm::ArrayRef<GenericEnvironment *> CapturedEnvs;
+
+  bool hasLocalArchetypes() const {
+    return !LocalArchetypeSubs.empty() || !CapturedEnvs.empty();
+  }
 
   SubstitutionMapWithLocalArchetypes() {}
   SubstitutionMapWithLocalArchetypes(SubstitutionMap subs) : SubsMap(subs) {}
 
   Type operator()(SubstitutableType *type) {
-    if (isa<LocalArchetypeType>(type))
-      return QueryTypeSubstitutionMap{LocalArchetypeSubs}(type);
+    if (auto *local = dyn_cast<LocalArchetypeType>(type)) {
+      auto *origEnv = local->getGenericEnvironment();
+
+      // Special handling of captured environments, which don't appear in
+      // LocalArchetypeSubs. This only happens with the LocalArchetypeTransform
+      // in SILGenLocalArchetype.cpp.
+      auto found = LocalArchetypeSubs.find(origEnv);
+      if (found == LocalArchetypeSubs.end()) {
+        if (std::find(CapturedEnvs.begin(), CapturedEnvs.end(), origEnv)
+            == CapturedEnvs.end()) {
+          return Type();
+        }
+
+        // Map the local archetype to an interface type in the new generic
+        // signature.
+        MapLocalArchetypesOutOfContext mapOutOfContext(BaseGenericSig,
+                                                       CapturedEnvs);
+        auto interfaceTy = mapOutOfContext(local);
+
+        // Map this interface type into the new generic environment to get
+        // a primary archetype.
+        return interfaceTy.subst(*SubsMap);
+      }
+
+      auto *newEnv = found->second;
+
+      auto interfaceTy = local->getInterfaceType();
+      return newEnv->mapTypeIntoContext(interfaceTy);
+    }
 
     if (SubsMap)
       return Type(type).subst(*SubsMap);
@@ -210,10 +244,12 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   }
 
   /// Register a re-mapping for local archetypes such as opened existentials.
-  void registerLocalArchetypeRemapping(ArchetypeType *From,
-                                       ArchetypeType *To) {
-    auto result = Functor.LocalArchetypeSubs.insert(
-        std::make_pair(CanArchetypeType(From), CanType(To)));
+  void registerLocalArchetypeRemapping(GenericEnvironment *From,
+                                       GenericEnvironment *To) {
+    ASSERT(From->getGenericSignature()->getMaxDepth()
+           == To->getGenericSignature()->getMaxDepth());
+
+    auto result = Functor.LocalArchetypeSubs.insert(std::make_pair(From, To));
     assert(result.second);
     (void)result;
   }
@@ -344,13 +380,17 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   }
 
   void remapRootOpenedType(CanOpenedArchetypeType archetypeTy) {
-    assert(archetypeTy->isRoot());
-
-    auto origExistentialTy = archetypeTy->getExistentialType()
+    auto *origEnv = archetypeTy->getGenericEnvironment();
+    auto subMap = origEnv->getOuterSubstitutions();
+    ASSERT(!subMap && "Transform the substitution map!");
+    auto origExistentialTy = origEnv->getOpenedExistentialType()
         ->getCanonicalType();
+
     auto substExistentialTy = getOpASTType(origExistentialTy);
-    auto replacementTy = OpenedArchetypeType::get(substExistentialTy);
-    registerLocalArchetypeRemapping(archetypeTy, replacementTy);
+    auto *newEnv = GenericEnvironment::forOpenedExistential(
+        substExistentialTy, subMap, UUID::fromTime());
+
+    registerLocalArchetypeRemapping(origEnv, newEnv);
   }
 
   /// SILCloner will take care of debug scope on the instruction
@@ -469,7 +509,7 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   SILType remapType(SILType Ty) {
     if (Functor.SubsMap || Ty.hasLocalArchetype()) {
       SubstOptions options(std::nullopt);
-      if (!Functor.LocalArchetypeSubs.empty())
+      if (Functor.hasLocalArchetypes())
         options |= SubstFlags::SubstituteLocalArchetypes;
 
       Ty = Ty.subst(Builder.getModule(), Functor, Functor,
@@ -494,7 +534,7 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   CanType remapASTType(CanType ty) {
     if (Functor.SubsMap || ty->hasLocalArchetype()) {
       SubstOptions options(std::nullopt);
-      if (!Functor.LocalArchetypeSubs.empty())
+      if (Functor.hasLocalArchetypes())
         options |= SubstFlags::SubstituteLocalArchetypes;
 
       ty = ty.subst(Functor, Functor, options)->getCanonicalType();
@@ -518,7 +558,7 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   ProtocolConformanceRef remapConformance(Type Ty, ProtocolConformanceRef C) {
     if (Functor.SubsMap || Ty->hasLocalArchetype()) {
       SubstOptions options(std::nullopt);
-      if (!Functor.LocalArchetypeSubs.empty())
+      if (Functor.hasLocalArchetypes())
         options |= SubstFlags::SubstituteLocalArchetypes;
 
       C = C.subst(Ty, Functor, Functor, options);
@@ -541,9 +581,9 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
 
   SubstitutionMap remapSubstitutionMap(SubstitutionMap Subs) {
     // If we have local archetypes to substitute, do so now.
-    if (Subs.getRecursiveProperties().hasLocalArchetype() || Functor.SubsMap) {
+    if (Functor.SubsMap || Subs.getRecursiveProperties().hasLocalArchetype()) {
       SubstOptions options(std::nullopt);
-      if (!Functor.LocalArchetypeSubs.empty())
+      if (Functor.hasLocalArchetypes())
         options |= SubstFlags::SubstituteLocalArchetypes;
 
       Subs = Subs.subst(Functor, Functor, options);
@@ -2909,20 +2949,7 @@ void SILCloner<ImplClass>::visitOpenPackElementInst(
                                          openedShapeClass,
                                          newContextSubs);
 
-  // Associate the old opened archetypes with the new ones.
-  SmallVector<ArchetypeType*, 4> oldOpenedArchetypes;
-  origEnv->forEachPackElementArchetype([&](ElementArchetypeType *oldType) {
-    oldOpenedArchetypes.push_back(oldType);
-  });
-  {
-    size_t nextOldIndex = 0;
-    newEnv->forEachPackElementArchetype([&](ElementArchetypeType *newType) {
-      ArchetypeType *oldType = oldOpenedArchetypes[nextOldIndex++];
-      registerLocalArchetypeRemapping(oldType, newType);
-    });
-    assert(nextOldIndex == oldOpenedArchetypes.size() &&
-           "different opened archetype count");
-  }
+  registerLocalArchetypeRemapping(origEnv, newEnv);
 
   recordClonedInstruction(
       Inst, getBuilder().createOpenPackElement(loc, newIndexValue, newEnv));

lib/SILGen/SILGenLocalArchetype.cpp

@@ -46,33 +46,10 @@ class LocalArchetypeTransform : public SILCloner<LocalArchetypeTransform> {
     // index and depth.
     Functor.SubsMap = env->getForwardingSubstitutionMap();
 
-    // Local archetypes map to generic parameters at higher depths.
-    MapLocalArchetypesOutOfContext mapOutOfContext(sig.baseGenericSig,
-                                                   sig.capturedEnvs);
-
-    // For each captured environment...
-    for (auto *capturedEnv : sig.capturedEnvs) {
-      // For each introduced generic parameter...
-      auto localParams = capturedEnv->getGenericSignature()
-          .getInnermostGenericParams();
-      for (auto *gp : localParams) {
-        // Get the local archetype from the captured environment.
-        auto origArchetypeTy = capturedEnv->mapTypeIntoContext(gp)
-            ->castTo<LocalArchetypeType>();
-
-        // Map the local archetype to an interface type in the new generic
-        // signature.
-        auto substInterfaceTy = mapOutOfContext(origArchetypeTy);
-
-        // Map this interface type into the new generic environment to get
-        // a primary archetype.
-        auto substArchetypeTy = env->mapTypeIntoContext(substInterfaceTy)
-            ->castTo<PrimaryArchetypeType>();
-
-        // Remember this correspondence.
-        registerLocalArchetypeRemapping(origArchetypeTy, substArchetypeTy);
-      }
-    }
+    // Captured local archetypes map to primary archetypes at higher
+    // depths.
+    Functor.BaseGenericSig = sig.baseGenericSig;
+    Functor.CapturedEnvs = sig.capturedEnvs;
   }
 
   void doIt() {

lib/SILOptimizer/Transforms/CSE.cpp

@@ -908,7 +908,8 @@ bool CSE::processOpenExistentialRef(OpenExistentialRefInst *Inst,
   // opened archetypes trivial.
   InstructionCloner Cloner(Inst->getFunction());
   Cloner.registerLocalArchetypeRemapping(
-      OldOpenedArchetype->castTo<ArchetypeType>(), NewOpenedArchetype);
+      OldOpenedArchetype->getGenericEnvironment(),
+      NewOpenedArchetype->getGenericEnvironment());
   auto &Builder = Cloner.getBuilder();
 
   // Now clone each candidate and replace the opened archetype