[lld-macho] Add swift support to ObjC category merger (llvm#95124)

Currently ObjC category merger only supports categories defined in ObjC.
But swift supports generating ObjC categories also. This change adds
supports for the later categories.

Author: alx32

lld/MachO/ObjC.cpp

@@ -675,15 +675,18 @@ void ObjcCategoryMerger::parseProtocolListInfo(const ConcatInputSection *isec,
   ptrList.structCount += protocolCount;
   ptrList.structSize = target->wordSize;
 
-  uint32_t expectedListSize =
+  [[maybe_unused]] uint32_t expectedListSize =
       (protocolCount * target->wordSize) +
       /*header(count)*/ protocolListHeaderLayout.totalSize +
       /*extra null value*/ target->wordSize;
-  assert(expectedListSize == ptrListSym->isec()->data.size() &&
-         "Protocol list does not match expected size");
 
-  // Suppress unsuded var warning
-  (void)expectedListSize;
+  // On Swift, the protocol list does not have the extra (unnecessary) null
+  [[maybe_unused]] uint32_t expectedListSizeSwift =
+      expectedListSize - target->wordSize;
+
+  assert((expectedListSize == ptrListSym->isec()->data.size() ||
+          expectedListSizeSwift == ptrListSym->isec()->data.size()) &&
+         "Protocol list does not match expected size");
 
   uint32_t off = protocolListHeaderLayout.totalSize;
   for (uint32_t inx = 0; inx < protocolCount; ++inx) {
@@ -1148,10 +1151,9 @@ void ObjcCategoryMerger::collectAndValidateCategoriesData() {
       if (nlCategories.count(categorySym))
         continue;
 
-      // We only support ObjC categories (no swift + @objc)
-      // TODO: Support swift + @objc categories also
-      if (!categorySym->getName().starts_with(objc::symbol_names::category))
-        continue;
+      assert(categorySym->getName().starts_with(objc::symbol_names::category) ||
+             categorySym->getName().starts_with(
+                 objc::symbol_names::swift_objc_category));
 
       auto *catBodyIsec = dyn_cast<ConcatInputSection>(categorySym->isec());
       assert(catBodyIsec &&

lld/MachO/ObjC.h

@@ -32,6 +32,8 @@ constexpr const char categoryInstanceMethods[] =
 constexpr const char categoryClassMethods[] =
     "__OBJC_$_CATEGORY_CLASS_METHODS_";
 constexpr const char categoryProtocols[] = "__OBJC_CATEGORY_PROTOCOLS_$_";
+
+constexpr const char swift_objc_category[] = "__CATEGORY_";
 } // namespace symbol_names
 
 // Check for duplicate method names within related categories / classes.

lld/test/MachO/objc-category-merging-minimal.s

@@ -23,6 +23,10 @@
 # RUN: llvm-objdump --objc-meta-data --macho merge_base_class_minimal_no_merge.dylib  | FileCheck %s --check-prefixes=NO_MERGE_INTO_BASE
 # RUN: llvm-objdump --objc-meta-data --macho merge_base_class_minimal_yes_merge.dylib | FileCheck %s --check-prefixes=YES_MERGE_INTO_BASE
 
+############ Test merging swift category into the base class ############
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-macos -o MyBaseClassSwiftExtension.o MyBaseClassSwiftExtension.s
+# RUN: %lld -arch arm64 -dylib -o merge_base_class_swift_minimal_yes_merge.dylib -objc_category_merging MyBaseClassSwiftExtension.o merge_base_class_minimal.o
+# RUN: llvm-objdump --objc-meta-data --macho merge_base_class_swift_minimal_yes_merge.dylib | FileCheck %s --check-prefixes=YES_MERGE_INTO_BASE_SWIFT
 
 #### Check merge categories enabled ###
 # Check that the original categories are not there
@@ -77,6 +81,21 @@ YES_MERGE_INTO_BASE-NEXT: name {{.*}} baseInstanceMethod
 YES_MERGE_INTO_BASE-NEXT: types {{.*}} v16@0:8
 YES_MERGE_INTO_BASE-NEXT: imp -[MyBaseClass baseInstanceMethod]
 
+
+#### Check merge swift category into base class ###
+YES_MERGE_INTO_BASE_SWIFT: _OBJC_CLASS_$_MyBaseClass
+YES_MERGE_INTO_BASE_SWIFT-NEXT: _OBJC_METACLASS_$_MyBaseClass
+YES_MERGE_INTO_BASE_SWIFT: baseMethods
+YES_MERGE_INTO_BASE_SWIFT-NEXT: entsize 24
+YES_MERGE_INTO_BASE_SWIFT-NEXT: count 2
+YES_MERGE_INTO_BASE_SWIFT-NEXT: name {{.*}} swiftMethod
+YES_MERGE_INTO_BASE_SWIFT-NEXT: types {{.*}} v16@0:8
+YES_MERGE_INTO_BASE_SWIFT-NEXT: imp _$sSo11MyBaseClassC0abC14SwiftExtensionE11swiftMethodyyFTo
+YES_MERGE_INTO_BASE_SWIFT-NEXT: name {{.*}} baseInstanceMethod
+YES_MERGE_INTO_BASE_SWIFT-NEXT: types {{.*}} v16@0:8
+YES_MERGE_INTO_BASE_SWIFT-NEXT: imp -[MyBaseClass baseInstanceMethod]
+
+
 #--- a64_fakedylib.s
 
     .section    __DATA,__objc_data
@@ -279,3 +298,74 @@ L_OBJC_IMAGE_INFO:
 	.long	0
 	.long	64
 .subsections_via_symbols
+
+
+#--- MyBaseClassSwiftExtension.s
+; xcrun -sdk macosx swiftc -emit-assembly MyBaseClassSwiftExtension.swift -import-objc-header YourProject-Bridging-Header.h -o MyBaseClassSwiftExtension.s
+;  ================== Generated from Swift: ==================
+; import Foundation
+; extension MyBaseClass {
+;     @objc func swiftMethod() {
+;     }
+; }
+;  ================== Generated from Swift ===================
+	.private_extern	_$sSo11MyBaseClassC0abC14SwiftExtensionE11swiftMethodyyF
+	.globl	_$sSo11MyBaseClassC0abC14SwiftExtensionE11swiftMethodyyF
+	.p2align	2
+_$sSo11MyBaseClassC0abC14SwiftExtensionE11swiftMethodyyF:
+	.cfi_startproc
+	mov	w0, #0
+	ret
+	.cfi_endproc
+
+	.p2align	2
+_$sSo11MyBaseClassC0abC14SwiftExtensionE11swiftMethodyyFTo:
+	.cfi_startproc
+	mov	w0, #0
+	ret
+	.cfi_endproc
+
+	.section	__TEXT,__cstring,cstring_literals
+	.p2align	4, 0x0
+l_.str.25.MyBaseClassSwiftExtension:
+	.asciz	"MyBaseClassSwiftExtension"
+
+	.section	__TEXT,__objc_methname,cstring_literals
+"L_selector_data(swiftMethod)":
+	.asciz	"swiftMethod"
+
+	.section	__TEXT,__cstring,cstring_literals
+"l_.str.7.v16@0:8":
+	.asciz	"v16@0:8"
+
+	.section	__DATA,__objc_data
+	.p2align	3, 0x0
+__CATEGORY_INSTANCE_METHODS_MyBaseClass_$_MyBaseClassSwiftExtension:
+	.long	24
+	.long	1
+	.quad	"L_selector_data(swiftMethod)"
+	.quad	"l_.str.7.v16@0:8"
+	.quad	_$sSo11MyBaseClassC0abC14SwiftExtensionE11swiftMethodyyFTo
+
+	.section	__DATA,__objc_const
+	.p2align	3, 0x0
+__CATEGORY_MyBaseClass_$_MyBaseClassSwiftExtension:
+	.quad	l_.str.25.MyBaseClassSwiftExtension
+	.quad	_OBJC_CLASS_$_MyBaseClass
+	.quad	__CATEGORY_INSTANCE_METHODS_MyBaseClass_$_MyBaseClassSwiftExtension
+	.quad	0
+	.quad	0
+	.quad	0
+	.quad	0
+	.long	60
+	.space	4
+
+	.section	__DATA,__objc_catlist,regular,no_dead_strip
+	.p2align	3, 0x0
+_objc_categories:
+	.quad	__CATEGORY_MyBaseClass_$_MyBaseClassSwiftExtension
+
+	.no_dead_strip	_main
+	.no_dead_strip	l_entry_point
+
+.subsections_via_symbols