[lld] [lld-macho] dead-strip objc stubs (PR #79726)

[Kyungwoo Lee via llvm-commits]

https://github.com/kyulee-com updated https://github.com/llvm/llvm-project/pull/79726

>From b6a0e671056bd427e8eea34efad4d3fec0c94407 Mon Sep 17 00:00:00 2001
From: Kyungwoo Lee <kyulee at meta.com>
Date: Sat, 27 Jan 2024 14:25:21 -0800
Subject: [PATCH] [lld-macho] dead-strip objc stubs

---
 lld/MachO/Driver.cpp                   |  5 ++---
 lld/MachO/SyntheticSections.cpp        | 14 +++++++++++--
 lld/MachO/SyntheticSections.h          |  2 ++
 lld/MachO/Writer.cpp                   | 10 +++++++++-
 lld/test/MachO/arm64-objc-stubs-dead.s | 27 ++++++++++++++++++++++++++
 5 files changed, 52 insertions(+), 6 deletions(-)
 create mode 100644 lld/test/MachO/arm64-objc-stubs-dead.s

diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
index 411fbcfcf233e..a57f60c5eed36 100644
--- a/lld/MachO/Driver.cpp
+++ b/lld/MachO/Driver.cpp
@@ -1275,11 +1275,10 @@ static void foldIdenticalLiterals() {
 static void addSynthenticMethnames() {
   std::string &data = *make<std::string>();
   llvm::raw_string_ostream os(data);
-  const int prefixLength = ObjCStubsSection::symbolPrefix.size();
   for (Symbol *sym : symtab->getSymbols())
     if (isa<Undefined>(sym))
-      if (sym->getName().starts_with(ObjCStubsSection::symbolPrefix))
-        os << sym->getName().drop_front(prefixLength) << '\0';
+      if (ObjCStubsSection::isObjCStubSymbol(sym))
+        os << ObjCStubsSection::getMethname(sym) << '\0';
 
   if (data.empty())
     return;
diff --git a/lld/MachO/SyntheticSections.cpp b/lld/MachO/SyntheticSections.cpp
index 53220ad04b842..d3c8cb02942bf 100644
--- a/lld/MachO/SyntheticSections.cpp
+++ b/lld/MachO/SyntheticSections.cpp
@@ -814,9 +814,19 @@ ObjCStubsSection::ObjCStubsSection()
               : target->objcStubsSmallAlignment;
 }
 
+bool ObjCStubsSection::isObjCStubSymbol(Symbol *sym) {
+  return sym->getName().starts_with(symbolPrefix);
+}
+
+StringRef ObjCStubsSection::getMethname(Symbol *sym) {
+  assert(isObjCStubSymbol(sym) && "not an objc stub");
+  auto name = sym->getName();
+  StringRef methname = name.drop_front(symbolPrefix.size());
+  return methname;
+}
+
 void ObjCStubsSection::addEntry(Symbol *sym) {
-  assert(sym->getName().starts_with(symbolPrefix) && "not an objc stub");
-  StringRef methname = sym->getName().drop_front(symbolPrefix.size());
+  StringRef methname = getMethname(sym);
   offsets.push_back(
       in.objcMethnameSection->getStringOffset(methname).outSecOff);
 
diff --git a/lld/MachO/SyntheticSections.h b/lld/MachO/SyntheticSections.h
index 5fb7b6e09e8e6..5ae97954ab3fe 100644
--- a/lld/MachO/SyntheticSections.h
+++ b/lld/MachO/SyntheticSections.h
@@ -332,6 +332,8 @@ class ObjCStubsSection final : public SyntheticSection {
   void setUp();
 
   static constexpr llvm::StringLiteral symbolPrefix = "_objc_msgSend$";
+  static bool isObjCStubSymbol(Symbol *sym);
+  static StringRef getMethname(Symbol *sym);
 
 private:
   std::vector<Defined *> symbols;
diff --git a/lld/MachO/Writer.cpp b/lld/MachO/Writer.cpp
index 1b0e64abe843a..3634c626f0692 100644
--- a/lld/MachO/Writer.cpp
+++ b/lld/MachO/Writer.cpp
@@ -736,8 +736,16 @@ void Writer::scanSymbols() {
       dysym->getFile()->refState =
           std::max(dysym->getFile()->refState, dysym->getRefState());
     } else if (isa<Undefined>(sym)) {
-      if (sym->getName().starts_with(ObjCStubsSection::symbolPrefix))
+      if (ObjCStubsSection::isObjCStubSymbol(sym)) {
+        // When -dead_strip is enabled, we don't want to emit any dead stubs.
+        // Although this stub symbol is yet undefined, addSym() was called
+        // during MarkLive.
+        if (config->deadStrip) {
+          if (!sym->isLive())
+            continue;
+        }
         in.objcStubs->addEntry(sym);
+      }
     }
   }
 
diff --git a/lld/test/MachO/arm64-objc-stubs-dead.s b/lld/test/MachO/arm64-objc-stubs-dead.s
new file mode 100644
index 0000000000000..5dcb171c17eac
--- /dev/null
+++ b/lld/test/MachO/arm64-objc-stubs-dead.s
@@ -0,0 +1,27 @@
+# REQUIRES: aarch64
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %s -o %t.o
+
+# RUN: %lld -arch arm64 -lSystem -U _objc_msgSend -o %t.out %t.o
+# RUN: llvm-nm %t.out | FileCheck %s
+# RUN: %lld -arch arm64 -lSystem -U _objc_msgSend -dead_strip -o %t.out %t.o
+# RUN: llvm-nm %t.out | FileCheck %s --check-prefix=DEAD
+
+# CHECK: _foo
+# CHECK: _objc_msgSend$length
+
+# DEAD-NOT: _foo
+# DEAD-NOT: _objc_msgSend$length
+
+.section __TEXT,__text
+
+.globl _foo
+_foo:
+  bl  _objc_msgSend$length
+  ret
+
+.globl _main
+_main:
+  ret
+
+.subsections_via_symbols