Merge pull request #79908 from eeckstein/fix-witness-table-specialization

embedded: fix specialization of associated conformance entries in witness tables

SwiftCompilerSources/Sources/AST/Conformance.swift

@@ -61,6 +61,11 @@ public struct Conformance: CustomStringConvertible, NoReflectionChildren {
     assert(isSpecialized)
     return SubstitutionMap(bridged: bridged.getSpecializedSubstitutions())
   }
+
+  public func getAssociatedConformance(ofAssociatedType assocType: Type, to proto: ProtocolDecl) -> Conformance {
+    assert(isConcrete)
+    return bridged.getAssociatedConformance(assocType.bridged, proto.bridged).conformance
+  }
 }
 
 public struct ConformanceArray : RandomAccessCollection, CustomReflectable {

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -144,11 +144,14 @@ func specializeWitnessTable(forConformance conformance: Conformance,
     case .associatedType(let requirement, let witness):
       let substType = witness.subst(with: conformance.specializedSubstitutions)
       return .associatedType(requirement: requirement, witness: substType)
-    case .associatedConformance(let requirement, let proto, let witness):
-      if witness.isSpecialized {
-        specializeWitnessTable(forConformance: witness, errorLocation: errorLocation, context, notifyNewWitnessTable)
+    case .associatedConformance(let requirement, let proto, _):
+      let concreteAssociateConf = conformance.getAssociatedConformance(ofAssociatedType: requirement.type, to: proto)
+      if concreteAssociateConf.isSpecialized {
+        specializeWitnessTable(forConformance: concreteAssociateConf,
+                               errorLocation: errorLocation,
+                               context, notifyNewWitnessTable)
       }
-      return .associatedConformance(requirement: requirement, protocol: proto, witness: witness)
+      return .associatedConformance(requirement: requirement, protocol: proto, witness: concreteAssociateConf)
     }
   }
   let newWT = context.createWitnessTable(entries: newEntries,conformance: conformance,

include/swift/AST/ASTBridging.h

@@ -3067,6 +3067,8 @@ struct BridgedConformance {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance getGenericConformance() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance getInheritedConformance() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedSubstitutionMap getSpecializedSubstitutions() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance getAssociatedConformance(BridgedASTType assocType,
+                                                                                 BridgedDeclObj proto) const;
 };
 
 struct BridgedConformanceArray {

include/swift/AST/ASTBridgingImpl.h

@@ -557,6 +557,11 @@ BridgedSubstitutionMap BridgedConformance::getSpecializedSubstitutions() const {
   return {specPC->getSubstitutionMap()};
 }
 
+BridgedConformance BridgedConformance::getAssociatedConformance(BridgedASTType assocType, BridgedDeclObj proto) const {
+  return {unbridged().getConcrete()->getAssociatedConformance(assocType.unbridged(),
+                                                              proto.getAs<swift::ProtocolDecl>())};
+}
+
 BridgedConformance BridgedConformanceArray::getAt(SwiftInt index) const {
   return pcArray.unbridged<swift::ProtocolConformanceRef>()[index];
 }

include/swift/IRGen/Linking.h

@@ -1123,7 +1123,7 @@ public:
 
   static LinkEntity forProtocolWitnessTable(const ProtocolConformance *C) {
     if (isEmbedded(C)) {
-      assert(C->getProtocol()->requiresClass());
+      ASSERT(C->getProtocol()->requiresClass());
     }
 
     LinkEntity entity;

lib/IRGen/GenProto.cpp

@@ -959,6 +959,13 @@ namespace {
     }
 
     void addAssociatedConformance(const AssociatedConformance &req) {
+      if (req.getAssociation()->getASTContext().LangOpts.hasFeature(Feature::Embedded) &&
+          !req.getAssociatedRequirement()->requiresClass()) {
+        // If it's not a class protocol, the associated type can never be used to create
+        // an existential. Therefore this witness entry is never used at runtime
+        // in embedded swift.
+        return;
+      }
       Entries.push_back(WitnessTableEntry::forAssociatedConformance(req));
     }
 
@@ -1748,6 +1755,14 @@ public:
       (void)entry;
       SILEntries = SILEntries.slice(1);
 
+      if (IGM.Context.LangOpts.hasFeature(Feature::Embedded) &&
+          !requirement.getAssociatedRequirement()->requiresClass()) {
+        // If it's not a class protocol, the associated type can never be used to create
+        // an existential. Therefore this witness entry is never used at runtime
+        // in embedded swift.
+        return;
+      }
+
       auto associate =
         ConformanceInContext.getAssociatedType(
           requirement.getAssociation())->getCanonicalType();
@@ -1775,7 +1790,10 @@ public:
       if (IGM.Context.LangOpts.hasFeature(Feature::Embedded)) {
         // In Embedded Swift associated-conformance entries simply point to the witness table
         // of the associated conformance.
-        llvm::Constant *witnessEntry = IGM.getAddrOfWitnessTable(associatedConformance.getConcrete());
+        ProtocolConformanceRef assocConf =
+          SILWT->getConformance()->getAssociatedConformance(requirement.getAssociation(),
+                                                            requirement.getAssociatedRequirement());
+        llvm::Constant *witnessEntry = IGM.getAddrOfWitnessTable(assocConf.getConcrete());
         auto &schema = IGM.getOptions().PointerAuth
                           .ProtocolAssociatedTypeWitnessTableAccessFunctions;
         Table.addSignedPointer(witnessEntry, schema, requirement);
@@ -3734,7 +3752,7 @@ llvm::Value *irgen::emitWitnessTableRef(IRGenFunction &IGF,
 
   // In Embedded Swift, only class-bound wtables are allowed.
   if (srcType->getASTContext().LangOpts.hasFeature(Feature::Embedded)) {
-    assert(proto->requiresClass());
+    ASSERT(proto->requiresClass());
   }
 
   assert(Lowering::TypeConverter::protocolRequiresWitnessTable(proto)

test/embedded/existential-class-bound1.swift

@@ -156,6 +156,82 @@ struct S: P2 {
   }
 }
 
+protocol Q3 {
+  func bar()
+}
+
+protocol P3<T>: AnyObject {
+  associatedtype T: Q3
+
+  var t: T { get }
+
+  func foo()
+}
+
+extension P3 {
+  func foo() {
+    t.bar()
+  }
+}
+
+final class C3<T: Q3>: P3 {
+  var t: T
+
+
+  init(t: T) { self.t = t }
+}
+
+struct S3<I: BinaryInteger>: Q3 {
+  var x: I
+
+  func bar() {
+    print(x)
+  }
+}
+
+@inline(never)
+func testP3() -> any P3 {
+  return C3(t: S3(x: 102))
+}
+
+protocol P4<T>: AnyObject {
+  associatedtype T: Q
+
+  var t: T { get }
+
+  func foo()
+}
+
+extension P4 {
+  func foo() {
+    t.bar()
+  }
+}
+
+final class C4<T: Q>: P4 {
+  var t: T
+
+
+  init(t: T) { self.t = t }
+}
+
+class K4<I: BinaryInteger>: Q {
+  var x: I
+
+  init(x: I) { self.x = x }
+
+  func bar() {
+    print(x)
+  }
+}
+
+@inline(never)
+func testP4() -> any P4 {
+  return C4(t: K4(x: 437))
+}
+
+
+
 @main
 struct Main {
     static func main() {
@@ -181,6 +257,10 @@ struct Main {
         // CHECK: Derived2.bar()
         testConditionalConformance(t: S(i: 27))
         // CHECK: 27
+        testP3().foo()
+        // CHECK: 102
+        testP4().foo()
+        // CHECK: 437
     }
 }