[lld/mac] Handle symbols from -U in treatUndefinedSymbol()

In ld64, `-U section$start$FOO$bar` handles `section$start$FOO$bar`
as a regular `section$start` symbol, that is section$start processing
happens before -U processing.

Likely, nobody uses that in practice so it doesn't seem very important
to be compatible with this, but it also moves the -U handling code next
to the `-undefined dynamic_lookup` handling code, which is nice because
they do the same thing. And, in fact, this did identify a bug in a corner
case in the intersection of `-undefined dynamic_lookup` and dead-stripping
(fix for that in D106565).

Vaguely related to PR50760.

No interesting behavior change.

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

Author: nico

lld/MachO/Config.h

@@ -13,6 +13,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/BinaryFormat/MachO.h"
 #include "llvm/Support/CachePruning.h"
 #include "llvm/Support/GlobPattern.h"
@@ -158,6 +159,7 @@ struct Configuration {
   std::vector<llvm::StringRef> runtimePaths;
   std::vector<std::string> astPaths;
   std::vector<Symbol *> explicitUndefineds;
+  llvm::StringSet<> explicitDynamicLookups;
   // There are typically few custom sectionAlignments or segmentProtections,
   // so use a vector instead of a map.
   std::vector<SectionAlign> sectionAlignments;

lld/MachO/Driver.cpp

@@ -1175,7 +1175,7 @@ bool macho::link(ArrayRef<const char *> argsArr, bool canExitEarly,
   }
 
   for (const Arg *arg : args.filtered(OPT_U))
-    symtab->addDynamicLookup(arg->getValue());
+    config->explicitDynamicLookups.insert(arg->getValue());
 
   config->mapFile = args.getLastArgValue(OPT_map);
   config->optimize = args::getInteger(args, OPT_O, 1);

lld/MachO/SymbolTable.cpp

@@ -197,6 +197,13 @@ Defined *SymbolTable::addSynthetic(StringRef name, InputSection *isec,
 }
 
 void lld::macho::treatUndefinedSymbol(const Undefined &sym, StringRef source) {
+  // Handle -U.
+  if (config->explicitDynamicLookups.count(sym.getName())) {
+    symtab->addDynamicLookup(sym.getName());
+    return;
+  }
+
+  // Handle -undefined.
   auto message = [source, &sym]() {
     std::string message = "undefined symbol";
     if (config->archMultiple)

lld/test/MachO/dead-strip.s

@@ -147,9 +147,9 @@
 # STRIPDYLIB-NEXT:  l {{.*}} __dyld_private
 # STRIPDYLIB-NEXT:  g {{.*}} _main
 # STRIPDYLIB-NEXT:  g {{.*}} __mh_execute_header
-# STRIPDYLIB-NEXT:  *UND* _ref_undef_fun
 # STRIPDYLIB-NEXT:  *UND* dyld_stub_binder
 # STRIPDYLIB-NEXT:  *UND* _ref_dylib_fun
+# STRIPDYLIB-NEXT:  *UND* _ref_undef_fun
 # STRIPDYLIB:      Bind table:
 # STRIPDYLIB:      Lazy bind table:
 # STRIPDYLIB:       __DATA   __la_symbol_ptr {{.*}} flat-namespace _ref_undef_fun
@@ -202,9 +202,9 @@
 # LIVESUPP-NEXT:   g {{.*}} _bar
 # LIVESUPP-NEXT:   g {{.*}} _foo
 # LIVESUPP-NEXT:   g {{.*}} __mh_execute_header
-# LIVESUPP-NEXT:   *UND* _ref_undef_fun
 # LIVESUPP-NEXT:   *UND* dyld_stub_binder
 # LIVESUPP-NEXT:   *UND* _ref_dylib_fun
+# LIVESUPP-NEXT:   *UND* _ref_undef_fun
 
 # RUN: llvm-mc -filetype=obj -triple=x86_64-apple-macos \
 # RUN:     %t/live-support-iterations.s -o %t/live-support-iterations.o