[ELF][RISCV] --wrap=foo: Correctly update st_value(foo)

With --wrap=foo, we may have `d->file != file` for a defined symbol `foo`.
For the object file defining `foo`, its symbol table may not contain
`foo` after `redirectSymbols` changed the `foo` entry to `__wrap_foo` (see D50569).

Therefore, skipping `foo` with the condition `if (!d || d->file != file)` may
cause `__wrap_foo` not to be updated. See `ab.o w.o --wrap=foo` in the new test
(originally reported by D150220).

We could adjust the condition to `if (!d)`, but that would leave many `anchors`
entries if a symbol is referenced by many files. Switch to iterating over
`symtab` instead.

Note: D149735 (actually not NFC) allowed duplicate `anchors` entries and fixed
`a.o bw.o --wrap=foo`.

Reviewed By: jobnoorman

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

2 files changed, 75 insertions, 1 deletions

diff --git a/lld/ELF/Arch/RISCV.cpp b/lld/ELF/Arch/RISCV.cpp
index 018da8a..d0d7511 100644
--- a/lld/ELF/Arch/RISCV.cpp
+++ b/lld/ELF/Arch/RISCV.cpp
@@ -550,10 +550,18 @@ static void initSymbolAnchors() {
   }
   // Store anchors (st_value and st_value+st_size) for symbols relative to text
   // sections.
+  //
+  // For a defined symbol foo, we may have `d->file != file` with --wrap=foo.
+  // We should process foo, as the defining object file's symbol table may not
+  // contain foo after redirectSymbols changed the foo entry to __wrap_foo. To
+  // avoid adding a Defined that is undefined in one object file, use
+  // `!d->scriptDefined` to exclude symbols that are definitely not wrapped.
+  //
+  // `relaxAux->anchors` may contain duplicate symbols, but that is fine.
   for (InputFile *file : ctx.objectFiles)
     for (Symbol *sym : file->getSymbols()) {
       auto *d = dyn_cast<Defined>(sym);
-      if (!d || d->file != file)
+      if (!d || (d->file != file && !d->scriptDefined))
         continue;
       if (auto *sec = dyn_cast_or_null<InputSection>(d->section))
         if (sec->flags & SHF_EXECINSTR && sec->relaxAux) {
diff --git a/lld/test/ELF/riscv-relax-wrap.s b/lld/test/ELF/riscv-relax-wrap.s
new file mode 100644
index 0000000..6b507de
--- /dev/null
+++ b/lld/test/ELF/riscv-relax-wrap.s
@@ -0,0 +1,66 @@
+# REQUIRES: riscv
+## Don't forget to update st_value(foo) when foo is defined in another relocatable object file.
+
+# RUN: rm -rf %t && split-file %s %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple=riscv64 -mattr=+relax a.s -o a.o
+# RUN: llvm-mc -filetype=obj -triple=riscv64 -mattr=+relax b.s -o b.o
+# RUN: llvm-mc -filetype=obj -triple=riscv64 -mattr=+relax w.s -o w.o
+# RUN: llvm-mc -filetype=obj -triple=riscv64 -mattr=+relax w2.s -o w2.o
+# RUN: llvm-mc -filetype=obj -triple=riscv64 -mattr=+relax call_foo.s -o call_foo.o
+
+# RUN: ld.lld -r b.o w.o -o bw.o
+# RUN: ld.lld -Ttext=0x10000 a.o bw.o --wrap=foo -o 1
+# RUN: llvm-objdump -d --no-show-raw-insn 1 | FileCheck %s
+
+# RUN: ld.lld -r a.o b.o -o ab.o
+# RUN: ld.lld -Ttext=0x10000 ab.o w.o --wrap=foo -o 2
+# RUN: llvm-objdump -d --no-show-raw-insn 2 | FileCheck %s
+
+# RUN: ld.lld -Ttext=0x10000 w2.o call_foo.o --wrap=foo -o 3
+# RUN: llvm-objdump -d --no-show-raw-insn 3 | FileCheck %s --check-prefix=CHECK2
+
+# CHECK-LABEL:  <_start>:
+# CHECK-NEXT:     10000: jal {{.*}} <__wrap_foo>
+# CHECK-EMPTY:
+# CHECK-NEXT:   <foo>:
+# CHECK-NEXT:     10004: jal {{.*}} <__wrap_foo>
+# CHECK-EMPTY:
+# CHECK-NEXT:   <__wrap_foo>:
+# CHECK-NEXT:     10008: jal {{.*}} <foo>
+
+# CHECK2-LABEL: <_start>:
+# CHECK2-NEXT:    jal {{.*}} <call_foo>
+# CHECK2-EMPTY:
+# CHECK2-NEXT:  <__wrap_foo>:
+# CHECK2-NEXT:    ret
+# CHECK2-EMPTY:
+# CHECK2-NEXT:  <call_foo>:
+# CHECK2-NEXT:    jal {{.*}} <__wrap_foo>
+
+#--- a.s
+.globl _start
+_start:
+  call foo
+
+#--- b.s
+.globl foo
+foo:
+  call __wrap_foo
+
+#--- w.s
+.globl __wrap_foo
+__wrap_foo:
+  call __real_foo
+
+#--- w2.s
+.globl _start, __wrap_foo
+_start:
+  call call_foo
+
+__wrap_foo:
+  ret
+
+#--- call_foo.s
+.globl call_foo
+call_foo:
+  call foo