Merge pull request #12354 from slavapestov/try-again-this-time-with-more-self

Re-apply "Fix issue with 'Self' metadata when class conforms to protocol with default implementations"

Author: slavapestov

include/swift/AST/DiagnosticsSema.def

@@ -1452,6 +1452,11 @@ ERROR(witness_requires_dynamic_self,none,
       "method %0 in non-final class %1 must return `Self` to conform to "
       "protocol %2",
       (DeclName, Type, Type))
+ERROR(witness_requires_class_implementation,none,
+      "method %0 in non-final class %1 cannot be implemented in a "
+      "protocol extension because it returns `Self` and has associated type "
+      "requirements",
+      (DeclName, Type))
 ERROR(witness_not_accessible_proto,none,
       "%select{initializer %1|method %1|%select{|setter for }2property %1"
       "|subscript%select{| setter}2}0 must be declared "

include/swift/AST/Types.h

@@ -3520,10 +3520,16 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
 
   CanType getSelfInstanceType() const;
 
-  /// If this a @convention(witness_method) function with an abstract
-  /// self parameter, return the protocol constraint for the Self type.
+  /// If this is a @convention(witness_method) function with a protocol
+  /// constrained self parameter, return the protocol constraint for
+  /// the Self type.
   ProtocolDecl *getDefaultWitnessMethodProtocol(ModuleDecl &M) const;
 
+  /// If this is a @convention(witness_method) function with a class
+  /// constrained self parameter, return the class constraint for the
+  /// Self type.
+  ClassDecl *getWitnessMethodClass(ModuleDecl &M) const;
+
   ExtInfo getExtInfo() const { return ExtInfo(SILFunctionTypeBits.ExtInfo); }
 
   /// \brief Returns the language-level calling convention of the function.

lib/IRGen/GenProto.cpp

@@ -314,18 +314,35 @@ void PolymorphicConvention::considerWitnessSelf(CanSILFunctionType fnType) {
                        selfTy);
 
   if (auto *proto = fnType->getDefaultWitnessMethodProtocol(M)) {
-    // The Self type is abstract, so we must pass in a witness table.
+    // The Self type is abstract, so we can fulfill its metadata from
+    // the Self metadata parameter.
     addSelfMetadataFulfillment(selfTy);
 
-    // Look at the witness table for the conformance.
+    // The witness table for the Self : P conformance can be
+    // fulfilled from the Self witness table parameter.
     Sources.emplace_back(MetadataSource::Kind::SelfWitnessTable,
                          MetadataSource::InvalidSourceIndex,
                          selfTy);
     addSelfWitnessTableFulfillment(selfTy, proto);
+  } else if (auto *classDecl = fnType->getWitnessMethodClass(M)) {
+    // The Self type is abstract, so we can fulfill its metadata from
+    // the Self metadata parameter.
+    addSelfMetadataFulfillment(selfTy);
+
+    // FIXME: We should fulfill the witness table too, but we don't
+    // have the original protocol anymore -- we should store it as part
+    // of the @convention(witness_method) bit in the SILFunctionType's
+    // ExtInfo
   } else {
     // If the Self type is concrete, we have a witness thunk with a
     // fully substituted Self type. The witness table parameter is not
     // used.
+    //
+    // FIXME: As above, we should fulfill the Self metadata and
+    // conformance from our two special paramaters here. However, the
+    // Self metadata will be inexact.
+    //
+    // For now, just fulfill the generic arguments of 'Self'.
     considerType(selfTy, IsInexact, Sources.size() - 1, MetadataPath());
   }
 }

lib/SIL/SILFunctionType.cpp

@@ -100,10 +100,10 @@ SILFunctionType::getDefaultWitnessMethodProtocol(ModuleDecl &M) const {
   assert(getRepresentation() == SILFunctionTypeRepresentation::WitnessMethod);
   auto selfTy = getSelfInstanceType();
   if (auto paramTy = dyn_cast<GenericTypeParamType>(selfTy)) {
+    assert(paramTy->getDepth() == 0 && paramTy->getIndex() == 0);
     auto superclass = GenericSig->getSuperclassBound(paramTy, M);
     if (superclass)
       return nullptr;
-    assert(paramTy->getDepth() == 0 && paramTy->getIndex() == 0);
     auto protos = GenericSig->getConformsTo(paramTy, M);
     assert(protos.size() == 1);
     return protos[0];
@@ -112,6 +112,20 @@ SILFunctionType::getDefaultWitnessMethodProtocol(ModuleDecl &M) const {
   return nullptr;
 }
 
+ClassDecl *
+SILFunctionType::getWitnessMethodClass(ModuleDecl &M) const {
+  auto selfTy = getSelfInstanceType();
+  auto genericSig = getGenericSignature();
+  if (auto paramTy = dyn_cast<GenericTypeParamType>(selfTy)) {
+    assert(paramTy->getDepth() == 0 && paramTy->getIndex() == 0);
+    auto superclass = genericSig->getSuperclassBound(paramTy, M);
+    if (superclass)
+      return superclass->getClassOrBoundGenericClass();
+  }
+
+  return nullptr;
+}
+
 static CanType getKnownType(Optional<CanType> &cacheSlot, ASTContext &C,
                             StringRef moduleName, StringRef typeName) {
   if (!cacheSlot) {

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -712,7 +712,8 @@ DevirtualizationResult swift::tryDevirtualizeClassMethod(FullApplySite AI,
 //===----------------------------------------------------------------------===//
 
 static SubstitutionMap
-getSubstitutionsForProtocolConformance(ProtocolConformanceRef CRef) {
+getSubstitutionsForProtocolConformance(ModuleDecl *M,
+                                       ProtocolConformanceRef CRef) {
   auto C = CRef.getConcrete();
 
   // Walk down to the base NormalProtocolConformance.
@@ -746,8 +747,7 @@ getSubstitutionsForProtocolConformance(ProtocolConformanceRef CRef) {
 
   if (Subs.empty()) {
     auto *DC = NormalC->getDeclContext();
-    return NormalC->getType()
-      ->getContextSubstitutionMap(DC->getParentModule(), DC);
+    return NormalC->getType()->getContextSubstitutionMap(M, DC);
   }
 
   return NormalC->getGenericSignature()->getSubstitutionMap(Subs);
@@ -764,17 +764,50 @@ getSubstitutionsForProtocolConformance(ProtocolConformanceRef CRef) {
 /// are written in terms of the requirement's generic signature need
 /// to be remapped to substitutions suitable for the witness signature.
 ///
+/// Supported remappings are:
+///
+/// - (Concrete witness thunk) Original substitutions:
+///   [Self := ConcreteType, R0 := X0, R1 := X1, ...]
+/// - Requirement generic signature:
+///   <Self : P, R0, R1, ...>
+/// - Witness thunk generic signature:
+///   <W0, W1, ...>
+/// - Remapped substitutions:
+///   [W0 := X0, W1 := X1, ...]
+///
+/// - (Class witness thunk) Original substitutions:
+///   [Self := C<A0, A1>, T0 := X0, T1 := X1, ...]
+/// - Requirement generic signature:
+///   <Self : P, R0, R1, ...>
+/// - Witness thunk generic signature:
+///   <Self : C<B0, B1>, B0, B1, W0, W1, ...>
+/// - Remapped substitutions:
+///   [Self := C<B0, B1>, B0 := A0, B1 := A1, W0 := X0, W1 := X1]
+///
+/// - (Default witness thunk) Original substitutions:
+///   [Self := ConcreteType, R0 := X0, R1 := X1, ...]
+/// - Requirement generic signature:
+///   <Self : P, R0, R1, ...>
+/// - Witness thunk generic signature:
+///   <Self : P, W0, W1, ...>
+/// - Remapped substitutions:
+///   [Self := ConcreteType, W0 := X0, W1 := X1, ...]
+///
 /// \param conformanceRef The (possibly-specialized) conformance
 /// \param requirementSig The generic signature of the requirement
 /// \param witnessThunkSig The generic signature of the witness method
 /// \param origSubs The substitutions from the call instruction
+/// \param isDefaultWitness True if this is a default witness method
+/// \param classWitness The ClassDecl if this is a class witness method
 static SubstitutionMap
 getWitnessMethodSubstitutions(
+    ModuleDecl *mod,
     ProtocolConformanceRef conformanceRef,
     GenericSignature *requirementSig,
     GenericSignature *witnessThunkSig,
     SubstitutionList origSubs,
-    bool isDefaultWitness) {
+    bool isDefaultWitness,
+    ClassDecl *classWitness) {
 
   if (witnessThunkSig == nullptr)
     return SubstitutionMap();
@@ -789,42 +822,67 @@ getWitnessMethodSubstitutions(
 
   // If `Self` maps to a bound generic type, this gives us the
   // substitutions for the concrete type's generic parameters.
-  auto baseSubMap = getSubstitutionsForProtocolConformance(conformanceRef);
+  auto baseSubMap = getSubstitutionsForProtocolConformance(mod, conformanceRef);
 
   unsigned baseDepth = 0;
   auto *rootConformance = conformance->getRootNormalConformance();
   if (auto *witnessSig = rootConformance->getGenericSignature())
     baseDepth = witnessSig->getGenericParams().back()->getDepth() + 1;
 
-  auto origDepth = 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;
+  }
 
   return SubstitutionMap::combineSubstitutionMaps(
       baseSubMap,
       origSubMap,
       CombineSubstitutionMaps::AtDepth,
-      baseDepth,
-      origDepth,
+      /*firstDepth=*/baseDepth,
+      /*secondDepth=*/1,
       witnessThunkSig);
 }
 
 static SubstitutionMap
 getWitnessMethodSubstitutions(SILModule &Module, ApplySite AI, SILFunction *F,
                               ProtocolConformanceRef CRef) {
+  auto witnessFnTy = F->getLoweredFunctionType();
+  assert(witnessFnTy->getRepresentation() ==
+         SILFunctionTypeRepresentation::WitnessMethod);
+
   auto requirementSig = AI.getOrigCalleeType()->getGenericSignature();
-  auto witnessThunkSig = F->getLoweredFunctionType()->getGenericSignature();
+  auto witnessThunkSig = witnessFnTy->getGenericSignature();
 
   SubstitutionList origSubs = AI.getSubstitutions();
 
+  auto *mod = Module.getSwiftModule();
   bool isDefaultWitness =
-    F->getLoweredFunctionType()->getRepresentation()
-      == SILFunctionTypeRepresentation::WitnessMethod &&
-    F->getLoweredFunctionType()->getDefaultWitnessMethodProtocol(
-                                                     *Module.getSwiftModule())
-      == CRef.getRequirement();
+    (witnessFnTy->getDefaultWitnessMethodProtocol(*mod)
+      == CRef.getRequirement());
+  auto *classWitness = witnessFnTy->getWitnessMethodClass(*mod);
 
   return getWitnessMethodSubstitutions(
-      CRef, requirementSig, witnessThunkSig,
-      origSubs, isDefaultWitness);
+      mod, CRef, requirementSig, witnessThunkSig,
+      origSubs, isDefaultWitness, classWitness);
 }
 
 /// Generate a new apply of a function_ref to replace an apply of a