[lld-macho] Use LC_LOAD_WEAK_DYLIB for dylibs with only weakrefs

Note that dylibs without *any* refs will still be loaded in the usual
(strong) fashion.

Reviewed By: #lld-macho, thakis

Differential Revision: https://reviews.llvm.org/D93435

lld/MachO/InputFiles.cpp

@@ -527,7 +527,7 @@ void loadReexport(StringRef path, DylibFile *umbrella) {
 }
 
 DylibFile::DylibFile(MemoryBufferRef mb, DylibFile *umbrella)
-    : InputFile(DylibKind, mb) {
+    : InputFile(DylibKind, mb), refState(RefState::Unreferenced) {
   if (umbrella == nullptr)
     umbrella = this;
 
@@ -580,7 +580,7 @@ DylibFile::DylibFile(MemoryBufferRef mb, DylibFile *umbrella)
 }
 
 DylibFile::DylibFile(const InterfaceFile &interface, DylibFile *umbrella)
-    : InputFile(DylibKind, interface) {
+    : InputFile(DylibKind, interface), refState(RefState::Unreferenced) {
   if (umbrella == nullptr)
     umbrella = this;
 

lld/MachO/InputFiles.h

@@ -38,6 +38,7 @@ namespace macho {
 class InputSection;
 class Symbol;
 struct Reloc;
+enum class RefState : uint8_t;
 
 // If --reproduce option is given, all input files are written
 // to this tar archive.
@@ -135,6 +136,7 @@ public:
   uint32_t compatibilityVersion = 0;
   uint32_t currentVersion = 0;
   uint64_t ordinal = 0; // Ordinal numbering starts from 1, so 0 is a sentinel
+  RefState refState;
   bool reexport = false;
   bool forceWeakImport = false;
 };

lld/MachO/Symbols.h

@@ -116,7 +116,11 @@ private:
   const bool external : 1;
 };
 
-// Indicates whether & how a dylib symbol is referenced.
+// This enum does double-duty: as a symbol property, it indicates whether & how
+// a dylib symbol is referenced. As a DylibFile property, it indicates the kind
+// of referenced symbols contained within the file. If there are both weak
+// and strong references to the same file, we will count the file as
+// strongly-referenced.
 enum class RefState : uint8_t { Unreferenced = 0, Weak = 1, Strong = 2 };
 
 class Undefined : public Symbol {

lld/MachO/Writer.cpp

@@ -43,6 +43,7 @@ public:
   Writer() : buffer(errorHandler().outputBuffer) {}
 
   void scanRelocations();
+  void scanSymbols();
   void createOutputSections();
   void createLoadCommands();
   void assignAddresses(OutputSegment *);
@@ -424,6 +425,17 @@ void Writer::scanRelocations() {
   }
 }
 
+void Writer::scanSymbols() {
+  for (const macho::Symbol *sym : symtab->getSymbols()) {
+    if (const auto *defined = dyn_cast<Defined>(sym)) {
+      if (defined->overridesWeakDef)
+        in.weakBinding->addNonWeakDefinition(defined);
+    } else if (const auto *dysym = dyn_cast<DylibSymbol>(sym)) {
+      dysym->file->refState = std::max(dysym->file->refState, dysym->refState);
+    }
+  }
+}
+
 void Writer::createLoadCommands() {
   in.header->addLoadCommand(make<LCDyldInfo>(
       in.rebase, in.binding, in.weakBinding, in.lazyBinding, in.exports));
@@ -463,10 +475,10 @@ void Writer::createLoadCommands() {
   uint64_t dylibOrdinal = 1;
   for (InputFile *file : inputFiles) {
     if (auto *dylibFile = dyn_cast<DylibFile>(file)) {
-      // TODO: dylibs that are only referenced by weak refs should also be
-      // loaded via LC_LOAD_WEAK_DYLIB.
       LoadCommandType lcType =
-          dylibFile->forceWeakImport ? LC_LOAD_WEAK_DYLIB : LC_LOAD_DYLIB;
+          dylibFile->forceWeakImport || dylibFile->refState == RefState::Weak
+              ? LC_LOAD_WEAK_DYLIB
+              : LC_LOAD_DYLIB;
       in.header->addLoadCommand(make<LCDylib>(lcType, dylibFile->dylibName,
                                               dylibFile->compatibilityVersion,
                                               dylibFile->currentVersion));
@@ -699,11 +711,7 @@ void Writer::run() {
   scanRelocations();
   if (in.stubHelper->isNeeded())
     in.stubHelper->setup();
-
-  for (const macho::Symbol *sym : symtab->getSymbols())
-    if (const auto *defined = dyn_cast<Defined>(sym))
-      if (defined->overridesWeakDef)
-        in.weakBinding->addNonWeakDefinition(defined);
+  scanSymbols();
 
   // Sort and assign sections to their respective segments. No more sections nor
   // segments may be created after these methods run.

lld/test/MachO/weak-import.s

@@ -1,35 +1,72 @@
 # REQUIRES: x86
 # RUN: split-file %s %t
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/test.s -o %t/test.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/weak-ref-only.s -o %t/weak-ref-only.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/weak-ref-sub-library.s -o %t/weak-ref-sub-library.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/mixed-ref.s -o %t/mixed-ref.o
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/foo.s -o %t/foo.o
-# RUN: %lld -lSystem -dylib %t/foo.o -o %t/libfoo.dylib
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-darwin %t/bar.s -o %t/bar.o
+# RUN: %lld -lSystem -dylib %t/bar.o -o %t/libbar.dylib
+# RUN: %lld -lSystem -dylib %t/foo.o %t/libbar.dylib -sub_library libbar -o %t/libfoo.dylib
 
 # RUN: %lld -weak-lSystem %t/test.o -weak_framework CoreFoundation -weak_library %t/libfoo.dylib -o %t/test
-# RUN: llvm-objdump --macho --all-headers %t/test | FileCheck %s -DDIR=%t
+# RUN: llvm-objdump --macho --all-headers %t/test | FileCheck %s -DDIR=%t --check-prefixes=WEAK-SYS,WEAK-FOO
 # RUN: %lld -weak-lSystem %t/test.o \
 # RUN:   -framework CoreFoundation -weak_framework CoreFoundation -framework CoreFoundation \
 # RUN:   %t/libfoo.dylib -weak_library %t/libfoo.dylib %t/libfoo.dylib -o %t/test
-# RUN: llvm-objdump --macho --all-headers %t/test | FileCheck %s -DDIR=%t
+# RUN: llvm-objdump --macho --all-headers %t/test | FileCheck %s -DDIR=%t --check-prefixes=WEAK-SYS,WEAK-FOO
+# RUN: %lld -lSystem -dylib %t/libfoo.dylib %t/weak-ref-only.o -o %t/weak-ref-only
+# RUN: llvm-objdump --macho --all-headers %t/weak-ref-only | FileCheck %s -DDIR=%t --check-prefixes=SYS,WEAK-FOO
+# RUN: %lld -lSystem -dylib %t/libfoo.dylib %t/weak-ref-sub-library.o -o %t/weak-ref-sub-library
+# RUN: llvm-objdump --macho --all-headers %t/weak-ref-sub-library | FileCheck %s -DDIR=%t --check-prefixes=SYS,WEAK-FOO
+# RUN: %lld -lSystem -dylib %t/libfoo.dylib %t/mixed-ref.o -o %t/mixed-ref
+# RUN: llvm-objdump --macho --all-headers %t/mixed-ref | FileCheck %s -DDIR=%t --check-prefixes=SYS,FOO
 
-# CHECK:          cmd LC_LOAD_WEAK_DYLIB
-# CHECK-NEXT: cmdsize
-# CHECK-NEXT:    name /usr/lib/libSystem.B.dylib
+# WEAK-SYS:          cmd LC_LOAD_WEAK_DYLIB
+# WEAK-SYS-NEXT: cmdsize
+# WEAK-SYS-NEXT:    name /usr/lib/libSystem.B.dylib
 
-# CHECK:          cmd LC_LOAD_WEAK_DYLIB
-# CHECK-NEXT: cmdsize
-# CHECK-NEXT:    name /System/Library/Frameworks/CoreFoundation.framework/CoreFoundation
+# WEAK-SYS:          cmd LC_LOAD_WEAK_DYLIB
+# WEAK-SYS-NEXT: cmdsize
+# WEAK-SYS-NEXT:    name /System/Library/Frameworks/CoreFoundation.framework/CoreFoundation
 
-# CHECK:          cmd LC_LOAD_WEAK_DYLIB
-# CHECK-NEXT: cmdsize
-# CHECK-NEXT:    name [[DIR]]/libfoo.dylib
+# SYS:               cmd LC_LOAD_DYLIB
+# SYS-NEXT:      cmdsize
+# SYS-NEXT:         name /usr/lib/libSystem.B.dylib
+
+# WEAK-FOO:          cmd LC_LOAD_WEAK_DYLIB
+# WEAK-FOO-NEXT: cmdsize
+# WEAK-FOO-NEXT:    name [[DIR]]/libfoo.dylib
+
+# FOO:               cmd LC_LOAD_DYLIB
+# FOO-NEXT:      cmdsize
+# FOO-NEXT:         name [[DIR]]/libfoo.dylib
 
 #--- foo.s
 .globl _foo
 _foo:
-  ret
+
+#--- bar.s
+.globl _bar
+_bar:
+
+#--- weak-ref-only.s
+.weak_reference _foo
+.data
+.quad _foo
+
+#--- weak-ref-sub-library.s
+.weak_reference _bar
+.data
+.quad _bar
+
+#--- mixed-ref.s
+.weak_definition _foo
+.data
+.quad _foo
+.quad _bar
 
 #--- test.s
 .globl _main
-.text
 _main:
   ret