AST: Change RequirementEnvironment::getRequirementToWitnessThunkSubs() to use contextual types

In the provided test case, the generic signature of the
protocol requirement is

    <Self, T where Self == T.A, T: P2>

The conformance requirement `Self: P1` is derived from `T: P2`
and `Self.A: P1` in P1. So given a substitution map
`{ Self := X, T := T }` that replaces T with a concrete type X
and T with a type parameter T, there are two ways to recover a
substituted conformance for `Self: P1`:

- We can apply the substitution map to Self to get X, and look
  up conformance of X to P1, to get a concrete conformance.

- We can evaluate the conformance path `(T: P2)(Self.A: P1)`,
  to get an abstract conformance.

Both answers are correct, but SILGenModule::emitProtocolWitness()
was assuming it would always get a concrete conformance back.

This was the case until e3c8f42,
but then we started returning an abstract conformance. SILGen
would then mangle the protocol witness thunk in a way that was
not sufficiently unique, and as a result, we could miscompile
a program where two witness tables both hit this same scenario.

By using contextual types in the getRequirementToWitnessThunkSubs()
substitution map, we ensure that evaluating the conformance path
against the substitution map produces the same result as performing
the global lookup.

Also, to prevent this from happening again, add a check to SILGen
to guard against emitting two protocol witness thunks with the
same mangled name.

Unfortunately, this is done intentionally as part of some
backward deployment logic for coroutine accessors. There is a
hack to allow duplicate thunks with the same name in this case,
but this should be revisited some day.

Fixes rdar://problem/155624135..

Author: slavapestov

include/swift/AST/RequirementEnvironment.h

@@ -114,7 +114,8 @@ class RequirementEnvironment {
   }
 
   /// Retrieve the substitution map that maps the interface types of the
-  /// requirement to the interface types of the witness thunk signature.
+  /// requirement to the archetypes of the witness thunk signature's
+  /// generic environment.
   SubstitutionMap getRequirementToWitnessThunkSubs() const {
     return reqToWitnessThunkSigMap;
   }

lib/AST/RequirementEnvironment.cpp

@@ -19,6 +19,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DeclContext.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeCheckRequests.h"
@@ -156,6 +157,11 @@ RequirementEnvironment::RequirementEnvironment(
       reqSig.getRequirements().size() == 1) {
     witnessThunkSig = conformanceDC->getGenericSignatureOfContext()
         .getCanonicalSignature();
+    if (witnessThunkSig) {
+      reqToWitnessThunkSigMap = reqToWitnessThunkSigMap.subst(
+          witnessThunkSig.getGenericEnvironment()
+              ->getForwardingSubstitutionMap());
+    }
     return;
   }
 
@@ -219,4 +225,7 @@ RequirementEnvironment::RequirementEnvironment(
                                           std::move(genericParamTypes),
                                           std::move(requirements),
                                           /*allowInverses=*/false);
+  reqToWitnessThunkSigMap = reqToWitnessThunkSigMap.subst(
+      witnessThunkSig.getGenericEnvironment()
+          ->getForwardingSubstitutionMap());
 }

lib/SILGen/SILGenType.cpp

@@ -727,7 +727,7 @@ SILGenModule::getWitnessTable(NormalProtocolConformance *conformance) {
 }
 
 SILFunction *SILGenModule::emitProtocolWitness(
-    ProtocolConformanceRef conformance, SILLinkage linkage,
+    ProtocolConformanceRef origConformance, SILLinkage linkage,
     SerializedKind_t serializedKind, SILDeclRef requirement,
     SILDeclRef witnessRef, IsFreeFunctionWitness_t isFree, Witness witness) {
   auto requirementInfo =
@@ -767,7 +767,8 @@ SILFunction *SILGenModule::emitProtocolWitness(
   // The type of the witness thunk.
   auto reqtSubstTy = cast<AnyFunctionType>(
     reqtOrigTy->substGenericArgs(reqtSubMap)
-      ->getReducedType(genericSig));
+              ->mapTypeOutOfContext()
+              ->getCanonicalType());
 
   // Rewrite the conformance in terms of the requirement environment's Self
   // type, which might have a different generic signature than the type
@@ -777,14 +778,18 @@ SILFunction *SILGenModule::emitProtocolWitness(
   // in a protocol extension, the generic signature will have an additional
   // generic parameter representing Self, so the generic parameters of the
   // class will all be shifted down by one.
-  if (reqtSubMap) {
-    auto requirement = conformance.getProtocol();
-    auto self = requirement->getSelfInterfaceType()->getCanonicalType();
+  auto conformance = reqtSubMap.lookupConformance(M.getASTContext().TheSelfType,
+                                                  origConformance.getProtocol())
+      .mapConformanceOutOfContext();
+  ASSERT(conformance.isAbstract() == origConformance.isAbstract());
 
-    conformance = reqtSubMap.lookupConformance(self, requirement);
-
-    if (genericEnv)
-      reqtSubMap = reqtSubMap.subst(genericEnv->getForwardingSubstitutionMap());
+  ProtocolConformance *manglingConformance = nullptr;
+  if (conformance.isConcrete()) {
+    manglingConformance = conformance.getConcrete();
+    if (auto *inherited = dyn_cast<InheritedProtocolConformance>(manglingConformance)) {
+      manglingConformance = inherited->getInheritedConformance();
+      conformance = ProtocolConformanceRef(manglingConformance);
+    }
   }
 
   // Generic signatures where all parameters are concrete are lowered away
@@ -806,20 +811,14 @@ SILFunction *SILGenModule::emitProtocolWitness(
   // But this is expensive, so we only do it for coroutine lowering.
   // When they're part of the AST function type, we can remove this
   // parameter completely.
+  bool allowDuplicateThunk = false;
   std::optional<SubstitutionMap> witnessSubsForTypeLowering;
   if (auto accessor = dyn_cast<AccessorDecl>(requirement.getDecl())) {
     if (accessor->isCoroutine()) {
       witnessSubsForTypeLowering =
         witness.getSubstitutions().mapReplacementTypesOutOfContext();
-    }
-  }
-
-  ProtocolConformance *manglingConformance = nullptr;
-  if (conformance.isConcrete()) {
-    manglingConformance = conformance.getConcrete();
-    if (auto *inherited = dyn_cast<InheritedProtocolConformance>(manglingConformance)) {
-      manglingConformance = inherited->getInheritedConformance();
-      conformance = ProtocolConformanceRef(manglingConformance);
+      if (accessor->isRequirementWithSynthesizedDefaultImplementation())
+        allowDuplicateThunk = true;
     }
   }
 
@@ -863,8 +862,9 @@ SILFunction *SILGenModule::emitProtocolWitness(
     InlineStrategy = AlwaysInline;
 
   SILFunction *f = M.lookUpFunction(nameBuffer);
-  if (f)
+  if (allowDuplicateThunk && f)
     return f;
+  ASSERT(!f);
 
   SILGenFunctionBuilder builder(*this);
   f = builder.createFunction(
@@ -891,7 +891,8 @@ SILFunction *SILGenModule::emitProtocolWitness(
   bool isPreconcurrency = false;
   if (conformance.isConcrete()) {
     if (auto *C =
-            dyn_cast<NormalProtocolConformance>(conformance.getConcrete()))
+          dyn_cast<NormalProtocolConformance>(
+            conformance.getConcrete()->getRootConformance()))
       isPreconcurrency = C->isPreconcurrency();
   }
 

lib/Sema/TypeCheckProtocol.cpp

@@ -1007,7 +1007,8 @@ findMissingGenericRequirementForSolutionFix(
   };
 
   auto selfTy = conformance->getProtocol()->getSelfInterfaceType()
-      .subst(reqEnvironment.getRequirementToWitnessThunkSubs());
+      .subst(reqEnvironment.getRequirementToWitnessThunkSubs())
+        ->mapTypeOutOfContext();
 
   auto sig = conformance->getGenericSignature();
   auto *env = conformance->getGenericEnvironment();
@@ -1162,14 +1163,9 @@ swift::matchWitness(WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
     // the required type and the witness type.
     cs.emplace(dc, ConstraintSystemFlags::AllowFixes);
 
-    auto syntheticSig = reqEnvironment.getWitnessThunkSignature();
-    auto *syntheticEnv = syntheticSig.getGenericEnvironment();
     auto reqSubMap = reqEnvironment.getRequirementToWitnessThunkSubs();
 
     Type selfTy = proto->getSelfInterfaceType().subst(reqSubMap);
-    if (syntheticEnv)
-      selfTy = syntheticEnv->mapTypeIntoContext(selfTy);
-
 
     // Open up the type of the requirement.
     SmallVector<OpenedType, 4> reqReplacements;
@@ -1194,10 +1190,6 @@ swift::matchWitness(WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
       if (replacedInReq->hasError())
         continue;
 
-      if (syntheticEnv) {
-        replacedInReq = syntheticEnv->mapTypeIntoContext(replacedInReq);
-      }
-
       if (auto packType = replacedInReq->getAs<PackType>()) {
         if (auto unwrapped = packType->unwrapSingletonPackExpansion())
           replacedInReq = unwrapped->getPatternType();
@@ -7295,7 +7287,8 @@ void TypeChecker::inferDefaultWitnesses(ProtocolDecl *proto) {
         RequirementEnvironment reqEnv(proto, reqSig, proto, nullptr, nullptr);
         auto match =
             RequirementMatch(asd, MatchKind::ExactMatch, asdTy, reqEnv);
-        match.WitnessSubstitutions = reqEnv.getRequirementToWitnessThunkSubs();
+        match.WitnessSubstitutions = reqEnv.getRequirementToWitnessThunkSubs()
+                                           .mapReplacementTypesOutOfContext();
         checker.recordWitness(asd, match);
       }
     }

test/SILGen/conditional_conformance.swift

@@ -20,7 +20,7 @@ extension Conformance: P1 where A: P2 {
 // This is defined below but is emitted before any witness tables.
 // Just make sure it does not have a generic signature.
 //
-// CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s23conditional_conformance16SameTypeConcreteVyxGAA2P1AASiRszlAaEP6normalyyFTW : $@convention(witness_method: P1) (@in_guaranteed SameTypeConcrete<Int>) -> ()
+// CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s23conditional_conformance11ConformanceVyxGAA2P1A2A2P2RzlAaEP6normalyyFTW : $@convention(witness_method: P1) <τ_0_0 where τ_0_0 : P2> (@in_guaranteed Conformance<τ_0_0>) -> () {
 
 
 // CHECK-LABEL: sil_witness_table hidden <A where A : P2> Conformance<A>: P1 module conditional_conformance {
@@ -48,8 +48,8 @@ extension SameTypeConcrete: P1 where B == Int {
 }
 
 // CHECK-LABEL: sil_witness_table hidden <B where B == Int> SameTypeConcrete<B>: P1 module conditional_conformance {
-// CHECK-NEXT:    method #P1.normal: <Self where Self : P1> (Self) -> () -> () : @$s23conditional_conformance16SameTypeConcreteVyxGAA2P1AASiRszlAaEP6normalyyFTW	// protocol witness for P1.normal() in conformance <A> SameTypeConcrete<A>
-// CHECK-NEXT:    method #P1.generic: <Self where Self : P1><T where T : P3> (Self) -> (T) -> () : @$s23conditional_conformance16SameTypeConcreteVyxGAA2P1AASiRszlAaEP7genericyyqd__AA2P3Rd__lFTW	// protocol witness for P1.generic<A>(_:) in conformance <A> SameTypeConcrete<A>
+// CHECK-NEXT:    method #P1.normal: <Self where Self : P1> (Self) -> () -> () : @$s23conditional_conformance16SameTypeConcreteVySiGAA2P1A2aEP6normalyyFTW	// protocol witness for P1.normal() in conformance SameTypeConcrete<Int>
+// CHECK-NEXT:    method #P1.generic: <Self where Self : P1><T where T : P3> (Self) -> (T) -> () : @$s23conditional_conformance16SameTypeConcreteVySiGAA2P1A2aEP7genericyyqd__AA2P3Rd__lFTW	// protocol witness for P1.generic<A>(_:) in conformance SameTypeConcrete<Int>
 // CHECK-NEXT:  }
 
 struct SameTypeGeneric<C, D> {}

test/SILGen/constrained_extensions.swift

@@ -254,7 +254,7 @@ extension S: HasX where T == Int {
   }
 }
 
-// CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s22constrained_extensions1SVyxGAA4HasXAASiRszlAaEP1xSivMTW : $@yield_once @convention(witness_method: HasX) @substituted <τ_0_0> (@inout τ_0_0) -> @yields @inout Int for <S<Int>> {
+// CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s22constrained_extensions1SVySiGAA4HasXA2aEP1xSivMTW : $@yield_once @convention(witness_method: HasX) @substituted <τ_0_0> (@inout τ_0_0) -> @yields @inout Int for <S<Int>> {
 // CHECK: [[WITNESS_FN:%.*]] = function_ref @$s22constrained_extensions1SVAASiRszlE1xSivM : $@yield_once @convention(method) (@inout S<Int>) -> @yields @inout Int
 // CHECK: begin_apply [[WITNESS_FN]](%0) : $@yield_once @convention(method) (@inout S<Int>) -> @yields @inout Int
 // CHECK: yield

test/SILGen/witness_same_type.swift

@@ -33,3 +33,58 @@ public struct LazySequenceOf<SS : Sequence, A> : Sequence where SS.Iterator.Elem
     } 
   }
 }
+
+// rdar://problem/155465011
+
+// Due to a series of unfortunate coincidences, we would assign the same mangling
+// to the witness thunks for P1.f() in the S1: P1 and S2: P1 conformance.
+
+public protocol P1 {
+  func f<T: P2>(_: T) where T.A == Self
+}
+
+public protocol P2 {
+  associatedtype A: P1
+}
+
+struct N1: P2 {
+  typealias A = S1
+}
+
+struct N2: P2 {
+  typealias A = S2
+}
+
+struct S1: P1 {
+  func f<T: P2>(_: T) where T.A == Self {
+    print("Hello from S1")
+  }
+}
+
+struct S2: P1 {
+  func f<T: P2>(_: T) where T.A == Self {
+    print("Hello from S2")
+  }
+}
+
+// CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17witness_same_type2S1VAA2P1A2aDP1fyyqd__1AQyd__RszAA2P2Rd__lFTW : $@convention(witness_method: P1) <τ_0_0 where τ_0_0 : P2, τ_0_0.A == S1> (@in_guaranteed τ_0_0, @in_guaranteed S1) -> () {
+
+// CHECK-LABEL: sil private [transparent] [thunk] [ossa] @$s17witness_same_type2S2VAA2P1A2aDP1fyyqd__1AQyd__RszAA2P2Rd__lFTW : $@convention(witness_method: P1) <τ_0_0 where τ_0_0 : P2, τ_0_0.A == S2> (@in_guaranteed τ_0_0, @in_guaranteed S2) -> () {
+
+// CHECK-LABEL: sil_witness_table hidden N1: P2 module witness_same_type {
+// CHECK-NEXT:    associated_conformance (A: P1): S1: P1 module witness_same_type
+// CHECK-NEXT:    associated_type A: S1
+// CHECK-NEXT:  }
+
+// CHECK-LABEL: sil_witness_table hidden N2: P2 module witness_same_type {
+// CHECK-NEXT:    associated_conformance (A: P1): S2: P1 module witness_same_type
+// CHECK-NEXT:    associated_type A: S2
+// CHECK-NEXT:  }
+
+// CHECK-LABEL: sil_witness_table hidden S1: P1 module witness_same_type {
+// CHECK-NEXT:    method #P1.f: <Self><T where Self == T.A, T : P2> (Self) -> (T) -> () : @$s17witness_same_type2S1VAA2P1A2aDP1fyyqd__1AQyd__RszAA2P2Rd__lFTW
+// CHECK-NEXT:  }
+
+// CHECK-LABEL: sil_witness_table hidden S2: P1 module witness_same_type {
+// CHECK-NEXT:    method #P1.f: <Self><T where Self == T.A, T : P2> (Self) -> (T) -> () : @$s17witness_same_type2S2VAA2P1A2aDP1fyyqd__1AQyd__RszAA2P2Rd__lFTW
+// CHECK-NEXT:  }