swiftlang · slavapestov · Oct 18, 2016 · Oct 17, 2016
diff --git a/lib/AST/Substitution.cpp b/lib/AST/Substitution.cpp
@@ -85,6 +85,7 @@ Substitution Substitution::subst(Module *module,
         proto->isSpecificProtocol(KnownProtocolKind::AnyObject)) {
       auto classDecl
         = substReplacement->getClassOrBoundGenericClass();
+      assert(classDecl);
       SmallVector<ProtocolConformance *, 1> lookupResults;
       classDecl->lookupConformance(classDecl->getParentModule(),
                                    proto, lookupResults);

diff --git a/lib/SILGen/SILGenDecl.cpp b/lib/SILGen/SILGenDecl.cpp
@@ -1662,6 +1662,9 @@ SILGenModule::getWitnessTable(ProtocolConformance *conformance) {
 static bool maybeOpenCodeProtocolWitness(SILGenFunction &gen,
                                          ProtocolConformance *conformance,
                                          SILLinkage linkage,
+                                         Type selfInterfaceType,
+                                         Type selfType,
+                                         GenericEnvironment *genericEnv,
                                          SILDeclRef requirement,
                                          SILDeclRef witness,
                                          ArrayRef<Substitution> witnessSubs) {
@@ -1670,7 +1673,8 @@ static bool maybeOpenCodeProtocolWitness(SILGenFunction &gen,
       auto reqFn = cast<FuncDecl>(requirement.getDecl());
       assert(reqFn->getAccessorKind() == AccessorKind::IsMaterializeForSet);
       return gen.maybeEmitMaterializeForSetThunk(conformance, linkage,
-                                                 reqFn, witnessFn,
+                                                 selfInterfaceType, selfType,
+                                                 genericEnv, reqFn, witnessFn,
                                                  witnessSubs);
     }
   }
@@ -1689,13 +1693,11 @@ static bool isSelfDerived(Type selfTy, Type t) {
 /// requirement's type to get the type of the witness.
 static CanAnyFunctionType
 substSelfTypeIntoProtocolRequirementType(SILModule &M,
+                                         GenericEnvironment *reqtEnv,
                                          CanGenericFunctionType reqtTy,
-                                         ProtocolConformance *conformance) {
-  if (conformance == nullptr) {
-    // Default witness thunks just get the requirement type without
-    // substituting Self.
-    return reqtTy;
-  }
+                                         ProtocolConformance *conformance,
+                                         GenericEnvironment *&genericEnv,
+                                         SubstitutionMap &thunkSubs) {
 
   auto &C = M.getASTContext();
 
@@ -1732,26 +1734,27 @@ substSelfTypeIntoProtocolRequirementType(SILModule &M,
 
   // Now we look at the generic signature of the requirement.
   // We are going to drop `Self`, and requirements rooted in `Self`.
-  for (auto *param : reqtTy->getGenericParams().slice(1)) {
+  for (auto *param : reqtEnv->getGenericParams().slice(1)) {
     allParams.push_back(param);
     builder.addGenericParameter(param);
   }
 
   RequirementSource source(RequirementSource::Explicit, SourceLoc());
 
   for (auto &reqt : reqtTy->getRequirements()) {
-    if (isSelfDerived(selfTy, reqt.getFirstType()))
-      continue;
-
     switch (reqt.getKind()) {
     case RequirementKind::Conformance:
     case RequirementKind::Superclass:
     case RequirementKind::WitnessMarker:
+      if (isSelfDerived(selfTy, reqt.getFirstType()))
+        continue;
+
       builder.addRequirement(reqt, source);
       break;
 
     case RequirementKind::SameType: {
-      if (isSelfDerived(selfTy, reqt.getSecondType()))
+      if (isSelfDerived(selfTy, reqt.getFirstType()) &&
+          isSelfDerived(selfTy, reqt.getSecondType()))
         continue;
 
       // Substitute the constrained types.
@@ -1774,67 +1777,60 @@ substSelfTypeIntoProtocolRequirementType(SILModule &M,
   auto input = reqtTy->getInput().subst(subs)->getCanonicalType();
   auto result = reqtTy->getResult().subst(subs)->getCanonicalType();
 
+  CanAnyFunctionType reqtSubstTy;
+
+  auto addArchetypeSubstitution = [&](ArchetypeType *archetypeTy, Type substTy) {
+    auto contextTy = ArchetypeBuilder::mapTypeIntoContext(
+        M.getSwiftModule(), genericEnv, substTy);
+
+    thunkSubs.addSubstitution(CanType(archetypeTy), contextTy);
+
+    SmallVector<ProtocolConformanceRef, 2> conformances;
+    for (auto proto : archetypeTy->getConformsTo())
+      conformances.push_back(ProtocolConformanceRef(proto));
+    thunkSubs.addConformances(CanType(archetypeTy),
+                              C.AllocateCopy(conformances));
+  };
+
   // The result might be fully concrete, if the witness had no generic
   // signature, and the requirement had no additional generic parameters
   // beyond `Self`.
   if (!allParams.empty()) {
     builder.finalize(SourceLoc());
 
     auto *sig = builder.getGenericSignature();
+    genericEnv = builder.getGenericEnvironment();
+
+    // Outer generic parameters come from the generic context of
+    // the conformance (which might not be the same as the generic
+    // context of the witness, if the witness is defined in a
+    // superclass, concrete extension or protocol extension).
+    if (auto *outerEnv = conformance->getGenericEnvironment()) {
+      for (auto pair : outerEnv->getArchetypeToInterfaceMap()) {
+        auto archetypeTy = pair.first->castTo<ArchetypeType>();
+        auto substTy = pair.second;
+        addArchetypeSubstitution(archetypeTy, substTy);
+      }
+    }
 
-    return cast<GenericFunctionType>(
+    reqtSubstTy = cast<GenericFunctionType>(
       GenericFunctionType::get(sig, input, result, reqtTy->getExtInfo())
         ->getCanonicalType());
-  }
-
-  return CanFunctionType::get(input, result, reqtTy->getExtInfo());
-}
-
-static GenericEnvironment *
-getSubstitutedGenericEnvironment(SILModule &M,
-                                 GenericEnvironment *reqtEnv,
-                                 CanGenericSignature witnessSig,
-                                 ProtocolConformance *conformance) {
-  if (conformance == nullptr) {
-    // Default witness thunks just use the context archetypes of the requirement.
-    return reqtEnv;
-  }
-
-  SmallVector<GenericTypeParamType *, 4> genericParamTypes;
-  TypeSubstitutionMap witnessContextParams;
-
-  auto selfTy = conformance->getProtocol()->getSelfInterfaceType()
-      ->getCanonicalType();
+  } else {
+    genericEnv = nullptr;
 
-  // Outer generic parameters come from the generic context of
-  // the conformance (which might not be the same as the generic
-  // context of the witness, if the witness is defined in a
-  // superclass, concrete extension or protocol extension).
-  if (auto *outerEnv = conformance->getGenericEnvironment()) {
-    witnessContextParams = outerEnv->getInterfaceToArchetypeMap();
-    for (auto *paramTy : outerEnv->getGenericParams())
-      genericParamTypes.push_back(paramTy);
+    reqtSubstTy = CanFunctionType::get(input, result, reqtTy->getExtInfo());
   }
 
-  for (auto *paramTy : reqtEnv->getGenericParams().slice(1))
-    genericParamTypes.push_back(paramTy);
-
   // Inner generic parameters come from the requirement and
   // also map to the archetypes of the requirement.
-  for (auto pair : reqtEnv->getInterfaceToArchetypeMap()) {
-    // Skip the 'Self' parameter and friends.
-    if (isSelfDerived(selfTy, pair.first))
-      continue;
-
-    auto result = witnessContextParams.insert(pair);
-    assert(result.second);
+  for (auto pair : reqtEnv->getArchetypeToInterfaceMap()) {
+    auto archetypeTy = pair.first->castTo<ArchetypeType>();
+    auto substTy = pair.second.subst(subs);
+    addArchetypeSubstitution(archetypeTy, substTy);
   }
 
-  if (!witnessContextParams.empty())
-    return GenericEnvironment::get(M.getASTContext(), genericParamTypes,
-                                   witnessContextParams);
-
-  return nullptr;
+  return reqtSubstTy;
 }
 
 SILFunction *
@@ -1845,7 +1841,6 @@ SILGenModule::emitProtocolWitness(ProtocolConformance *conformance,
                                   IsFreeFunctionWitness_t isFree,
                                   ArrayRef<Substitution> witnessSubs) {
   auto requirementInfo = Types.getConstantInfo(requirement);
-  auto witnessInfo = Types.getConstantInfo(witness);
   unsigned witnessUncurryLevel = witness.uncurryLevel;
 
   // If the witness is a free function, consider the self argument
@@ -1859,11 +1854,37 @@ SILGenModule::emitProtocolWitness(ProtocolConformance *conformance,
   assert(requirement.uncurryLevel == witnessUncurryLevel &&
          "uncurry level of requirement and witness do not match");
 
+  GenericEnvironment *genericEnv = nullptr;
+
   // Work out the lowered function type of the SIL witness thunk.
   auto reqtOrigTy
     = cast<GenericFunctionType>(requirementInfo.LoweredInterfaceType);
-  auto reqtSubstTy
-    = substSelfTypeIntoProtocolRequirementType(M, reqtOrigTy, conformance);
+  CanAnyFunctionType reqtSubstTy;
+  if (conformance == nullptr) {
+    reqtSubstTy = reqtOrigTy;
+    genericEnv = requirementInfo.GenericEnv;
+  } else {
+    SubstitutionMap thunkSubs;
+    reqtSubstTy
+      = substSelfTypeIntoProtocolRequirementType(M, requirementInfo.GenericEnv,
+                                                 reqtOrigTy, conformance,
+                                                 genericEnv, thunkSubs);
+
+    // Remap witness substitutions to use the correct archetypes.
+    // Sema produces witness substitutions using a mix of archetypes
+    // from the requirement and the witness. To solve downstream
+    // issues in the SIL optimizer, we build all-new archetypes
+    // using the correct generic signature here.
+    //
+    // FIXME: Once Sema records witness substitutions in terms of
+    // interface types, this code as well as the code that builds
+    // thunkSubs in substSelfTypeIntoProtocolRequirementType() can
+    // be removed.
+    auto newWitnessSubs = getASTContext().AllocateCopy(witnessSubs);
+    for (unsigned i = 0, e = witnessSubs.size(); i < e; i++)
+      newWitnessSubs[i] = witnessSubs[i].subst(M.getSwiftModule(), thunkSubs);
+    witnessSubs = newWitnessSubs;
+  }
 
   // Lower the witness thunk type with the requirement's abstraction level.
   auto witnessSILFnType = getNativeSILFunctionType(M,
@@ -1899,16 +1920,6 @@ SILGenModule::emitProtocolWitness(ProtocolConformance *conformance,
     nameBuffer = mangler.finalize();
   }
 
-  CanGenericSignature witnessSig;
-  if (auto gft = dyn_cast<GenericFunctionType>(
-          witnessInfo.LoweredInterfaceType))
-    witnessSig = gft.getGenericSignature();
-
-  // Collect the generic environment for the witness.
-  GenericEnvironment *requirementEnv = requirementInfo.GenericEnv;
-  GenericEnvironment *witnessEnv = getSubstitutedGenericEnvironment(
-      M, requirementEnv, witnessSig, conformance);
-
   // If the thunked-to function is set to be always inlined, do the
   // same with the witness, on the theory that the user wants all
   // calls removed if possible, e.g. when we're able to devirtualize
@@ -1927,7 +1938,7 @@ SILGenModule::emitProtocolWitness(ProtocolConformance *conformance,
 
   auto *f = M.createFunction(
       linkage, nameBuffer, witnessSILFnType,
-      witnessEnv, SILLocation(witness.getDecl()),
+      genericEnv, SILLocation(witness.getDecl()),
       IsNotBare, IsTransparent, isFragile, IsThunk,
       SILFunction::NotRelevant, InlineStrategy);
 
@@ -1937,27 +1948,28 @@ SILGenModule::emitProtocolWitness(ProtocolConformance *conformance,
   PrettyStackTraceSILFunction trace("generating protocol witness thunk", f);
 
   // Create the witness.
+  Type selfInterfaceType;
   Type selfType;
 
   // If the witness is a free function, there is no Self type.
   if (!isFree) {
-    // If we are emitting a witness thunk for a concrete conformance, Self is
-    // just the conforming type.
     if (conformance) {
-      selfType = conformance->getType();
-
-    // For default implementations, Self is the protocol archetype.
+      selfInterfaceType = conformance->getInterfaceType();
     } else {
       auto *proto = cast<ProtocolDecl>(requirement.getDecl()->getDeclContext());
-      selfType = proto->getSelfTypeInContext();
+      selfInterfaceType = proto->getSelfInterfaceType();
     }
+
+    selfType = ArchetypeBuilder::mapTypeIntoContext(
+        M.getSwiftModule(), genericEnv, selfInterfaceType);
   }
 
   SILGenFunction gen(*this, *f);
 
   // Open-code certain protocol witness "thunks".
-  if (maybeOpenCodeProtocolWitness(gen, conformance, linkage, requirement,
-                                   witness, witnessSubs)) {
+  if (maybeOpenCodeProtocolWitness(gen, conformance, linkage,
+                                   selfInterfaceType, selfType, genericEnv,
+                                   requirement, witness, witnessSubs)) {
     assert(!isFree);
     return f;
   }

diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -1122,6 +1122,9 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                                 SILValue buffer, SILValue callbackStorage);
   bool maybeEmitMaterializeForSetThunk(ProtocolConformance *conformance,
                                        SILLinkage linkage,
+                                       Type selfInterfaceType,
+                                       Type selfType,
+                                       GenericEnvironment *genericEnv,
                                        FuncDecl *requirement,
                                        FuncDecl *witness,
                                        ArrayRef<Substitution> witnessSubs);

diff --git a/lib/SILGen/SILGenMaterializeForSet.cpp b/lib/SILGen/SILGenMaterializeForSet.cpp
@@ -160,31 +160,22 @@ struct MaterializeForSetEmitter {
 
   static MaterializeForSetEmitter
   forWitnessThunk(SILGenModule &SGM,
-                  ProtocolConformance *conformance,
-                  SILLinkage linkage,
+                  ProtocolConformance *conformance, SILLinkage linkage,
+                  Type selfInterfaceType, Type selfType,
+                  GenericEnvironment *genericEnv,
                   FuncDecl *requirement, FuncDecl *witness,
                   ArrayRef<Substitution> witnessSubs) {
-    Type selfInterfaceType, selfType;
-    if (conformance) {
-      selfInterfaceType = conformance->getInterfaceType();
-      selfType = conformance->getType();
-    } else {
-      auto *proto = cast<ProtocolDecl>(requirement->getDeclContext());
-      selfInterfaceType = proto->getSelfInterfaceType();
-      selfType = proto->getSelfTypeInContext();
-    }
-
     MaterializeForSetEmitter emitter(SGM, linkage, witness, witnessSubs,
                                      selfInterfaceType, selfType);
 
     if (conformance) {
       if (auto signature = conformance->getGenericSignature())
         emitter.GenericSig = signature->getCanonicalSignature();
-      emitter.GenericEnv = conformance->getGenericEnvironment();
+      emitter.GenericEnv = genericEnv;
     } else {
       auto signature = requirement->getGenericSignatureOfContext();
       emitter.GenericSig = signature->getCanonicalSignature();
-      emitter.GenericEnv = requirement->getGenericEnvironmentOfContext();
+      emitter.GenericEnv = genericEnv;
     }
 
     emitter.RequirementStorage = requirement->getAccessorStorageDecl();
@@ -816,14 +807,17 @@ MaterializeForSetEmitter::createSetterCallback(SILFunction &F,
 bool SILGenFunction::
 maybeEmitMaterializeForSetThunk(ProtocolConformance *conformance,
                                 SILLinkage linkage,
+                                Type selfInterfaceType,
+                                Type selfType,
+                                GenericEnvironment *genericEnv,
                                 FuncDecl *requirement,
                                 FuncDecl *witness,
                                 ArrayRef<Substitution> witnessSubs) {
 
   MaterializeForSetEmitter emitter
     = MaterializeForSetEmitter::forWitnessThunk(
-        SGM, conformance, linkage, requirement, witness,
-        witnessSubs);
+        SGM, conformance, linkage, selfInterfaceType, selfType,
+        genericEnv, requirement, witness, witnessSubs);
 
   if (!emitter.shouldOpenCode())
     return false;

diff --git a/test/SILGen/generic_witness.swift b/test/SILGen/generic_witness.swift
@@ -1,4 +1,5 @@
 // RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s
+// RUN: %target-swift-frontend -emit-ir %s
 
 protocol Runcible {
   func runce<A>(_ x: A)
@@ -49,7 +50,7 @@ struct Canvas<I : Ink> where I.Paint : Pen {
 
 extension Canvas : Medium {}
 
-// CHECK-LABEL: sil hidden [transparent] [thunk] @_TTWuRx15generic_witness3Inkwx5PaintS_3PenrGVS_6Canvasx_S_6MediumS_FS4_4drawuRd__S_6Pencilwd__6StrokezWx7TextureS1__rfT5paintWxS7_S1__6pencilqd___T_ : $@convention(witness_method) <I where I : Ink, I.Paint : Pen><P where P : Pencil> (@in I.Paint, @in P, @in_guaranteed Canvas<I>) -> () {
+// CHECK-LABEL: sil hidden [transparent] [thunk] @_TTWuRx15generic_witness3Inkwx5PaintS_3PenrGVS_6Canvasx_S_6MediumS_FS4_4drawuRd__S_6Pencilwd__6StrokezWx7TextureS1__rfT5paintWxS7_S1__6pencilqd___T_ : $@convention(witness_method) <I where I : Ink, I.Paint : Pen><P where P : Pencil, I.Paint == P.Stroke> (@in I.Paint, @in P, @in_guaranteed Canvas<I>) -> ()
 // CHECK: [[FN:%.*]] = function_ref @_TFV15generic_witness6Canvas4drawuRd__S_6Pencilwx5Paintzwd__6StrokerfT5paintwxS2_6pencilqd___T_ : $@convention(method) <τ_0_0 where τ_0_0 : Ink, τ_0_0.Paint : Pen><τ_1_0 where τ_1_0 : Pencil, τ_0_0.Paint == τ_1_0.Stroke> (@in τ_0_0.Paint, @in τ_1_0, Canvas<τ_0_0>) -> ()
 // CHECK: apply [[FN]]<I, P, I.Paint>({{.*}}) : $@convention(method) <τ_0_0 where τ_0_0 : Ink, τ_0_0.Paint : Pen><τ_1_0 where τ_1_0 : Pencil, τ_0_0.Paint == τ_1_0.Stroke> (@in τ_0_0.Paint, @in τ_1_0, Canvas<τ_0_0>) -> ()
 // CHECK: }
diff --git a/test/SILGen/witness_same_type.swift b/test/SILGen/witness_same_type.swift
@@ -1,4 +1,5 @@
 // RUN: %target-swift-frontend -emit-silgen %s | %FileCheck %s
+// RUN: %target-swift-frontend -emit-ir %s
 
 protocol Fooable {
   associatedtype Bar
@@ -10,7 +11,7 @@ struct X {}
 
 // Ensure that the protocol witness for requirements with same-type constraints
 // is set correctly. <rdar://problem/16369105>
-// CHECK-LABEL: sil hidden [transparent] [thunk] @_TTWV17witness_same_type3FooS_7FooableS_FS1_3foo{{.*}} : $@convention(witness_method) <T where T : Fooable> (@in T, @in_guaranteed Foo) -> @out X
+// CHECK-LABEL: sil hidden [transparent] [thunk] @_TTWV17witness_same_type3FooS_7FooableS_FS1_3foouRd__S1_wx3Barzwd__S2_rfT1xqd___wxS2_ : $@convention(witness_method) <T where T : Fooable, T.Bar == X> (@in T, @in_guaranteed Foo) -> @out X
 struct Foo: Fooable {
   typealias Bar = X