[SIL] Witness methods store the conformance from which they come.

include/swift/AST/Attr.h

@@ -57,6 +57,7 @@ class TypeAttributes {
   /// If this is an empty attribute specifier, then this will be an invalid loc.
   SourceLoc AtLoc;
   Optional<StringRef> convention = None;
+  Optional<StringRef> conventionWitnessMethodProtocol = None;
 
   // For an opened existential type, the known ID.
   Optional<UUID> OpenedID;

include/swift/AST/DiagnosticsParse.def

@@ -1339,6 +1339,10 @@ ERROR(convention_attribute_expected_name,none,
       "expected convention name identifier in 'convention' attribute", ())
 ERROR(convention_attribute_expected_rparen,none,
       "expected ')' after convention name for 'convention' attribute", ())
+ERROR(convention_attribute_witness_method_expected_colon,none,
+      "expected ':' after 'witness_method' for 'convention' attribute", ())
+ERROR(convention_attribute_witness_method_expected_protocol,none,
+      "expected protocol name in 'witness_method' 'convention' attribute", ())
 
 // objc
 ERROR(attr_objc_missing_colon,none,

include/swift/AST/Types.h

@@ -19,13 +19,14 @@
 
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/GenericParamKey.h"
+#include "swift/AST/Identifier.h"
 #include "swift/AST/Ownership.h"
+#include "swift/AST/ProtocolConformanceRef.h"
 #include "swift/AST/Requirement.h"
 #include "swift/AST/SILLayout.h"
 #include "swift/AST/SubstitutionList.h"
 #include "swift/AST/Type.h"
 #include "swift/AST/TypeAlignments.h"
-#include "swift/AST/Identifier.h"
 #include "swift/Basic/ArrayRefView.h"
 #include "swift/Basic/UUID.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -3347,6 +3348,7 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   //   CanType?                       // if NumResults > 1, all result cache
 
   CanGenericSignature GenericSig;
+  Optional<ProtocolConformanceRef> WitnessMethodConformance;
 
   MutableArrayRef<SILParameterInfo> getMutableParameters() {
     return {getTrailingObjects<SILParameterInfo>(), NumParameters};
@@ -3380,21 +3382,22 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
     return *(reinterpret_cast<CanType *>(ptr) + 1);
   }
 
-  SILFunctionType(GenericSignature *genericSig, ExtInfo ext,
-                  ParameterConvention calleeConvention,
-                  ArrayRef<SILParameterInfo> params,
-                  ArrayRef<SILResultInfo> normalResults,
-                  Optional<SILResultInfo> errorResult, const ASTContext &ctx,
-                  RecursiveTypeProperties properties);
+  SILFunctionType(
+      GenericSignature *genericSig, ExtInfo ext,
+      ParameterConvention calleeConvention, ArrayRef<SILParameterInfo> params,
+      ArrayRef<SILResultInfo> normalResults,
+      Optional<SILResultInfo> errorResult, const ASTContext &ctx,
+      RecursiveTypeProperties properties,
+      Optional<ProtocolConformanceRef> witnessMethodConformance = None);
 
 public:
-  static CanSILFunctionType get(GenericSignature *genericSig,
-                                ExtInfo ext,
-                                ParameterConvention calleeConvention,
-                                ArrayRef<SILParameterInfo> interfaceParams,
-                                ArrayRef<SILResultInfo> interfaceResults,
-                                Optional<SILResultInfo> interfaceErrorResult,
-                                const ASTContext &ctx);
+  static CanSILFunctionType
+  get(GenericSignature *genericSig, ExtInfo ext,
+      ParameterConvention calleeConvention,
+      ArrayRef<SILParameterInfo> interfaceParams,
+      ArrayRef<SILResultInfo> interfaceResults,
+      Optional<SILResultInfo> interfaceErrorResult, const ASTContext &ctx,
+      Optional<ProtocolConformanceRef> witnessMethodConformance = None);
 
   /// Given that this function type uses a C-language convention, return its
   /// formal semantic result type.
@@ -3551,6 +3554,20 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   /// Self type.
   ClassDecl *getWitnessMethodClass(ModuleDecl &M) const;
 
+  /// If this is a @convention(witness_method) function, return the conformance
+  /// for which the method is a witness.
+  ProtocolConformanceRef getWitnessMethodConformance() const {
+    assert(getRepresentation() == Representation::WitnessMethod);
+    return *WitnessMethodConformance;
+  }
+
+  /// If this is a @convention(witness_method) function, return the conformance
+  /// for which the method is a witness, if it isn't that convention, return
+  /// None.
+  Optional<ProtocolConformanceRef> getWitnessMethodConformanceOrNone() const {
+    return WitnessMethodConformance;
+  }
+
   ExtInfo getExtInfo() const { return ExtInfo(SILFunctionTypeBits.ExtInfo); }
 
   /// \brief Returns the language-level calling convention of the function.

include/swift/SIL/SILType.h

@@ -572,10 +572,10 @@ NON_SIL_TYPE(AnyFunction)
 NON_SIL_TYPE(LValue)
 #undef NON_SIL_TYPE
 
-CanSILFunctionType getNativeSILFunctionType(SILModule &M,
-                        Lowering::AbstractionPattern origType,
-                        CanAnyFunctionType substType,
-                        Optional<SILDeclRef> constant = None);
+CanSILFunctionType getNativeSILFunctionType(
+    SILModule &M, Lowering::AbstractionPattern origType,
+    CanAnyFunctionType substType, Optional<SILDeclRef> constant = None,
+    Optional<ProtocolConformanceRef> witnessMethodConformance = None);
 
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, SILType T) {
   T.print(OS);

include/swift/SIL/TypeLowering.h

@@ -67,21 +67,30 @@ inline CanAnyFunctionType adjustFunctionType(CanAnyFunctionType t,
 
 /// Given a SIL function type, return a type that is identical except
 /// for using the given ExtInfo.
-CanSILFunctionType adjustFunctionType(CanSILFunctionType type,
-                                      SILFunctionType::ExtInfo extInfo,
-                                      ParameterConvention calleeConv);
-inline CanSILFunctionType adjustFunctionType(CanSILFunctionType type,
-                                      SILFunctionType::ExtInfo extInfo) {
-  return adjustFunctionType(type, extInfo, type->getCalleeConvention());
+CanSILFunctionType
+adjustFunctionType(CanSILFunctionType type, SILFunctionType::ExtInfo extInfo,
+                   ParameterConvention calleeConv,
+                   Optional<ProtocolConformanceRef> witnessMethodConformance);
+inline CanSILFunctionType
+adjustFunctionType(CanSILFunctionType type, SILFunctionType::ExtInfo extInfo,
+                   Optional<ProtocolConformanceRef> witnessMethodConformance) {
+  return adjustFunctionType(type, extInfo, type->getCalleeConvention(),
+                            witnessMethodConformance);
 }
-inline CanSILFunctionType adjustFunctionType(CanSILFunctionType t,
-                                         SILFunctionType::Representation rep) {
+inline CanSILFunctionType
+adjustFunctionType(CanSILFunctionType t, SILFunctionType::Representation rep,
+                   Optional<ProtocolConformanceRef> witnessMethodConformance) {
   if (t->getRepresentation() == rep) return t;
   auto extInfo = t->getExtInfo().withRepresentation(rep);
-
-  return adjustFunctionType(t, extInfo,
-                    extInfo.hasContext() ? DefaultThickCalleeConvention
-                                         : ParameterConvention::Direct_Unowned);
+
+  return adjustFunctionType(
+      t, extInfo, extInfo.hasContext() ? DefaultThickCalleeConvention
+                                       : ParameterConvention::Direct_Unowned,
+      witnessMethodConformance);
+}
+inline CanSILFunctionType
+adjustFunctionType(CanSILFunctionType t, SILFunctionType::Representation rep) {
+  return adjustFunctionType(t, rep, t->getWitnessMethodConformanceOrNone());
 }
 
 
@@ -658,11 +667,10 @@ class TypeConverter {
   /// whose type cannot be represented in the AST because it is
   /// a polymorphic function value. This function returns the
   /// unsubstituted lowered type of this callback.
-  CanSILFunctionType
-  getMaterializeForSetCallbackType(AbstractStorageDecl *storage,
-                                   CanGenericSignature genericSig,
-                                   Type selfType,
-                                   SILFunctionTypeRepresentation rep);
+  CanSILFunctionType getMaterializeForSetCallbackType(
+      AbstractStorageDecl *storage, CanGenericSignature genericSig,
+      Type selfType, SILFunctionTypeRepresentation rep,
+      Optional<ProtocolConformanceRef> witnessMethodConformance);
 
   /// Return the SILFunctionType for a native function value of the
   /// given type.

include/swift/Serialization/ModuleFormat.h

@@ -54,7 +54,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// in source control, you should also update the comment to briefly
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
-const uint16_t VERSION_MINOR = 376; // Last change: Outlined CopyAddr
+const uint16_t VERSION_MINOR = 377; // Last change: SILFunctionType witness_method conformances
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -742,6 +742,7 @@ namespace decls_block {
                            // followed by result types/conventions, alternating
                            // followed by error result type/convention
                            // followed by generic parameter types
+    // Optionally a protocol conformance (for witness_methods)
     // Trailed by its generic requirements, if any.
   >;
 

lib/AST/ASTContext.cpp

@@ -3805,15 +3805,14 @@ void SILFunctionType::Profile(llvm::FoldingSetNodeID &id,
   if (errorResult) errorResult->profile(id);
 }
 
-SILFunctionType::SILFunctionType(GenericSignature *genericSig, ExtInfo ext,
-                                 ParameterConvention calleeConvention,
-                                 ArrayRef<SILParameterInfo> params,
-                                 ArrayRef<SILResultInfo> normalResults,
-                                 Optional<SILResultInfo> errorResult,
-                                 const ASTContext &ctx,
-                                 RecursiveTypeProperties properties)
-    : TypeBase(TypeKind::SILFunction, &ctx, properties),
-      GenericSig(genericSig) {
+SILFunctionType::SILFunctionType(
+    GenericSignature *genericSig, ExtInfo ext,
+    ParameterConvention calleeConvention, ArrayRef<SILParameterInfo> params,
+    ArrayRef<SILResultInfo> normalResults, Optional<SILResultInfo> errorResult,
+    const ASTContext &ctx, RecursiveTypeProperties properties,
+    Optional<ProtocolConformanceRef> witnessMethodConformance)
+    : TypeBase(TypeKind::SILFunction, &ctx, properties), GenericSig(genericSig),
+      WitnessMethodConformance(witnessMethodConformance) {
 
   SILFunctionTypeBits.HasErrorResult = errorResult.hasValue();
   SILFunctionTypeBits.ExtInfo = ext.Bits;
@@ -3840,6 +3839,13 @@ SILFunctionType::SILFunctionType(GenericSignature *genericSig, ExtInfo ext,
     getMutableAllResultsCache() = CanType();
   }
 #ifndef NDEBUG
+  if (ext.getRepresentation() == Representation::WitnessMethod)
+    assert(WitnessMethodConformance &&
+           "witness_method SIL function without a conformance");
+  else
+    assert(!WitnessMethodConformance &&
+           "non-witness_method SIL function with a conformance");
+
   // Make sure the interface types are sane.
   if (genericSig) {
     for (auto gparam : genericSig->getGenericParams()) {
@@ -3885,12 +3891,11 @@ CanSILBlockStorageType SILBlockStorageType::get(CanType captureType) {
   return CanSILBlockStorageType(storageTy);
 }
 
-CanSILFunctionType SILFunctionType::get(GenericSignature *genericSig,
-                                        ExtInfo ext, ParameterConvention callee,
-                                        ArrayRef<SILParameterInfo> params,
-                                        ArrayRef<SILResultInfo> normalResults,
-                                        Optional<SILResultInfo> errorResult,
-                                        const ASTContext &ctx) {
+CanSILFunctionType SILFunctionType::get(
+    GenericSignature *genericSig, ExtInfo ext, ParameterConvention callee,
+    ArrayRef<SILParameterInfo> params, ArrayRef<SILResultInfo> normalResults,
+    Optional<SILResultInfo> errorResult, const ASTContext &ctx,
+    Optional<ProtocolConformanceRef> witnessMethodConformance) {
   llvm::FoldingSetNodeID id;
   SILFunctionType::Profile(id, genericSig, ext, callee, params, normalResults,
                            errorResult);
@@ -3928,9 +3933,9 @@ CanSILFunctionType SILFunctionType::get(GenericSignature *genericSig,
     properties.removeHasDependentMember();
   }
 
-  auto fnType =
-      new (mem) SILFunctionType(genericSig, ext, callee, params, normalResults,
-                                errorResult, ctx, properties);
+  auto fnType = new (mem)
+      SILFunctionType(genericSig, ext, callee, params, normalResults,
+                      errorResult, ctx, properties, witnessMethodConformance);
   ctx.Impl.SILFunctionTypes.InsertNode(fnType, insertPos);
   return CanSILFunctionType(fnType);
 }

lib/AST/ASTPrinter.cpp

@@ -3500,7 +3500,9 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     }
   }
 
-  void printFunctionExtInfo(SILFunctionType::ExtInfo info) {
+  void printFunctionExtInfo(
+      SILFunctionType::ExtInfo info,
+      Optional<ProtocolConformanceRef> witnessMethodConformance) {
     if (Options.SkipAttributes)
       return;
 
@@ -3529,7 +3531,9 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
         Printer << "objc_method";
         break;
       case SILFunctionType::Representation::WitnessMethod:
-        Printer << "witness_method";
+        Printer << "witness_method: ";
+        printTypeDeclName(
+            witnessMethodConformance->getRequirement()->getDeclaredType());
         break;
       case SILFunctionType::Representation::Closure:
         Printer << "closure";
@@ -3647,7 +3651,8 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
   }
 
   void visitSILFunctionType(SILFunctionType *T) {
-    printFunctionExtInfo(T->getExtInfo());
+    printFunctionExtInfo(T->getExtInfo(),
+                         T->getWitnessMethodConformanceOrNone());
     printCalleeConvention(T->getCalleeConvention());
     if (auto sig = T->getGenericSignature()) {
       printGenericSignature(sig,

lib/AST/ProtocolConformance.cpp

@@ -86,9 +86,18 @@ ProtocolConformanceRef::subst(Type origType,
 
   // If we have a concrete conformance, we need to substitute the
   // conformance to apply to the new type.
-  if (isConcrete())
+  if (isConcrete()) {
+    auto concrete = getConcrete();
+    if (auto classDecl = concrete->getType()->getClassOrBoundGenericClass()) {
+      // If this is a class, we need to traffic in the actual type that
+      // implements the protocol, not 'Self' and not any subclasses (with their
+      // inherited conformances).
+      substType =
+          substType->eraseDynamicSelfType()->getSuperclassForDecl(classDecl);
+    }
     return ProtocolConformanceRef(
       getConcrete()->subst(substType, subs, conformances));
+  }
 
   // Opened existentials trivially conform and do not need to go through
   // substitution map lookup.

lib/AST/Type.cpp

@@ -3583,13 +3583,11 @@ case TypeKind::Id:
 
     if (!changed) return *this;
 
-    return SILFunctionType::get(fnTy->getGenericSignature(),
-                                fnTy->getExtInfo(),
+    return SILFunctionType::get(fnTy->getGenericSignature(), fnTy->getExtInfo(),
                                 fnTy->getCalleeConvention(),
-                                transInterfaceParams,
-                                transInterfaceResults,
-                                transErrorResult,
-                                Ptr->getASTContext());
+                                transInterfaceParams, transInterfaceResults,
+                                transErrorResult, Ptr->getASTContext(),
+                                fnTy->getWitnessMethodConformanceOrNone());
   }
 
   case TypeKind::UnownedStorage:

lib/IRGen/GenProto.cpp

@@ -307,6 +307,8 @@ bool PolymorphicConvention::considerType(CanType type, IsExact_t isExact,
 
 void PolymorphicConvention::considerWitnessSelf(CanSILFunctionType fnType) {
   CanType selfTy = fnType->getSelfInstanceType();
+  auto conformance = fnType->getWitnessMethodConformance();
+  auto protocol = conformance.getRequirement();
 
   // First, bind type metadata for Self.
   Sources.emplace_back(MetadataSource::Kind::SelfMetadata,
@@ -328,11 +330,7 @@ void PolymorphicConvention::considerWitnessSelf(CanSILFunctionType fnType) {
     // 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
+    addSelfWitnessTableFulfillment(selfTy, protocol);
   } else {
     // If the Self type is concrete, we have a witness thunk with a
     // fully substituted Self type. The witness table parameter is not
@@ -344,6 +342,8 @@ void PolymorphicConvention::considerWitnessSelf(CanSILFunctionType fnType) {
     //
     // For now, just fulfill the generic arguments of 'Self'.
     considerType(selfTy, IsInexact, Sources.size() - 1, MetadataPath());
+
+    addSelfWitnessTableFulfillment(selfTy, protocol);
   }
 }
 

lib/IRGen/LoadableByAddress.cpp

@@ -180,7 +180,8 @@ getNewSILFunctionTypePtr(GenericEnvironment *GenericEnv,
       currSILFunctionType->getCalleeConvention(), newArgTys,
       getNewResults(GenericEnv, currSILFunctionType, Mod),
       currSILFunctionType->getOptionalErrorResult(),
-      currSILFunctionType->getASTContext());
+      currSILFunctionType->getASTContext(),
+      currSILFunctionType->getWitnessMethodConformanceOrNone());
   return newSILFunctionType;
 }
 
@@ -2105,7 +2106,8 @@ static bool rewriteFunctionReturn(StructLoweringState &pass) {
         loweredTy->getGenericSignature(), loweredTy->getExtInfo(),
         loweredTy->getCalleeConvention(), loweredTy->getParameters(),
         newSILResultInfo, loweredTy->getOptionalErrorResult(),
-        F->getModule().getASTContext());
+        F->getModule().getASTContext(),
+        loweredTy->getWitnessMethodConformanceOrNone());
     F->rewriteLoweredTypeUnsafe(NewTy);
     return true;
   } else if (isLargeLoadableType(genEnv, resultTy, pass.Mod)) {