SIL: Sink local archetype substitution into remapSubstitutionMap()

Author: slavapestov

include/swift/SIL/SILCloner.h

@@ -29,6 +29,35 @@
 
 namespace swift {
 
+struct SubstitutionMapWithLocalArchetypes {
+  std::optional<SubstitutionMap> SubsMap;
+  TypeSubstitutionMap LocalArchetypeSubs;
+
+  SubstitutionMapWithLocalArchetypes() {}
+  SubstitutionMapWithLocalArchetypes(SubstitutionMap subs) : SubsMap(subs) {}
+
+  Type operator()(SubstitutableType *type) {
+    if (isa<LocalArchetypeType>(type))
+      return QueryTypeSubstitutionMap{LocalArchetypeSubs}(type);
+
+    if (SubsMap)
+      return Type(type).subst(*SubsMap);
+
+    return Type(type);
+  }
+
+  ProtocolConformanceRef operator()(CanType origType,
+                                    Type substType,
+                                    ProtocolDecl *proto) {
+    if (isa<LocalArchetypeType>(origType))
+      return proto->getParentModule()->lookupConformance(substType, proto);
+    if (SubsMap)
+      return SubsMap->lookupConformance(origType, proto);
+
+    return ProtocolConformanceRef(proto);
+  }
+};
+
 /// SILCloner - Abstract SIL visitor which knows how to clone instructions and
 /// whose behavior can be customized by subclasses via the CRTP. This is meant
 /// to be subclassed to implement inlining, function specialization, and other
@@ -49,7 +78,8 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
 
   SILBuilder Builder;
   DominanceInfo *DomTree = nullptr;
-  TypeSubstitutionMap LocalArchetypeSubs;
+  SubstitutionMapWithLocalArchetypes Functor;
+  TypeSubstitutionMap &LocalArchetypeSubs;
 
   // The old-to-new value map.
   llvm::DenseMap<SILValue, SILValue> ValueMap;
@@ -76,9 +106,10 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   using SILInstructionVisitor<ImplClass>::asImpl;
 
   explicit SILCloner(SILFunction &F, DominanceInfo *DT = nullptr)
-      : Builder(F), DomTree(DT) {}
+      : Builder(F), DomTree(DT), LocalArchetypeSubs(Functor.LocalArchetypeSubs) {}
 
-  explicit SILCloner(SILGlobalVariable *GlobVar) : Builder(GlobVar) {}
+  explicit SILCloner(SILGlobalVariable *GlobVar)
+      : Builder(GlobVar), LocalArchetypeSubs(Functor.LocalArchetypeSubs) {}
 
   void clearClonerState() {
     ValueMap.clear();
@@ -189,12 +220,6 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   }
 
   SubstitutionMap getOpSubstitutionMap(SubstitutionMap Subs) {
-    // If we have local archetypes to substitute, do so now.
-    if (Subs.hasLocalArchetypes() && !LocalArchetypeSubs.empty()) {
-      Subs = Subs.subst(QueryTypeSubstitutionMapOrIdentity{LocalArchetypeSubs},
-                        MakeAbstractConformanceForGenericType());
-    }
-
     return asImpl().remapSubstitutionMap(Subs)
                    .getCanonical(/*canonicalizeSignature*/false);
   }
@@ -441,7 +466,13 @@ class SILCloner : protected SILInstructionVisitor<ImplClass> {
   SILBasicBlock *remapBasicBlock(SILBasicBlock *BB);
   void postProcess(SILInstruction *Orig, SILInstruction *Cloned);
 
-  SubstitutionMap remapSubstitutionMap(SubstitutionMap Subs) { return Subs; }
+  SubstitutionMap remapSubstitutionMap(SubstitutionMap Subs) {
+    // If we have local archetypes to substitute, do so now.
+    if (Subs.hasLocalArchetypes())
+      Subs = Subs.subst(Functor, Functor);
+
+    return Subs;
+  }
 
   /// This is called by either of the top-level visitors, cloneReachableBlocks
   /// or cloneSILFunction, after all other visitors are have been called.

include/swift/SIL/TypeSubstCloner.h

@@ -148,6 +148,8 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
   using SILClonerWithScopes<ImplClass>::getOpBasicBlock;
   using SILClonerWithScopes<ImplClass>::recordClonedInstruction;
   using SILClonerWithScopes<ImplClass>::recordFoldedValue;
+  using SILClonerWithScopes<ImplClass>::LocalArchetypeSubs;
+  using SILClonerWithScopes<ImplClass>::Functor;
 
   TypeSubstCloner(SILFunction &To,
                   SILFunction &From,
@@ -159,6 +161,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
       SubsMap(ApplySubs),
       Original(From),
       Inlining(Inlining) {
+    Functor.SubsMap = ApplySubs;
   }
 
 protected:
@@ -205,9 +208,10 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
   }
 
   SubstitutionMap remapSubstitutionMap(SubstitutionMap Subs) {
+    Subs = Subs.subst(Functor, Functor);
+
     auto context = getBuilder().getTypeExpansionContext();
 
-    Subs = Subs.subst(SubsMap);
     if (!Subs.hasOpaqueArchetypes() ||
         !context.shouldLookThroughOpaqueTypeArchetypes())
       return Subs;

lib/SILOptimizer/IPO/CrossModuleOptimization.cpp

@@ -143,6 +143,9 @@ class InstructionVisitor : public SILCloner<InstructionVisitor> {
   }
 
   SubstitutionMap remapSubstitutionMap(SubstitutionMap Subs) {
+    if (Subs.hasLocalArchetypes())
+      Subs = Subs.subst(Functor, Functor);
+
     CMS.makeSubstUsableFromInline(Subs);
     return Subs;
   }

test/SILOptimizer/opened_existential_concrete_member_type.swift

@@ -0,0 +1,23 @@
+// RUN: %target-swift-frontend -emit-ir %s
+
+public protocol P1 {}
+
+public protocol P2<A> {
+  associatedtype A
+}
+
+public struct S: P1 {}
+
+public struct G<A>: P2 {}
+
+public func callee<T: P2>(_: T) where T.A: P1 {}
+
+@_transparent
+public func caller<A: P1>(_ p: any P2<A>) {
+  callee(p)
+}
+
+public func test(_ a: G<S>, _ b: G<S>) {
+  return caller(b)
+}
+