Support class extensions in objcImpl

beccadax · beccadax · commit 68ec6b68de06 · 2024-02-09T21:32:06.000-08:00
Their requirements are now included when typechecking the main body extension, and their conformances are emitted in IRGen.

Fixes rdar://118535473.
diff --git a/lib/IRGen/GenClass.cpp b/lib/IRGen/GenClass.cpp
@@ -1242,7 +1242,7 @@ namespace {
         const ClassLayout &fieldLayout)
         : IGM(IGM), TheEntity(theUnion), TheExtension(nullptr),
           FieldLayout(&fieldLayout) {
-      visitConformances(getClass());
+      visitConformances(getClass()->getImplementationContext());
 
       if (getClass()->isRootDefaultActor()) {
         Ivars.push_back(Field::DefaultActorStorage);
@@ -1310,12 +1310,12 @@ namespace {
     /// Gather protocol records for all of the explicitly-specified Objective-C
     /// protocol conformances.
     void visitConformances(const IterableDeclContext *idc) {
-      auto dc = idc->getAsGenericContext();
-      if (dc->getImplementedObjCContext() != dc) {
-        // We want to use the conformance list imported from the ObjC header.
-        auto importedIDC = cast<IterableDeclContext>(
-                                 dc->getImplementedObjCContext()->getAsDecl());
-        visitConformances(importedIDC);
+      auto decl = idc->getDecl();
+      if (decl->getImplementedObjCDecl()) {
+        // We want to use the conformance lists imported from the ObjC header.
+        for (auto interface : decl->getAllImplementedObjCDecls()) {
+          visitConformances(cast<IterableDeclContext>(interface));
+        }
         return;
       }
 
diff --git a/lib/Sema/TypeCheckDeclObjC.cpp b/lib/Sema/TypeCheckDeclObjC.cpp
@@ -2930,16 +2930,18 @@ class ObjCImplementationChecker {
 
     // Did we actually match this extension to an interface? (In invalid code,
     // we might not have.)
-    auto interfaceDecl = ext->getImplementedObjCDecl();
-    if (!interfaceDecl)
+    auto interfaceDecls = ext->getAllImplementedObjCDecls();
+    if (interfaceDecls.empty())
       return;
 
     // Add the @_objcImplementation extension's members as candidates.
     addCandidates(ext);
 
     // Add its interface's members as requirements.
-    auto interface = cast<IterableDeclContext>(interfaceDecl);
-    addRequirements(interface);
+    for (auto interfaceDecl : interfaceDecls) {
+      auto interface = cast<IterableDeclContext>(interfaceDecl);
+      addRequirements(interface);
+    }
   }
 
 private:
@@ -3422,9 +3424,9 @@ class ObjCImplementationChecker {
 
     // Check only applies to members of implementations, not implementations in
     // their own right.
-    if (!cand->isObjCImplementation()
-          && cand->getDeclContext()->getImplementedObjCContext()
-                 != req->getDeclContext())
+    if (!cand->isObjCImplementation() &&
+          getCategoryName(cand->getDeclContext()->getImplementedObjCContext())
+               != getCategoryName(req->getDeclContext()))
       return MatchOutcome::WrongCategory;
 
     if (cand->getKind() != req->getKind())
diff --git a/test/IRGen/Inputs/objc_implementation.h b/test/IRGen/Inputs/objc_implementation.h
@@ -10,6 +10,12 @@
 
 @end
 
+@interface ImplClass () <NSMutableCopying>
+
+- (void)extensionMethod:(int)param;
+
+@end
+
 
 @interface ImplClass (Category1)
 
diff --git a/test/IRGen/objc_implementation.swift b/test/IRGen/objc_implementation.swift
@@ -17,13 +17,13 @@
 
 // Metaclass
 // CHECK: @"OBJC_METACLASS_$_ImplClass" = global %objc_class { ptr @"OBJC_METACLASS_$_NSObject", ptr @"OBJC_METACLASS_$_NSObject", ptr @_objc_empty_cache, ptr null, {{.*}}ptr [[_METACLASS_DATA_ImplClass:@[^, ]+]]{{.*}}, align 8
-// CHECK: @_PROTOCOLS_ImplClass = internal constant { i64, [1 x ptr] } { i64 1, [1 x ptr] [ptr @_PROTOCOL_NSCopying] }, section "__DATA, __objc_const", align 8
+// CHECK: @_PROTOCOLS_ImplClass = internal constant { i64, [2 x ptr] } { i64 2, [2 x ptr] [ptr @_PROTOCOL_NSCopying, ptr @_PROTOCOL_NSMutableCopying] }, section "__DATA, __objc_const", align 8
 // CHECK: [[_METACLASS_DATA_ImplClass]] = internal constant { i32, i32, i32, i32, ptr, ptr, ptr, ptr, ptr, ptr, ptr } { i32 129, i32 40, i32 40, i32 0, ptr null, ptr @.str.9.ImplClass, ptr null, ptr @_PROTOCOLS_ImplClass, ptr null, ptr null, ptr null }, section "__DATA, __objc_const", align 8
 // TODO: Why the extra i32 field above?
 
 // Class
 // CHECK: [[selector_data__cxx_destruct:@[^, ]+]] = private global [14 x i8] c".cxx_destruct\00", section "__TEXT,__objc_methname,cstring_literals", align 1
-// CHECK: [[_INSTANCE_METHODS_ImplClass:@[^, ]+]] = internal constant { i32, i32, [7 x { ptr, ptr, ptr }] } { i32 24, i32 7, [7 x { ptr, ptr, ptr }] [{ ptr, ptr, ptr } { ptr @"\01L_selector_data(init)", ptr @".str.7.@16@0:8", ptr @"$sSo9ImplClassC19objc_implementationEABycfcTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data_implProperty]], ptr @".str.7.i16@0:8", ptr @"$sSo9ImplClassC19objc_implementationE12implPropertys5Int32VvgTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data_setImplProperty_]], ptr @".str.10.v20@0:8i16", ptr @"$sSo9ImplClassC19objc_implementationE12implPropertys5Int32VvsTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data_mainMethod_]], ptr @".str.10.v20@0:8i16", ptr @"$sSo9ImplClassC19objc_implementationE10mainMethodyys5Int32VFTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr @"\01L_selector_data(copyWithZone:)", ptr @".str.11.@24@0:8^v16", ptr @"$sSo9ImplClassC19objc_implementationE4copy4withypSg10ObjectiveC6NSZoneVSg_tFTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr @"\01L_selector_data(dealloc)", ptr @".str.7.v16@0:8", ptr @"$sSo9ImplClassC19objc_implementationEfDTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data__cxx_destruct]], ptr @".str.7.v16@0:8", ptr @"$sSo9ImplClassCfETo{{(\.ptrauth)?}}" }] }, section "__DATA, __objc_data", align 8
+// CHECK: [[_INSTANCE_METHODS_ImplClass:@[^, ]+]] = internal constant { i32, i32, [9 x { ptr, ptr, ptr }] } { i32 24, i32 9, [9 x { ptr, ptr, ptr }] [{ ptr, ptr, ptr } { ptr @"\01L_selector_data(init)", ptr @".str.7.@16@0:8", ptr @"$sSo9ImplClassC19objc_implementationEABycfcTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data_implProperty]], ptr @".str.7.i16@0:8", ptr @"$sSo9ImplClassC19objc_implementationE12implPropertys5Int32VvgTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data_setImplProperty_]], ptr @".str.10.v20@0:8i16", ptr @"$sSo9ImplClassC19objc_implementationE12implPropertys5Int32VvsTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data_mainMethod_]], ptr @".str.10.v20@0:8i16", ptr @"$sSo9ImplClassC19objc_implementationE10mainMethodyys5Int32VFTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr @"\01L_selector_data(extensionMethod:)", ptr @".str.10.v20@0:8i16", ptr @"$sSo9ImplClassC19objc_implementationE15extensionMethodyys5Int32VFTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr @"\01L_selector_data(copyWithZone:)", ptr @".str.11.@24@0:8^v16", ptr @"$sSo9ImplClassC19objc_implementationE4copy4withypSg10ObjectiveC6NSZoneVSg_tFTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr @"\01L_selector_data(mutableCopyWithZone:)", ptr @".str.11.@24@0:8^v16", ptr @"$sSo9ImplClassC19objc_implementationE11mutableCopy4withypSg10ObjectiveC6NSZoneVSg_tFTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr @"\01L_selector_data(dealloc)", ptr @".str.7.v16@0:8", ptr @"$sSo9ImplClassC19objc_implementationEfDTo{{(\.ptrauth)?}}" }, { ptr, ptr, ptr } { ptr [[selector_data__cxx_destruct]], ptr @".str.7.v16@0:8", ptr @"$sSo9ImplClassCfETo{{(\.ptrauth)?}}" }] }, section "__DATA, __objc_data", align 8
 // CHECK: [[_IVARS_ImplClass:@[^, ]+]] = internal constant { i32, i32, [2 x { ptr, ptr, ptr, i32, i32 }] } { i32 32, i32 2, [2 x { ptr, ptr, ptr, i32, i32 }] [{ ptr, ptr, ptr, i32, i32 } { ptr @"$sSo9ImplClassC19objc_implementationE12implPropertys5Int32VvpWvd", ptr @.str.12.implProperty, ptr @.str.0., i32 2, i32 4 }, { ptr, ptr, ptr, i32, i32 } { ptr @"$sSo9ImplClassC19objc_implementationE13implProperty2So8NSObjectCSgvpWvd", ptr @.str.13.implProperty2, ptr @.str.0., i32 3, i32 8 }] }, section "__DATA, __objc_const", align 8
 // CHECK: [[_PROPERTIES_ImplClass:@[^, ]+]] = internal constant { i32, i32, [1 x { ptr, ptr }] } { i32 16, i32 1, [1 x { ptr, ptr }] [{ ptr, ptr } { ptr @.str.12.implProperty, ptr @".str.18.Ti,N,VimplProperty" }] }, section "__DATA, __objc_const", align 8
 // FIXME: Should just be @_PROTOCOLS_ImplClass, but an IRGen bug causes the protocol list to be emitted twice.
@@ -50,6 +50,7 @@
   final var implProperty2: NSObject?
   
   @objc func mainMethod(_: Int32) { print(implProperty) }
+  @objc func extensionMethod(_: Int32) {}
 
   @objc(copyWithZone:)
   func copy(with zone: NSZone?) -> Any? {
@@ -59,6 +60,11 @@
     return copy
   }
   
+  @objc(mutableCopyWithZone:)
+  func mutableCopy(with zone: NSZone?) -> Any? {
+    return copy(with: zone)
+  }
+
   deinit { print(implProperty) }
 }
 
@@ -199,12 +205,24 @@ public func fn(impl: ImplClass, swiftSub: SwiftSubclass) {
 // ObjC calling convention -[ImplClass mainMethod:]
 // CHECK-LABEL: define internal void @"$sSo9ImplClassC19objc_implementationE10mainMethodyys5Int32VFTo"
 
+// Swift calling convention -[ImplClass extensionMethod:]
+// CHECK-LABEL: define hidden swiftcc void @"$sSo9ImplClassC19objc_implementationE15extensionMethodyys5Int32VF"
+
+// ObjC calling convention -[ImplClass extensionMethod:]
+// CHECK-LABEL: define internal void @"$sSo9ImplClassC19objc_implementationE15extensionMethodyys5Int32VFTo"
+
 // Swift calling convention -[ImplClass copyWithZone:]
 // CHECK-LABEL: define hidden swiftcc void @"$sSo9ImplClassC19objc_implementationE4copy4withypSg10ObjectiveC6NSZoneVSg_tF"
 
 // ObjC calling convention -[ImplClass copyWithZone:]
 // CHECK-LABEL: define internal ptr @"$sSo9ImplClassC19objc_implementationE4copy4withypSg10ObjectiveC6NSZoneVSg_tFTo"
 
+// Swift calling convention -[ImplClass mutableCopyWithZone:]
+// CHECK-LABEL: define hidden swiftcc void @"$sSo9ImplClassC19objc_implementationE11mutableCopy4withypSg10ObjectiveC6NSZoneVSg_tF"
+
+// ObjC calling convention -[ImplClass mutableCopyWithZone:]
+// CHECK-LABEL: define internal ptr @"$sSo9ImplClassC19objc_implementationE11mutableCopy4withypSg10ObjectiveC6NSZoneVSg_tFTo"
+
 // Swift calling convention -[ImplClass dealloc]
 // CHECK-LABEL: define linkonce_odr hidden swiftcc void @"$sSo9ImplClassC19objc_implementationEfD"
   // Access in this function should be DirectToStorage
diff --git a/test/Inputs/clang-importer-sdk/usr/include/Foundation.h b/test/Inputs/clang-importer-sdk/usr/include/Foundation.h
@@ -156,6 +156,10 @@ __attribute__((availability(ios,introduced=8.0)))
 - (id)copyWithZone:(nullable NSZone *)zone;
 @end
 
+@protocol NSMutableCopying
+- (id)mutableCopyWithZone:(nullable NSZone *)zone;
+@end
+
 @interface NSDictionary<KeyType, ObjectType> : NSObject /*<NSCopying, NSMutableCopying, NSSecureCoding, NSFastEnumeration>*/
 @property (readonly) NSUInteger count;
 - (nullable ObjectType)objectForKey:(nonnull KeyType)aKey;
diff --git a/test/decl/ext/Inputs/objc_implementation.h b/test/decl/ext/Inputs/objc_implementation.h
@@ -59,6 +59,13 @@
 
 @end
 
+@interface ObjCClass () <NSCopying>
+
+- (void)extensionMethodFromHeader1:(int)param;
+- (void)extensionMethodFromHeader2:(int)param;
+
+@end
+
 @interface ObjCClass (PresentAdditions)
 
 - (void)categoryMethodFromHeader1:(int)param;
diff --git a/test/decl/ext/objc_implementation.swift b/test/decl/ext/objc_implementation.swift
@@ -12,6 +12,7 @@ protocol EmptySwiftProto {}
   // FIXME: give better diagnostic expected-warning@-6 {{extension for main class interface should provide implementation for instance method 'method(fromHeader3:)'}}
   // expected-warning@-7 {{'@_objcImplementation' extension cannot add conformance to 'EmptySwiftProto'; add this conformance with an ordinary extension}}
   // expected-warning@-8 {{'@_objcImplementation' extension cannot add conformance to 'EmptyObjCProto'; add this conformance in the Objective-C header}}
+  // expected-warning@-9 {{extension for main class interface should provide implementation for instance method 'extensionMethod(fromHeader2:)'}}
 
   func method(fromHeader1: CInt) {
     // OK, provides an implementation for the header's method.
@@ -215,6 +216,15 @@ protocol EmptySwiftProto {}
   @nonobjc public convenience init(notFromHeader6: CInt) {
     // OK
   }
+
+  @objc func extensionMethod(fromHeader1: CInt) {
+    // OK
+  }
+
+  @objc(copyWithZone:) func copy(with zone: NSZone?) -> Any {
+    // OK
+    return self
+  }
 }
 
 @_objcImplementation(PresentAdditions) extension ObjCClass {
diff --git a/test/decl/ext/objc_implementation_conflicts.swift b/test/decl/ext/objc_implementation_conflicts.swift
@@ -170,6 +170,11 @@ import objc_implementation_private
 
   func instanceMethod1(_: CInt) {}
   func instanceMethod2(_: CInt) {}
+
+  @objc func extensionMethod(fromHeader1: CInt) {}
+  @objc func extensionMethod(fromHeader2: CInt) {}
+
+  @objc(copyWithZone:) func copy(with zone: NSZone?) -> Any { self }
 }
 
 @_objcImplementation(PresentAdditions) extension ObjCClass {
