Merge pull request #19210 from DougGregor/keypath-protocol-overrides

Author: swift-ci

include/swift/SIL/SILDeclRef.h

@@ -359,6 +359,11 @@ struct SILDeclRef {
   /// entry overridden by this method.
   SILDeclRef getOverriddenWitnessTableEntry() const;
 
+  /// Return the original protocol requirement that introduced the witness table
+  /// entry overridden by this method.
+  static AbstractFunctionDecl *getOverriddenWitnessTableEntry(
+                                                    AbstractFunctionDecl *func);
+
   /// True if the referenced entity is some kind of thunk.
   bool isThunk() const;
 

lib/IRGen/GenKeyPath.cpp

@@ -940,6 +940,8 @@ emitKeyPathComponent(IRGenModule &IGM,
       } else {
         if (auto overridden = declRef.getOverriddenVTableEntry())
           declRef = overridden;
+        if (auto overridden = declRef.getOverriddenWitnessTableEntry())
+          declRef = overridden;
 
         auto dc = declRef.getDecl()->getDeclContext();
 

lib/IRGen/GenThunk.cpp

@@ -126,6 +126,8 @@ llvm::Constant *
 IRGenModule::getAddrOfMethodDescriptor(SILDeclRef declRef,
                                        ForDefinition_t forDefinition) {
   assert(forDefinition == NotForDefinition);
+  assert(declRef.getOverriddenWitnessTableEntry() == declRef &&
+         "Overriding protocol requirements do not have method descriptors");
   LinkEntity entity = LinkEntity::forMethodDescriptor(declRef);
   return getAddrOfLLVMVariable(entity, Alignment(4), forDefinition,
                                MethodDescriptorStructTy, DebugTypeInfo());

lib/SIL/SILDeclRef.cpp

@@ -782,12 +782,22 @@ SILDeclRef SILDeclRef::getNextOverriddenVTableEntry() const {
 }
 
 SILDeclRef SILDeclRef::getOverriddenWitnessTableEntry() const {
-  ValueDecl *bestOverridden = nullptr;
+  auto bestOverridden =
+    getOverriddenWitnessTableEntry(cast<AbstractFunctionDecl>(getDecl()));
+  return SILDeclRef(bestOverridden, kind, isCurried);
+}
+
+AbstractFunctionDecl *SILDeclRef::getOverriddenWitnessTableEntry(
+                                                 AbstractFunctionDecl *func) {
+  if (!isa<ProtocolDecl>(func->getDeclContext()))
+    return func;
 
-  SmallVector<ValueDecl *, 4> stack;
-  SmallPtrSet<ValueDecl *, 4> visited;
-  stack.push_back(getDecl());
-  visited.insert(getDecl());
+  AbstractFunctionDecl *bestOverridden = nullptr;
+
+  SmallVector<AbstractFunctionDecl *, 4> stack;
+  SmallPtrSet<AbstractFunctionDecl *, 4> visited;
+  stack.push_back(func);
+  visited.insert(func);
 
   while (!stack.empty()) {
     auto current = stack.back();
@@ -802,20 +812,21 @@ SILDeclRef SILDeclRef::getOverriddenWitnessTableEntry() const {
                         cast<ProtocolDecl>(current->getDeclContext()),
                         cast<ProtocolDecl>(bestOverridden->getDeclContext()))
             < 0) {
-        bestOverridden = current;
+        bestOverridden = cast<AbstractFunctionDecl>(current);
       }
 
       continue;
     }
 
     // Add overridden declarations to the stack.
     for (auto overridden : overriddenDecls) {
-      if (visited.insert(overridden).second)
-        stack.push_back(overridden);
+      auto overriddenFunc = cast<AbstractFunctionDecl>(overridden);
+      if (visited.insert(overriddenFunc).second)
+        stack.push_back(overriddenFunc);
     }
   }
 
-  return SILDeclRef(bestOverridden, kind, isCurried);
+  return bestOverridden;
 }
 
 SILDeclRef SILDeclRef::getOverriddenVTableEntry() const {

lib/SILGen/SILGen.h

@@ -442,6 +442,13 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   SubstitutionMap
   getNonMemberVarDeclSubstitutions(VarDecl *var);
 
+  /// Map the substitutions for the original declaration to substitutions for
+  /// the overridden declaration.
+  static SubstitutionMap mapSubstitutionsForWitnessOverride(
+                                               AbstractFunctionDecl *original,
+                                               AbstractFunctionDecl *overridden,
+                                               SubstitutionMap subs);
+
   /// Emit a property descriptor for the given storage decl if it needs one.
   void tryEmitPropertyDescriptor(AbstractStorageDecl *decl);
 

lib/SILGen/SILGenApply.cpp

@@ -44,6 +44,21 @@ using namespace Lowering;
 //                             Utility Functions
 //===----------------------------------------------------------------------===//
 
+SubstitutionMap SILGenModule::mapSubstitutionsForWitnessOverride(
+                                              AbstractFunctionDecl *original,
+                                              AbstractFunctionDecl *overridden,
+                                              SubstitutionMap subs) {
+  // Substitute the 'Self' type of the base protocol.
+  auto origProto = cast<ProtocolDecl>(original->getDeclContext());
+  Type origProtoSelfType = origProto->getProtocolSelfType();
+  auto baseProto = cast<ProtocolDecl>(overridden->getDeclContext());
+  return SubstitutionMap::getProtocolSubstitutions(
+           baseProto,
+           origProtoSelfType.subst(subs),
+           *subs.lookupConformance(origProtoSelfType->getCanonicalType(),
+                                   baseProto));
+}
+
 /// Return the abstraction pattern to use when calling a function value.
 static AbstractionPattern
 getIndirectApplyAbstractionPattern(SILGenFunction &SGF,
@@ -390,23 +405,16 @@ class Callee {
                                  SILLocation l) {
     // Find a witness that has an entry in the witness table.
     if (!c.requiresNewWitnessTableEntry()) {
-      auto origProto =
-        cast<ProtocolDecl>(c.getDecl()->getDeclContext());
-
       // Retrieve the constant that has an entry in the witness table.
+      auto original = cast<AbstractFunctionDecl>(c.getDecl());
       c = c.getOverriddenWitnessTableEntry();
       c = c.asForeign(c.getDecl()->isObjC());
+      auto overridden = cast<AbstractFunctionDecl>(c.getDecl());
 
       // Substitute the 'Self' type of the base protocol.
-      auto baseProto = cast<ProtocolDecl>(c.getDecl()->getDeclContext());
-      Type origProtoSelfType = origProto->getProtocolSelfType();
-      auto baseProtoSubs =
-          SubstitutionMap::getProtocolSubstitutions(
-            baseProto,
-            origProtoSelfType.subst(subs),
-            *subs.lookupConformance(origProtoSelfType->getCanonicalType(),
-                                    baseProto));
-      subs = baseProtoSubs;
+      subs = SILGenModule::mapSubstitutionsForWitnessOverride(original,
+                                                              overridden,
+                                                              subs);
     }
 
     auto &ci = SGF.getConstantInfo(c);

lib/SILGen/SILGenExpr.cpp

@@ -2852,6 +2852,23 @@ static SILFunction *getOrCreateKeyPathGetter(SILGenModule &SGM,
                          ArrayRef<IndexTypePair> indexes,
                          CanType baseType,
                          CanType propertyType) {
+  // If the storage declaration is from a protocol, chase the override chain
+  // back to the declaration whose getter introduced the witness table
+  // entry.
+  if (auto origProto = dyn_cast<ProtocolDecl>(property->getDeclContext())) {
+    auto getter = property->getGetter();
+    if (!SILDeclRef::requiresNewWitnessTableEntry(getter)) {
+      // Find the getter that does have a witness table entry.
+      auto wtableGetter =
+        cast<AccessorDecl>(SILDeclRef::getOverriddenWitnessTableEntry(getter));
+
+      // Substitute the 'Self' type of the base protocol.
+      subs = SILGenModule::mapSubstitutionsForWitnessOverride(
+              getter, wtableGetter, subs);
+      property = wtableGetter->getStorage();
+    }
+  }
+
   auto genericSig = genericEnv
     ? genericEnv->getGenericSignature()->getCanonicalSignature()
     : nullptr;
@@ -2966,6 +2983,23 @@ static SILFunction *getOrCreateKeyPathSetter(SILGenModule &SGM,
                           ArrayRef<IndexTypePair> indexes,
                           CanType baseType,
                           CanType propertyType) {
+  // If the storage declaration is from a protocol, chase the override chain
+  // back to the declaration whose setter introduced the witness table
+  // entry.
+  if (auto origProto = dyn_cast<ProtocolDecl>(property->getDeclContext())) {
+    auto setter = property->getSetter();
+    if (!SILDeclRef::requiresNewWitnessTableEntry(setter)) {
+      // Find the setter that does have a witness table entry.
+      auto wtableSetter =
+        cast<AccessorDecl>(SILDeclRef::getOverriddenWitnessTableEntry(setter));
+
+      // Substitute the 'Self' type of the base protocol.
+      subs = SILGenModule::mapSubstitutionsForWitnessOverride(
+              setter, wtableSetter, subs);
+      property = wtableSetter->getStorage();
+    }
+  }
+
   auto genericSig = genericEnv
     ? genericEnv->getGenericSignature()->getCanonicalSignature()
     : nullptr;

test/IRGen/keypath_witness_overrides.swift

@@ -0,0 +1,13 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -module-name protocol_overrides -emit-module -enable-sil-ownership -enable-resilience -emit-module-path=%t/protocol_overrides.swiftmodule %S/../SILGen/Inputs/protocol_overrides.swift
+// RUN: %target-swift-frontend -module-name keypath_witness_overrides -emit-ir %s -I %t | %FileCheck %s
+
+import protocol_overrides
+
+// CHECK: @keypath = private global
+// CHECK-SAME: %swift.method_descriptor* @"$S18protocol_overrides14OriginalGetterPy7ElementQz5IndexQzcigTq"
+public func getWritableKeyPath<OS: OverridesSetter>(_ c: OS, index: OS.Index) -> AnyKeyPath
+where OS.Index: Hashable {
+  let keypath = \OS.[index]
+  return keypath
+}

test/SILGen/Inputs/protocol_overrides.swift

@@ -0,0 +1,19 @@
+public protocol OriginalGetter {
+  associatedtype Index
+  associatedtype Element
+  
+  subscript (index: Index) -> Element { get }
+}
+
+public protocol OverridesGetter: OriginalGetter {
+  override subscript (index: Index) -> Element { get }
+}
+
+public protocol AddsSetter: OverridesGetter {
+  override subscript (index: Index) -> Element { get set }
+}
+
+public protocol OverridesSetter: AddsSetter {
+  override subscript (index: Index) -> Element { get set }
+}
+

test/SILGen/keypath_witness_overrides.swift

@@ -0,0 +1,26 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -module-name protocol_overrides -emit-module -enable-sil-ownership -enable-resilience -emit-module-path=%t/protocol_overrides.swiftmodule %S/Inputs/protocol_overrides.swift
+// RUN: %target-swift-emit-silgen %s -I %t | %FileCheck %s
+
+// Check that keypath formation properly records the point at which the witness
+// table entry for a protocol requirement is introduced.
+import protocol_overrides
+
+// CHECK-LABEL: sil hidden @$S25keypath_witness_overrides18getWritableKeyPath_5indexs03AnyfG0Cx_5IndexQzt09protocol_C015OverridesSetterRzSHAGRQlF
+func getWritableKeyPath<OS: OverridesSetter>(_ c: OS, index: OS.Index) -> AnyKeyPath
+where OS.Index: Hashable {
+  // CHECK: keypath $WritableKeyPath<OS, OS.Element>,
+  // CHECK-SAME: id #OverridesSetter.subscript!getter.1
+  // CHECK-SAME: getter @$S18protocol_overrides14OriginalGetterPy7ElementQz5IndexQzcipAA15OverridesSetterRzSHAGRQlxxTK
+  // CHECK-SAME: setter @$S18protocol_overrides10AddsSetterPy7ElementQz5IndexQzcipAA09OverridesD0RzSHAGRQlxxTk
+  // CHECK-SAME: indices_equals @$S5Index18protocol_overrides14OriginalGetterPQzAB15OverridesSetterRzSHAERQlTH
+  // CHECK-SAME: indices_hash @$S5Index18protocol_overrides14OriginalGetterPQzAB15OverridesSetterRzSHAERQlTh
+  let keypath = \OS.[index]
+  return keypath
+}
+
+// CHECK-LABEL: sil shared [thunk] @$S18protocol_overrides14OriginalGetterPy7ElementQz5IndexQzcipAA15OverridesSetterRzSHAGRQlxxTK
+// CHECK: witness_method $OS, #OriginalGetter.subscript!getter.1
+
+// CHECK-LABEL: sil shared [thunk] @$S18protocol_overrides10AddsSetterPy7ElementQz5IndexQzcipAA09OverridesD0RzSHAGRQlxxTk
+// CHECK-LABEL: witness_method $OS, #AddsSetter.subscript!setter.1