Merge pull request #59434 from slavapestov/devirtualize-class-witness-method-fix

SILOptimizer: Fix invariant violation in getWitnessMethodSubstitutions() with class witness methods

Author: slavapestov

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -900,6 +900,8 @@ getWitnessMethodSubstitutions(
     bool isSelfAbstract,
     ClassDecl *classWitness) {
 
+  auto &ctx = mod->getASTContext();
+
   if (witnessThunkSig.isNull())
     return SubstitutionMap();
 
@@ -909,46 +911,100 @@ getWitnessMethodSubstitutions(
   assert(!conformanceRef.isAbstract());
   auto conformance = conformanceRef.getConcrete();
 
+  auto selfType = conformance->getProtocol()->getSelfInterfaceType();
+
   // If `Self` maps to a bound generic type, this gives us the
   // substitutions for the concrete type's generic parameters.
   auto baseSubMap = conformance->getSubstitutions(mod);
 
   unsigned baseDepth = 0;
   auto *rootConformance = conformance->getRootConformance();
-  if (auto witnessSig = rootConformance->getGenericSignature())
-    baseDepth = witnessSig.getGenericParams().back()->getDepth() + 1;
-
-  // If the witness has a class-constrained 'Self' generic parameter,
-  // we have to build a new substitution map that shifts all generic
-  // parameters down by one.
-  if (classWitness != nullptr) {
-    auto *proto = conformance->getProtocol();
-    auto selfType = proto->getSelfInterfaceType();
-
-    auto selfSubMap = SubstitutionMap::getProtocolSubstitutions(
-        proto, selfType.subst(origSubMap), conformanceRef);
-    if (baseSubMap.empty()) {
-      assert(baseDepth == 0);
-      baseSubMap = selfSubMap;
-    } else {
-      baseSubMap = SubstitutionMap::combineSubstitutionMaps(
-          selfSubMap,
-          baseSubMap,
-          CombineSubstitutionMaps::AtDepth,
-          /*firstDepth=*/1,
-          /*secondDepth=*/0,
-          witnessThunkSig);
-    }
-    baseDepth += 1;
-  }
+  if (auto conformingTypeSig = rootConformance->getGenericSignature())
+    baseDepth = conformingTypeSig.getGenericParams().back()->getDepth() + 1;
 
-  return SubstitutionMap::combineSubstitutionMaps(
-      baseSubMap,
-      origSubMap,
-      CombineSubstitutionMaps::AtDepth,
-      /*firstDepth=*/baseDepth,
-      /*secondDepth=*/1,
-      witnessThunkSig);
+  // witnessThunkSig begins with the optional class 'Self', followed by the
+  // generic parameters of the concrete conforming type, followed by the
+  // generic parameters of the protocol requirement, if any.
+  //
+  // - The 'Self' parameter is replaced with the conforming type.
+  // - The conforming type's generic parameters are replaced by the
+  //   conformance substitutions.
+  // - The protocol requirement's generic parameters are replaced from the
+  //   substitution map at the call site.
+  return SubstitutionMap::get(
+      witnessThunkSig,
+      [&](SubstitutableType *type) {
+        auto *paramType = type->castTo<GenericTypeParamType>();
+        unsigned depth = paramType->getDepth();
+
+        if (classWitness != nullptr) {
+          if (depth == 0) {
+            assert(paramType->getIndex() == 0);
+            return selfType.subst(origSubMap);
+          }
+
+          --depth;
+        }
+
+        if (depth < baseDepth) {
+          paramType = GenericTypeParamType::get(
+              paramType->isTypeSequence(),
+              depth, paramType->getIndex(), ctx);
+
+          return Type(paramType).subst(baseSubMap);
+        }
+
+        depth = depth - baseDepth + 1;
+
+        paramType = GenericTypeParamType::get(
+            paramType->isTypeSequence(),
+            depth, paramType->getIndex(), ctx);
+        return Type(paramType).subst(origSubMap);
+      },
+      [&](CanType type, Type substType, ProtocolDecl *proto) {
+        auto *paramType = type->getRootGenericParam();
+        unsigned depth = paramType->getDepth();
+
+        if (classWitness != nullptr) {
+          if (depth == 0) {
+            assert(type->isEqual(paramType));
+            assert(paramType->getIndex() == 0);
+            return conformanceRef;
+          }
+
+          --depth;
+        }
+
+        if (depth < baseDepth) {
+          type = CanType(type.transform([&](Type t) -> Type {
+            if (t->isEqual(paramType)) {
+              return GenericTypeParamType::get(
+                  paramType->isTypeSequence(),
+                  depth, paramType->getIndex(), ctx);
+            }
+
+            assert(!t->is<GenericTypeParamType>());
+            return t;
+          }));
+
+          return baseSubMap.lookupConformance(type, proto);
+        }
+
+        depth = depth - baseDepth + 1;
+
+        type = CanType(type.transform([&](Type t) -> Type {
+          if (t->isEqual(paramType)) {
+            return GenericTypeParamType::get(
+                paramType->isTypeSequence(),
+                depth, paramType->getIndex(), ctx);
+          }
+
+          assert(!t->is<GenericTypeParamType>());
+          return t;
+        }));
+
+        return origSubMap.lookupConformance(type, proto);
+      });
 }
 
 SubstitutionMap

test/SILOptimizer/devirt_class_witness_method_generic.swift

@@ -0,0 +1,23 @@
+// RUN: %target-swift-frontend -emit-sil -O %s | %FileCheck %s
+
+public protocol P {}
+
+public protocol Q {
+  func foo<T: P>(t: T)
+  func bar<T: P>(t: T)
+}
+extension Q {
+  @inline(__always)
+  public func foo<T: P>(t: T) {
+    bar(t: t)
+  }
+
+  @_optimize(none)
+  public func bar<T: P>(t: T) {}
+}
+
+public class C<T>: Q {}
+
+// CHECK-LABEL: sil shared [transparent] [thunk] @$s35devirt_class_witness_method_generic1CCyqd__GAA1QA2aEP3bar1tyqd___tAA1PRd__lFTW : $@convention(witness_method: Q) <τ_0_0><τ_1_0 where τ_0_0 : C<τ_1_0>><τ_2_0 where τ_2_0 : P> (@in_guaranteed τ_2_0, @in_guaranteed τ_0_0) -> () {
+// CHECK: [[FN:%.*]] = function_ref @$s35devirt_class_witness_method_generic1QPAAE3bar1tyqd___tAA1PRd__lF : $@convention(method) <τ_0_0 where τ_0_0 : Q><τ_1_0 where τ_1_0 : P> (@in_guaranteed τ_1_0, @in_guaranteed τ_0_0) -> ()
+// CHECK: apply [[FN]]<τ_0_0, τ_2_0>(%0, %1) : $@convention(method) <τ_0_0 where τ_0_0 : Q><τ_1_0 where τ_1_0 : P> (@in_guaranteed τ_1_0, @in_guaranteed τ_0_0) -> ()