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

Witness thunks where the conforming type is a class and the witness is in a
protocol extension have an extra generic parameter constrained to the class
type that is passed as the 'Self' parameter for the protocol extension
method.

This means the substitution map for the devirtualized call must be
assembled from three sources:

- The 'Self' substitution
- The generic parameters of the concrete conforming type, if any
- The generic parameters of the protocol requirement, if any

This was previously done by making two calls to combineSubstitutionMaps(),
the first call combined the first two maps and the second call combined the
result of the first call with the third map.

Unfortunately, both calls were made with the generic signature of the
witness thunk, and the result of combining the first two substitution maps
does not provide sufficient replacements for all generic parameters and
conformance requirements in the witness thunk's signature.

This was apparently fine with the GenericSignatureBuilder, but the
Requirement Machine flags the missing generic parameters in assert builds.

Fixes #59193.

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) -> ()