From deb154001d805a2ba5563cad58e55305483b65bc Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@google.com>
Date: Thu, 7 Jan 2016 17:20:07 +0000
Subject: [PATCH] ELF: Implement --wrap.

In this patch, all symbols are resolved normally and then wrap options
are applied. Renaming is implemented by mutating `Body` pointers of
Symbols. (As a result, Symtab.find(SymbolName)->getName() may return
a string that's different from SymbolName, but that is by design.
I designed the symbol and the symbol table to allow this kind of
operations.)

http://reviews.llvm.org/D15896

llvm-svn: 257075
---
 lld/ELF/Driver.cpp         |  3 +++
 lld/ELF/Options.td         |  4 ++++
 lld/ELF/SymbolTable.cpp    | 14 ++++++++++++++
 lld/ELF/SymbolTable.h      |  1 +
 lld/ELF/Symbols.h          |  1 +
 lld/test/ELF/Inputs/wrap.s |  4 ++++
 lld/test/ELF/wrap.s        | 19 +++++++++++++++++++
 7 files changed, 46 insertions(+)
 create mode 100644 lld/test/ELF/Inputs/wrap.s
 create mode 100644 lld/test/ELF/wrap.s

diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
index da666cfb1107..12ad8ecac3ef 100644
--- a/lld/ELF/Driver.cpp
+++ b/lld/ELF/Driver.cpp
@@ -310,6 +310,9 @@ template <class ELFT> void LinkerDriver::link(opt::InputArgList &Args) {
   for (StringRef S : Config->Undefined)
     Symtab.addUndefinedOpt(S);
 
+  for (auto *Arg : Args.filtered(OPT_wrap))
+    Symtab.wrap(Arg->getValue());
+
   if (Config->OutputFile.empty())
     Config->OutputFile = "a.out";
 
diff --git a/lld/ELF/Options.td b/lld/ELF/Options.td
index 622cbb93bf11..1b02c5c8b795 100644
--- a/lld/ELF/Options.td
+++ b/lld/ELF/Options.td
@@ -109,6 +109,9 @@ def verbose : Flag<["--"], "verbose">;
 def whole_archive : Flag<["--", "-"], "whole-archive">,
   HelpText<"Force load of all members in a static library">;
 
+def wrap : Separate<["--", "-"], "wrap">, MetaVarName<"<symbol>">,
+  HelpText<"Use wrapper functions for symbol">;
+
 def z : JoinedOrSeparate<["-"], "z">, MetaVarName<"<option>">,
   HelpText<"Linker option extensions">;
 
@@ -136,6 +139,7 @@ def alias_soname_soname : Separate<["-"], "soname">, Alias<soname>;
 def alias_script_T : Separate<["-"], "T">, Alias<script>;
 def alias_strip_all: Flag<["-"], "s">, Alias<strip_all>;
 def alias_undefined_u : Separate<["-"], "u">, Alias<undefined>;
+def alias_wrap_wrap : Joined<["--", "-"], "wrap=">, Alias<wrap>;
 
 // Our symbol resolution algorithm handles symbols in archive files differently
 // than traditional linkers, so we don't need --start-group and --end-group.
diff --git a/lld/ELF/SymbolTable.cpp b/lld/ELF/SymbolTable.cpp
index 74951bad410f..01738d07e29b 100644
--- a/lld/ELF/SymbolTable.cpp
+++ b/lld/ELF/SymbolTable.cpp
@@ -18,6 +18,7 @@
 #include "Config.h"
 #include "Error.h"
 #include "Symbols.h"
+#include "llvm/Support/StringSaver.h"
 
 using namespace llvm;
 using namespace llvm::object;
@@ -124,6 +125,19 @@ SymbolBody *SymbolTable<ELFT>::addIgnored(StringRef Name) {
   return Sym;
 }
 
+// Rename SYM as __wrap_SYM. The original symbol is preserved as __real_SYM.
+// Used to implement --wrap.
+template <class ELFT> void SymbolTable<ELFT>::wrap(StringRef Name) {
+  if (Symtab.count(Name) == 0)
+    return;
+  StringSaver Saver(Alloc);
+  Symbol *Sym = addUndefined(Name)->getSymbol();
+  Symbol *Real = addUndefined(Saver.save("__real_" + Name))->getSymbol();
+  Symbol *Wrap = addUndefined(Saver.save("__wrap_" + Name))->getSymbol();
+  Real->Body = Sym->Body;
+  Sym->Body = Wrap->Body;
+}
+
 // Returns a file from which symbol B was created.
 // If B does not belong to any file, returns a nullptr.
 template <class ELFT>
diff --git a/lld/ELF/SymbolTable.h b/lld/ELF/SymbolTable.h
index 4f8f5afa420a..f240d81f5e7b 100644
--- a/lld/ELF/SymbolTable.h
+++ b/lld/ELF/SymbolTable.h
@@ -57,6 +57,7 @@ public:
   SymbolBody *addIgnored(StringRef Name);
   void scanShlibUndefined();
   SymbolBody *find(StringRef Name);
+  void wrap(StringRef Name);
   ELFFileBase<ELFT> *findFile(SymbolBody *B);
 
 private:
diff --git a/lld/ELF/Symbols.h b/lld/ELF/Symbols.h
index 16a3b338b3cd..59135bfe1f5e 100644
--- a/lld/ELF/Symbols.h
+++ b/lld/ELF/Symbols.h
@@ -105,6 +105,7 @@ public:
   // you can access P->Backref->Body to get the resolver's result.
   void setBackref(Symbol *P) { Backref = P; }
   SymbolBody *repl() { return Backref ? Backref->Body : this; }
+  Symbol *getSymbol() { return Backref; }
 
   // Decides which symbol should "win" in the symbol table, this or
   // the Other. Returns 1 if this wins, -1 if the Other wins, or 0 if
diff --git a/lld/test/ELF/Inputs/wrap.s b/lld/test/ELF/Inputs/wrap.s
new file mode 100644
index 000000000000..584e27033d5c
--- /dev/null
+++ b/lld/test/ELF/Inputs/wrap.s
@@ -0,0 +1,4 @@
+.globl foo, __wrap_foo, __real_foo
+foo = 0x11000
+__wrap_foo = 0x11010
+__real_foo = 0x11020
diff --git a/lld/test/ELF/wrap.s b/lld/test/ELF/wrap.s
new file mode 100644
index 000000000000..17aac2db6e18
--- /dev/null
+++ b/lld/test/ELF/wrap.s
@@ -0,0 +1,19 @@
+// REQUIRES: x86
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %s -o %t
+// RUN: llvm-mc -filetype=obj -triple=x86_64-pc-linux %p/Inputs/wrap.s -o %t2
+
+// RUN: ld.lld -o %t3 %t %t2 -wrap foo -wrap nosuchsym
+// RUN: llvm-objdump -d -print-imm-hex %t3 | FileCheck %s
+// RUN: ld.lld -o %t3 %t %t2 --wrap foo -wrap=nosuchsym
+// RUN: llvm-objdump -d -print-imm-hex %t3 | FileCheck %s
+
+// CHECK: _start:
+// CHECK-NEXT: movl $0x11010, %edx
+// CHECK-NEXT: movl $0x11010, %edx
+// CHECK-NEXT: movl $0x11000, %edx
+
+.global _start
+_start:
+  movl $foo, %edx
+  movl $__wrap_foo, %edx
+  movl $__real_foo, %edx