[lld-macho] Support re-exports of individual symbols

Specifically, we support this:

  ld64.lld -dylib foo.o libbar.dylib -exported_symbol _bar -o libfoo.dylib

Where `_bar` is defined in libbar.dylib.

Reviewed By: #lld-macho, smeenai

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

4 files changed, 65 insertions, 16 deletions

diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
index 0f2326b305b1..94c9b49268ac 100644
--- a/lld/MachO/Driver.cpp
+++ b/lld/MachO/Driver.cpp
@@ -1331,8 +1331,7 @@ static void handleExplicitExports() {
   if (config->hasExplicitExports) {
     parallelForEach(symtab->getSymbols(), [](Symbol *sym) {
       if (auto *defined = dyn_cast<Defined>(sym)) {
-        StringRef symbolName = defined->getName();
-        if (config->exportedSymbols.match(symbolName)) {
+        if (config->exportedSymbols.match(sym->getName())) {
           if (defined->privateExtern) {
             if (defined->weakDefCanBeHidden) {
               // weak_def_can_be_hidden symbols behave similarly to
@@ -1348,6 +1347,8 @@ static void handleExplicitExports() {
         } else {
           defined->privateExtern = true;
         }
+      } else if (auto *dysym = dyn_cast<DylibSymbol>(sym)) {
+        dysym->shouldReexport = config->exportedSymbols.match(sym->getName());
       }
     });
   } else if (!config->unexportedSymbols.empty()) {
diff --git a/lld/MachO/ExportTrie.cpp b/lld/MachO/ExportTrie.cpp
index 3ca8d350be14..e5fff36c4b6d 100644
--- a/lld/MachO/ExportTrie.cpp
+++ b/lld/MachO/ExportTrie.cpp
@@ -58,21 +58,24 @@ struct Edge {
 
 struct ExportInfo {
   uint64_t address;
+  uint64_t ordinal = 0;
   uint8_t flags = 0;
   ExportInfo(const Symbol &sym, uint64_t imageBase)
       : address(sym.getVA() - imageBase) {
     using namespace llvm::MachO;
-    // Set the symbol type.
     if (sym.isWeakDef())
       flags |= EXPORT_SYMBOL_FLAGS_WEAK_DEFINITION;
-    // TODO: Add proper support for re-exports & stub-and-resolver flags.
-
-    // Set the symbol kind.
-    if (sym.isTlv()) {
+    if (sym.isTlv())
       flags |= EXPORT_SYMBOL_FLAGS_KIND_THREAD_LOCAL;
-    } else if (auto *defined = dyn_cast<Defined>(&sym)) {
+    // TODO: Add proper support for stub-and-resolver flags.
+
+    if (auto *defined = dyn_cast<Defined>(&sym)) {
       if (defined->isAbsolute())
         flags |= EXPORT_SYMBOL_FLAGS_KIND_ABSOLUTE;
+    } else if (auto *dysym = dyn_cast<DylibSymbol>(&sym)) {
+      flags |= EXPORT_SYMBOL_FLAGS_REEXPORT;
+      if (!dysym->isDynamicLookup())
+        ordinal = dysym->getFile()->ordinal;
     }
   }
 };
@@ -87,19 +90,57 @@ struct macho::TrieNode {
   // fixpoint.
   size_t offset = 0;
 
+  uint32_t getTerminalSize() const;
   // Returns whether the new estimated offset differs from the old one.
   bool updateOffset(size_t &nextOffset);
   void writeTo(uint8_t *buf) const;
 };
 
+// For regular symbols, the node layout (excluding the children) is
+//
+//   uleb128 terminalSize;
+//   uleb128 flags;
+//   uleb128 address;
+//
+// For re-exported symbols, the layout is
+//
+//   uleb128 terminalSize;
+//   uleb128 flags;
+//   uleb128 ordinal;
+//   char[] originalName;
+//
+// If libfoo.dylib is linked against libbar.dylib, and libfoo exports an alias
+// _foo to a symbol _bar in libbar, then originalName will be "_bar". If libfoo
+// re-exports _bar directly (i.e. not via an alias), then originalName will be
+// the empty string.
+//
+// TODO: Support aliased re-exports. (Since we don't yet support these,
+// originalName will always be the empty string.)
+//
+// For stub-and-resolver nodes, the layout is
+//
+//   uleb128 terminalSize;
+//   uleb128 flags;
+//   uleb128 stubAddress;
+//   uleb128 resolverAddress;
+//
+// TODO: Support stub-and-resolver nodes.
+uint32_t TrieNode::getTerminalSize() const {
+  uint32_t size = getULEB128Size(info->flags);
+  if (info->flags & MachO::EXPORT_SYMBOL_FLAGS_REEXPORT)
+    size += getULEB128Size(info->ordinal) + 1; // + 1 for the null-terminator
+  else
+    size += getULEB128Size(info->address);
+  return size;
+}
+
 bool TrieNode::updateOffset(size_t &nextOffset) {
   // Size of the whole node (including the terminalSize and the outgoing edges.)
   // In contrast, terminalSize only records the size of the other data in the
   // node.
   size_t nodeSize;
   if (info) {
-    uint32_t terminalSize =
-        getULEB128Size(info->flags) + getULEB128Size(info->address);
+    uint32_t terminalSize = getTerminalSize();
     // Overall node size so far is the uleb128 size of the length of the symbol
     // info + the symbol info itself.
     nodeSize = terminalSize + getULEB128Size(terminalSize);
@@ -123,12 +164,15 @@ bool TrieNode::updateOffset(size_t &nextOffset) {
 void TrieNode::writeTo(uint8_t *buf) const {
   buf += offset;
   if (info) {
-    // TrieNodes with Symbol info: size, flags address
-    uint32_t terminalSize =
-        getULEB128Size(info->flags) + getULEB128Size(info->address);
+    uint32_t terminalSize = getTerminalSize();
     buf += encodeULEB128(terminalSize, buf);
     buf += encodeULEB128(info->flags, buf);
-    buf += encodeULEB128(info->address, buf);
+    if (info->flags & MachO::EXPORT_SYMBOL_FLAGS_REEXPORT) {
+      buf += encodeULEB128(info->ordinal, buf);
+      *buf++ = 0; // empty originalName string
+    } else {
+      buf += encodeULEB128(info->address, buf);
+    }
   } else {
     // TrieNode with no Symbol info.
     *buf++ = 0; // terminalSize
diff --git a/lld/MachO/Symbols.h b/lld/MachO/Symbols.h
index d8e2f86e525b..c30c4bc16ee0 100644
--- a/lld/MachO/Symbols.h
+++ b/lld/MachO/Symbols.h
@@ -254,8 +254,8 @@ class DylibSymbol : public Symbol {
 public:
   DylibSymbol(DylibFile *file, StringRefZ name, bool isWeakDef,
               RefState refState, bool isTlv)
-      : Symbol(DylibKind, name, file), refState(refState), weakDef(isWeakDef),
-        tlv(isTlv) {
+      : Symbol(DylibKind, name, file), shouldReexport(false),
+        refState(refState), weakDef(isWeakDef), tlv(isTlv) {
     if (file && refState > RefState::Unreferenced)
       file->numReferencedSymbols++;
   }
@@ -297,6 +297,7 @@ public:
     }
   }
 
+  bool shouldReexport : 1;
 private:
   RefState refState : 2;
   const bool weakDef : 1;
diff --git a/lld/MachO/SyntheticSections.cpp b/lld/MachO/SyntheticSections.cpp
index 95da58ca23c2..72ec67ad1e30 100644
--- a/lld/MachO/SyntheticSections.cpp
+++ b/lld/MachO/SyntheticSections.cpp
@@ -978,6 +978,9 @@ void ExportSection::finalizeContents() {
         continue;
       trieBuilder.addSymbol(*defined);
       hasWeakSymbol = hasWeakSymbol || sym->isWeakDef();
+    } else if (auto *dysym = dyn_cast<DylibSymbol>(sym)) {
+      if (dysym->shouldReexport)
+        trieBuilder.addSymbol(*dysym);
     }
   }
   size = trieBuilder.build();