[lld] [lld-macho][ObjC] Implement category merging into base class (PR #92448)

[Ellis Hoag via llvm-commits]

================
@@ -1302,3 +1340,100 @@ void objc::mergeCategories() {
 }
 
 void objc::doCleanup() { ObjcCategoryMerger::doCleanup(); }
+
+void ObjcCategoryMerger::mergeCategoriesIntoBaseClass(
+    const Defined *baseClass, std::vector<InfoInputCategory> &categories) {
+  assert(categories.size() >= 1 && "Expected at least one category to merge");
+
+  // Collect all the info from the categories
+  ClassExtensionInfo extInfo(catLayout);
+  for (auto &catInfo : categories) {
+    parseCatInfoToExtInfo(catInfo, extInfo);
+  }
+
+  // Get metadata for the base class
+  Defined *metaRo = getClassRo(baseClass, /*getMetaRo=*/true);
+  ConcatInputSection *metaIsec = dyn_cast<ConcatInputSection>(metaRo->isec());
+  Defined *classRo = getClassRo(baseClass, /*getMetaRo=*/false);
+  ConcatInputSection *classIsec = dyn_cast<ConcatInputSection>(classRo->isec());
+
+  // Now collect the info from the base class from the various lists in the
+  // class metadata
+  parseProtocolListInfo(classIsec, roClassLayout.baseProtocolsOffset,
+                        extInfo.protocols);
+
+  parsePointerListInfo(metaIsec, roClassLayout.baseMethodsOffset,
+                       extInfo.classMethods);
+
+  parsePointerListInfo(metaIsec, roClassLayout.basePropertiesOffset,
+                       extInfo.classProps);
+
+  parsePointerListInfo(classIsec, roClassLayout.baseMethodsOffset,
+                       extInfo.instanceMethods);
+
+  parsePointerListInfo(classIsec, roClassLayout.basePropertiesOffset,
+                       extInfo.instanceProps);
+
+  // Erase the old lists - these will be generated and replaced
+  eraseSymbolAtIsecOffset(metaIsec, roClassLayout.baseMethodsOffset);
+  eraseSymbolAtIsecOffset(metaIsec, roClassLayout.baseProtocolsOffset);
+  eraseSymbolAtIsecOffset(metaIsec, roClassLayout.basePropertiesOffset);
+  eraseSymbolAtIsecOffset(classIsec, roClassLayout.baseMethodsOffset);
+  eraseSymbolAtIsecOffset(classIsec, roClassLayout.baseProtocolsOffset);
+  eraseSymbolAtIsecOffset(classIsec, roClassLayout.basePropertiesOffset);
+
+  // Emit the newly merged lists - first into the meta RO then into the class RO
+  emitAndLinkPointerList(metaRo, roClassLayout.baseMethodsOffset, extInfo,
+                         extInfo.classMethods);
+
+  // Protocols are a special case - the single list is referenced by both the
+  // class RO and meta RO. Here we emit it and link it into the meta RO
+  Defined *protoListSym = emitAndLinkProtocolList(
+      metaRo, roClassLayout.baseProtocolsOffset, extInfo, extInfo.protocols);
+
+  emitAndLinkPointerList(metaRo, roClassLayout.basePropertiesOffset, extInfo,
+                         extInfo.classProps);
+
+  emitAndLinkPointerList(classRo, roClassLayout.baseMethodsOffset, extInfo,
+                         extInfo.instanceMethods);
+
+  // If we emitted a new protocol list, link it to the class RO also
+  if (protoListSym) {
+    createSymbolReference(classRo, protoListSym,
+                          roClassLayout.baseProtocolsOffset,
+                          infoCategoryWriter.catBodyInfo.relocTemplate);
+  }
+
+  emitAndLinkPointerList(classRo, roClassLayout.basePropertiesOffset, extInfo,
+                         extInfo.instanceProps);
+
+  // Mark all the categories as merged - this will be used to erase them later
+  for (auto &catInfo : categories)
+    catInfo.wasMerged = true;
+}
+
+// Erase the symbol at a given offset in an InputSection
+void ObjcCategoryMerger::eraseSymbolAtIsecOffset(ConcatInputSection *isec,
+                                                 uint32_t offset) {
+  Defined *sym = tryGetDefinedAtIsecOffset(isec, offset);
+  if (!sym)
+    return;
+
+  // Remove the symbol from isec->symbols
+  assert(isa<Defined>(sym) && "Can only erase a Defined");
+  isec->symbols.erase(
+      std::remove(isec->symbols.begin(), isec->symbols.end(), sym),
+      isec->symbols.end());
----------------
ellishg wrote:

You might be able to use `llvm::erase_if()`/`llvm::erase()` here.

https://github.com/llvm/llvm-project/pull/92448