[lld] [lld-macho] Add support for non-lazy categories to ObjC category merger (PR #91548)

[via llvm-commits]

https://github.com/alx32 created https://github.com/llvm/llvm-project/pull/91548

In ObjC we can have categories that define a `+load` method that is called when the category is loaded. In such cases, we shouldn't optimize the category. These categories are present in the `__objc_nlcatlist` section. So we scan these section for such categories and ignore them from optimization. 

>From ebf03bbdc302bb6cac914f4010a02ceb491ce1ed Mon Sep 17 00:00:00 2001
From: Alex B <alexborcan at meta.com>
Date: Thu, 25 Apr 2024 10:55:56 -0700
Subject: [PATCH] [lld-macho] Add support for non-lazy categories to ObjC
 category merger

---
 lld/MachO/ObjC.cpp                            | 24 ++++++++++
 .../objc-category-merging-complete-test.s     | 47 +++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/lld/MachO/ObjC.cpp b/lld/MachO/ObjC.cpp
index 15b89a808b05e..96ec646095be8 100644
--- a/lld/MachO/ObjC.cpp
+++ b/lld/MachO/ObjC.cpp
@@ -428,6 +428,7 @@ class ObjcCategoryMerger {
   static void doCleanup();
 
 private:
+  DenseSet<const Symbol *> collectNlCategories();
   void collectAndValidateCategoriesData();
   void
   mergeCategoriesIntoSingleCategory(std::vector<InfoInputCategory> &categories);
@@ -1060,7 +1061,27 @@ void ObjcCategoryMerger::createSymbolReference(Defined *refFrom,
   refFrom->isec()->relocs.push_back(r);
 }
 
+// Get the list of categories in the '__objc_nlcatlist' section. We can't
+// optimize these as they have a '+load' method that has to be called at
+// runtime.
+DenseSet<const Symbol *> ObjcCategoryMerger::collectNlCategories() {
+  DenseSet<const Symbol *> nlCategories;
+
+  for (InputSection *sec : allInputSections) {
+    if (sec->getName() != section_names::objcNonLazyCatList)
+      continue;
+
+    for (auto &r : sec->relocs) {
+      const Symbol *sym = r.referent.dyn_cast<Symbol *>();
+      nlCategories.insert(sym);
+    }
+  }
+  return nlCategories;
+}
+
 void ObjcCategoryMerger::collectAndValidateCategoriesData() {
+  auto nlCategories = collectNlCategories();
+
   for (InputSection *sec : allInputSections) {
     if (sec->getName() != section_names::objcCatList)
       continue;
@@ -1074,6 +1095,9 @@ void ObjcCategoryMerger::collectAndValidateCategoriesData() {
       assert(categorySym &&
              "Failed to get a valid category at __objc_catlit offset");
 
+      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))
diff --git a/lld/test/MachO/objc-category-merging-complete-test.s b/lld/test/MachO/objc-category-merging-complete-test.s
index d2d264a3f26c2..74400177b550d 100644
--- a/lld/test/MachO/objc-category-merging-complete-test.s
+++ b/lld/test/MachO/objc-category-merging-complete-test.s
@@ -88,6 +88,7 @@ MERGE_CATS-NEXT:                 name {{.*}} MyProtocol02Prop
 MERGE_CATS-NEXT:            attributes {{.*}} Ti,R,D
 MERGE_CATS-NEXT:                 name {{.*}} MyProtocol03Prop
 MERGE_CATS-NEXT:            attributes {{.*}} Ti,R,D
+MERGE_CATS:        __OBJC_$_CATEGORY_MyBaseClass_$_Category04
 
 
 NO_MERGE_CATS-NOT: __OBJC_$_CATEGORY_MyBaseClass(Category02|Category03)
@@ -431,6 +432,15 @@ L_OBJC_IMAGE_INFO:
 ## @dynamic MyProtocol03Prop;
 ## @end
 ##
+## // This category shouldn't be merged
+## @interface MyBaseClass(Category04)
+## + (void)load;
+## @end
+##
+## @implementation MyBaseClass(Category04)
+## + (void)load {}
+## @end
+##
 ## int main() {
 ##     return 0;
 ## }
@@ -493,6 +503,12 @@ L_OBJC_IMAGE_INFO:
 	b	_OUTLINED_FUNCTION_0
 	.cfi_endproc
                                         ; -- End function
+	.p2align	2
+"+[MyBaseClass(Category04) load]":
+	.cfi_startproc
+; %bb.0:
+	ret
+	.cfi_endproc
 	.globl	_main                           ; -- Begin function main
 	.p2align	2
 _main:                                  ; @main
@@ -746,11 +762,42 @@ __OBJC_$_CATEGORY_MyBaseClass_$_Category03:
 	.quad	0
 	.long	64                              ; 0x40
 	.space	4
+	.section	__TEXT,__objc_classname,cstring_literals
+l_OBJC_CLASS_NAME_.15:
+	.asciz	"Category04"
+	.section	__TEXT,__objc_methname,cstring_literals
+l_OBJC_METH_VAR_NAME_.16:
+	.asciz	"load"
+	.section	__DATA,__objc_const
+	.p2align	3, 0x0
+__OBJC_$_CATEGORY_CLASS_METHODS_MyBaseClass_$_Category04:
+	.long	24
+	.long	1
+	.quad	l_OBJC_METH_VAR_NAME_.16
+	.quad	l_OBJC_METH_VAR_TYPE_
+	.quad	"+[MyBaseClass(Category04) load]"
+	.p2align	3, 0x0
+__OBJC_$_CATEGORY_MyBaseClass_$_Category04:
+	.quad	l_OBJC_CLASS_NAME_.15
+	.quad	_OBJC_CLASS_$_MyBaseClass
+	.quad	0
+	.quad	__OBJC_$_CATEGORY_CLASS_METHODS_MyBaseClass_$_Category04
+	.quad	0
+	.quad	0
+	.quad	0
+	.long	64
+	.space	4
 	.section	__DATA,__objc_catlist,regular,no_dead_strip
 	.p2align	3, 0x0                          ; @"OBJC_LABEL_CATEGORY_$"
 l_OBJC_LABEL_CATEGORY_$:
 	.quad	__OBJC_$_CATEGORY_MyBaseClass_$_Category02
 	.quad	__OBJC_$_CATEGORY_MyBaseClass_$_Category03
+	.quad	__OBJC_$_CATEGORY_MyBaseClass_$_Category04
+	.section	__DATA,__objc_nlcatlist,regular,no_dead_strip
+	.p2align	3, 0x0
+l_OBJC_LABEL_NONLAZY_CATEGORY_$:
+	.quad	__OBJC_$_CATEGORY_MyBaseClass_$_Category04
+
 	.no_dead_strip	__OBJC_LABEL_PROTOCOL_$_MyProtocol02
 	.no_dead_strip	__OBJC_LABEL_PROTOCOL_$_MyProtocol03
 	.no_dead_strip	__OBJC_PROTOCOL_$_MyProtocol02