[LLD][ELF] --wrap: __real_foo references should trigger archive extraction for foo

A reference to __real_foo should trigger archive extraction of the input file that defines foo, otherwise a link using --wrap=foo might fail to link with an undefined reference to foo.
This matches bfd linker behaviour.

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

3 files changed, 116 insertions, 1 deletions

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index dd6a77a6da33..48990d00f44b 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -2243,9 +2243,16 @@ static std::vector<WrappedSymbol> addWrappedSymbols(opt::InputArgList &args) {
     if (!sym)
       continue;
 
-    Symbol *real = addUnusedUndefined(saver().save("__real_" + name));
     Symbol *wrap =
         addUnusedUndefined(saver().save("__wrap_" + name), sym->binding);
+
+    // If __real_ is referenced, pull in the symbol if it is lazy. Do this after
+    // processing __wrap_ as that may have referenced __real_.
+    StringRef realName = saver().save("__real_" + name);
+    if (symtab.find(realName))
+      addUnusedUndefined(name, sym->binding);
+
+    Symbol *real = addUnusedUndefined(realName);
     v.push_back({sym, real, wrap});
 
     // We want to tell LTO not to inline symbols to be overwritten
diff --git a/lld/test/ELF/wrap-extract-real.ll b/lld/test/ELF/wrap-extract-real.ll
new file mode 100644
index 000000000000..2ab6b3187ae3
--- /dev/null
+++ b/lld/test/ELF/wrap-extract-real.ll
@@ -0,0 +1,66 @@
+# REQUIRES: x86
+## --wrap=xxx should trigger archive extraction for symbol xxx for references to __real_xxx.
+
+# RUN: rm -rf %t && split-file %s %t && cd %t
+# RUN: llvm-as _start.ll -o _start.o
+# RUN: llvm-as _start_ref__real_foo.ll -o _start_ref__real_foo.o
+# RUN: llvm-as wrap.ll -o wrap.o
+# RUN: llvm-as foo.ll -o foo.o
+
+## Test when the reference to __real_foo is not in __wrap_foo.
+# RUN: ld.lld _start_ref__real_foo.o --start-lib foo.o --end-lib --wrap foo -o %t_real.elf
+# RUN: llvm-readelf --symbols %t_real.elf | FileCheck %s --check-prefix=REAL
+
+# REAL:      Symbol table '.symtab' contains 5 entries:
+# REAL-NEXT:        Value              Size Type    Bind   Vis       Ndx Name
+# REAL-NEXT:        {{.*}}           {{.*}} NOTYPE  LOCAL  DEFAULT   UND
+# REAL-NEXT:        {{.*}}           {{.*}} FILE    LOCAL  DEFAULT   ABS ld-temp.o
+# REAL-NEXT:        {{.*}}           {{.*}} FUNC    GLOBAL DEFAULT [[#]] _start
+# REAL-NEXT:        {{.*}}           {{.*}} FUNC    WEAK   DEFAULT [[#]] foo
+# REAL-NEXT:        {{.*}}           {{.*}} NOTYPE  GLOBAL DEFAULT   UND __wrap_foo
+
+## Test when the reference to __real_foo is in __wrap_foo.
+# RUN: ld.lld _start.o --start-lib wrap.o --end-lib --start-lib foo.o --end-lib --wrap foo -o %t_wrap_real.elf
+# RUN: llvm-readelf --symbols %t_wrap_real.elf | FileCheck %s --check-prefix=WRAP_REAL
+
+# WRAP_REAL:      Symbol table '.symtab' contains 5 entries:
+# WRAP_REAL-NEXT:    Value              Size Type    Bind   Vis       Ndx Name
+# WRAP_REAL-NEXT:    {{.*}}           {{.*}} NOTYPE  LOCAL  DEFAULT   UND
+# WRAP_REAL-NEXT:    {{.*}}           {{.*}} FILE    LOCAL  DEFAULT   ABS ld-temp.o
+# WRAP_REAL-NEXT:    {{.*}}           {{.*}} FUNC    GLOBAL DEFAULT [[#]] _start
+# WRAP_REAL-NEXT:    {{.*}}           {{.*}} FUNC    GLOBAL DEFAULT [[#]] __wrap_foo
+# WRAP_REAL-NEXT:    {{.*}}           {{.*}} FUNC    WEAK   DEFAULT [[#]] foo
+
+#--- _start.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-elf"
+define void @_start() {
+  ret void
+}
+
+#--- _start_ref__real_foo.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-elf"
+define void @_start() {
+  call void @__real_foo()
+  ret void
+}
+
+declare void @__real_foo()
+
+#--- wrap.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-elf"
+define void @__wrap_foo() {
+  call void @__real_foo()
+  ret void
+}
+
+declare void @__real_foo()
+
+#--- foo.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-elf"
+define void @foo() {
+  ret void
+}
diff --git a/lld/test/ELF/wrap-extract-real.s b/lld/test/ELF/wrap-extract-real.s
new file mode 100644
index 000000000000..52f98e5a82fc
--- /dev/null
+++ b/lld/test/ELF/wrap-extract-real.s
@@ -0,0 +1,42 @@
+# REQUIRES: x86
+## --wrap=xxx should trigger archive extraction for symbol xxx for references to __real_xxx.
+
+# RUN: rm -rf %t && split-file %s %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple=x86_64 _start.s -o _start.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64 ref__real_foo.s -o ref__real_foo.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64 wrap.s -o wrap.o
+# RUN: llvm-mc -filetype=obj -triple=x86_64 foo.s -o foo.o
+
+## Test when the reference to __real_foo is not in __wrap_foo.
+# RUN: ld.lld _start.o ref__real_foo.o --start-lib foo.o --end-lib --wrap foo -o %t_real.elf
+# RUN: llvm-readelf --symbols %t_real.elf | FileCheck %s --check-prefix=REAL
+
+# REAL:      Symbol table '.symtab' contains 4 entries:
+# REAL-NEXT:        Value             Size Type    Bind   Vis       Ndx Name
+# REAL-NEXT:        {{.*}}               0 NOTYPE  LOCAL  DEFAULT   UND
+# REAL-NEXT:        {{.*}}               0 NOTYPE  GLOBAL DEFAULT [[#]] _start
+# REAL-NEXT:        {{.*}}               0 NOTYPE  GLOBAL DEFAULT [[#]] foo
+# REAL-NEXT:        {{.*}}               0 NOTYPE  GLOBAL DEFAULT   UND __wrap_foo
+
+## Test when the reference to __real_foo is in __wrap_foo.
+# RUN: ld.lld _start.o --start-lib wrap.o --end-lib --start-lib foo.o --end-lib --wrap foo -o %t_wrap_real.elf
+# RUN: llvm-readelf --symbols %t_wrap_real.elf | FileCheck %s --check-prefix=WRAP_REAL
+
+# WRAP_REAL:      Symbol table '.symtab' contains 4 entries:
+# WRAP_REAL-NEXT:        Value             Size Type    Bind   Vis       Ndx Name
+# WRAP_REAL-NEXT:        {{.*}}               0 NOTYPE  LOCAL  DEFAULT   UND
+# WRAP_REAL-NEXT:        {{.*}}               0 NOTYPE  GLOBAL DEFAULT [[#]] _start
+# WRAP_REAL-NEXT:        {{.*}}               0 NOTYPE  GLOBAL DEFAULT [[#]] __wrap_foo
+# WRAP_REAL-NEXT:        {{.*}}               0 NOTYPE  GLOBAL DEFAULT [[#]] foo
+
+#--- _start.s
+.global _start; _start:; ret
+
+#--- ref__real_foo.s
+call __real_foo
+
+#--- wrap.s
+.global __wrap_foo; __wrap_foo:; call __real_foo
+
+#--- foo.s
+.global foo; foo:; ret