[ELF] Support expressions with -defsym option

[phosek]

Summary

Fixes PR34948.

[grimar]

I think I like this. Few comments:

ELF/ScriptLexer.cpp

206

What is this code for ? I tried following and no testcases failed for me:

if (E == StringRef::npos) {
  assert(false);
  Ret.push_back(S);
  return Ret;
}

231–232

I think instead of passing function as argument it could probably be better to always call maybeSplit
which would do nothing for case when '!InExpr' or '!InDefsym' and it could call either tokenizeDefsym or tokenizeExpr
inside.

ELF/ScriptLexer.h

41

Should this be single new enum variable ?

[ruiu]

ELF/ScriptLexer.cpp

200

Please expand the comment to explain what form of input is expected.

ELF/ScriptParser.cpp

274

You create a new instance of ScriptParser for each -defsym string, so Orig is always the default value which is false.

[ruiu]

ELF/ScriptParser.cpp

1342

Actually, I think a better way of doing this is to pass two StringRefs, one for symbol name and the other for expression to this function. Introducing a different tokenization scheme only to separate "<sym>=<expr>" into <sym> and <expr> seems a bit too much. You probably should split a -defsym argument into two StringRefs in the driver.

[phosek]

ELF/ScriptLexer.cpp

206

It's the -defsym=foo case, I've added test case for this.

ELF/ScriptParser.cpp

274

I've changed this to use enum to represent the context.

1342

That's what I've done in the first version: https://reviews.llvm.org/D39511?id=121190. Would you prefer that approach? The current approach is more akin to what gold and BFD.ld does. The main advantage of this approach is that the error messages that come out of script parser have the right offset, i.e. for -defsym=foo=bar, the error at offset 1 points to f while with the original approach it'd be b.

[ruiu]

ELF/ScriptParser.cpp

1342

Yes, I think I prefer that over this patch. (A reason why you added defsym'ed symbols to the symbol table using addSymbolAlias in that patch isn't clear to me though.)

[phosek]

ELF/ScriptParser.cpp

1342

That was just a leftover code, that should be fixed now.

[ruiu]

LGTM

I think this is much better than the previous version.

ELF/Driver.cpp

1047–1048

nit: you can eliminate by MB by readDefsym(From, MemoryBufferRef(To, "-defsym")).

ELF/ScriptParser.cpp

274

After this, you should call atEOF() to make sure that you've read all characters. Otherwise, I think you'll miss an error in `-defsym="(1+2)++++" for example, as readExpr will return at the closing parenthesis.

ELF/Driver.cpp

if (!to_integer(S, V)) {
     error("-image-base: number expected, but got " + S);
     return 0;
   }
   if ((V % Config->MaxPageSize) != 0)
     warn("-image-base: address isn't multiple of page size: " + S);
   return V;
 }
 
-// Parses --defsym=alias option.
-static std::vector<std::pair<StringRef, StringRef>>
-getDefsym(opt::InputArgList &Args) {
-  std::vector<std::pair<StringRef, StringRef>> Ret;
-  for (auto *Arg : Args.filtered(OPT_defsym)) {
-    StringRef From;
-    StringRef To;
-    std::tie(From, To) = StringRef(Arg->getValue()).split('=');
-    if (!isValidCIdentifier(To))
-      error("--defsym: symbol name expected, but got " + To);
-    Ret.push_back({From, To});
-  }
-  return Ret;
-}
-
 // Parses `--exclude-libs=lib,lib,...`.
 // The library names may be delimited by commas or colons.
 static DenseSet<StringRef> getExcludeLibs(opt::InputArgList &Args) {
   DenseSet<StringRef> Ret;
   for (auto *Arg : Args.filtered(OPT_exclude_libs)) {
     StringRef S = Arg->getValue();
     for (;;) {
       size_t Pos = S.find_first_of(",:");
template <class ELFT> void LinkerDriver::link(opt::InputArgList &Args) {
   for (auto *Arg : Args.filtered(OPT_trace_symbol))
     Symtab->trace(Arg->getValue());
 
   // Add all files to the symbol table. This will add almost all
   // symbols that we need to the symbol table.
   for (InputFile *F : Files)
     Symtab->addFile<ELFT>(F);
 
+  // Process -defsym option.
+  for (auto *Arg : Args.filtered(OPT_defsym)) {
+    MemoryBufferRef MB(Arg->getValue(), "-defsym");
+    readDefsym(MB);
+  }
+
   // Now that we have every file, we can decide if we will need a
   // dynamic symbol table.
   // We need one if we were asked to export dynamic symbols or if we are
   // producing a shared library.
   // We also need one if any shared libraries are used and for pie executables
   // (probably because the dynamic linker needs it).
   Config->HasDynSymTab =
       !SharedFiles.empty() || Config->Pic || Config->ExportDynamic;
 
template <class ELFT> void LinkerDriver::link(opt::InputArgList &Args) {
 
   // Apply version scripts.
   Symtab->scanVersionScript();
 
   // Create wrapped symbols for -wrap option.
   for (auto *Arg : Args.filtered(OPT_wrap))
     Symtab->addSymbolWrap<ELFT>(Arg->getValue());
 
-  // Create alias symbols for -defsym option.
-  for (std::pair<StringRef, StringRef> &Def : getDefsym(Args))
-    Symtab->addSymbolAlias<ELFT>(Def.first, Def.second);
-
   Symtab->addCombinedLTOObject<ELFT>();
   if (errorCount())
     return;
 
-  // Apply symbol renames for -wrap and -defsym
+  // Apply symbol renames for -wrap.
   Symtab->applySymbolRenames();
 
   // Now that we have a complete list of input files.
   // Beyond this point, no new files are added.
   // Aggregate all input sections into one place.
   for (InputFile *F : ObjectFiles)
     for (InputSectionBase *S : F->getSections())
       if (S && S != &InputSection::Discarded)

ELF/ScriptLexer.h

public:
   void skip();
   bool consume(StringRef Tok);
   void expect(StringRef Expect);
   bool consumeLabel(StringRef Tok);
   std::string getCurrentLocation();
 
   std::vector<MemoryBufferRef> MBs;
   std::vector<StringRef> Tokens;
-  bool InExpr = false;
+  enum class LexerContext { Default, Defsym, Expr };
+  LexerContext Context = LexerContext::Default;
   size_t Pos = 0;
 
 private:
-  void maybeSplitExpr();
+  void maybeSplit();
   StringRef getLine();
   size_t getLineNumber();
   size_t getColumnNumber();
 
   MemoryBufferRef getCurrentMB();
 };
 
 } // namespace elf
 } // namespace lld
 
 #endif

ELF/ScriptLexer.cpp

while (!S.empty()) {
     } else {
       Ret.push_back(S.substr(E, 1));
       S = S.substr(E + 1);
     }
   }
   return Ret;
 }
 
+// Split a given string as a defsym.
+// The expected format is <symbol name>=<expression> which will get tokenized
+// as "<symbol name>", "=", "<expression>".
+static std::vector<StringRef> tokenizeDefsym(StringRef S) {
+  std::vector<StringRef> Ret;
+
+  size_t E = S.find("=");
+  if (E == StringRef::npos) {
+    Ret.push_back(S);
+    return Ret;
+  }
+
+  if (E != 0)
+    Ret.push_back(S.substr(0, E));
+
+  Ret.push_back(S.substr(E, 1));
+  S = S.substr(E + 1);
+  Ret.push_back(S);
+
+  return Ret;
+}
+
 // In contexts where expressions are expected, the lexer should apply
 // different tokenization rules than the default one. By default,
 // arithmetic operator characters are regular characters, but in the
 // expression context, they should be independent tokens.
 //
 // For example, "foo*3" should be tokenized to "foo", "*" and "3" only
 // in the expression context.
 //
 // This function may split the current token into multiple tokens.
-void ScriptLexer::maybeSplitExpr() {
-  if (!InExpr || errorCount() || atEOF())
+void ScriptLexer::maybeSplit() {
+  if (errorCount() || atEOF())
     return;
 
-  std::vector<StringRef> V = tokenizeExpr(Tokens[Pos]);
+  std::vector<StringRef> V;
+  switch (Context) {
+  case LexerContext::Defsym:
+    V = tokenizeDefsym(Tokens[Pos]);
+    break;
+  case LexerContext::Expr:
+    V = tokenizeExpr(Tokens[Pos]);
+    break;
+  default:
+    return;
+  }
   if (V.size() == 1)
     return;
   Tokens.erase(Tokens.begin() + Pos);
   Tokens.insert(Tokens.begin() + Pos, V.begin(), V.end());
 }
 
 StringRef ScriptLexer::next() {
-  maybeSplitExpr();
+  maybeSplit();
 
   if (errorCount())
     return "";
   if (atEOF()) {
     setError("unexpected EOF");
     return "";
   }
   return Tokens[Pos++];

ELF/ScriptParser.h

 // Config and ScriptConfig.
 void readLinkerScript(MemoryBufferRef MB);
 
 // Parses a version script.
 void readVersionScript(MemoryBufferRef MB);
 
 void readDynamicList(MemoryBufferRef MB);
 
+// Parses the defsym expression.
+void readDefsym(MemoryBufferRef MB);
+
 } // namespace elf
 } // namespace lld
 
 #endif

ELF/ScriptParser.cpp

 public:
   ScriptParser(MemoryBufferRef MB)
       : ScriptLexer(MB),
         IsUnderSysroot(isUnderSysroot(MB.getBufferIdentifier())) {}
 
   void readLinkerScript();
   void readVersionScript();
   void readDynamicList();
+  void readDefsym();
 
 private:
   void addFile(StringRef Path);
 
   void readAsNeeded();
   void readEntry();
   void readExtern();
   void readGroup();
while (!atEOF()) {
     } else if (SymbolAssignment *Cmd = readProvideOrAssignment(Tok)) {
       Script->SectionCommands.push_back(Cmd);
     } else {
       setError("unknown directive: " + Tok);
     }
   }
 }
 
+void ScriptParser::readDefsym() {
+  auto Orig = Context;
+  Context = LexerContext::Defsym;
+  StringRef Tok = next();
+  Context = Orig;
+  expect("=");
+  Expr E = readExpr();
+  SymbolAssignment *Cmd = make<SymbolAssignment>(Tok, E, getCurrentLocation());
+  Script->SectionCommands.push_back(Cmd);
+}
+
 void ScriptParser::addFile(StringRef S) {
   if (IsUnderSysroot && S.startswith("/")) {
     SmallString<128> PathData;
     StringRef Path = (Config->Sysroot + S).toStringRef(PathData);
     if (sys::fs::exists(Path)) {
       Driver->addFile(Saver.save(Path), /*WithLOption=*/false);
       return;
     }
SymbolAssignment *ScriptParser::readAssignment(StringRef Name) {
   return make<SymbolAssignment>(Name, E, getCurrentLocation());
 }
 
 // This is an operator-precedence parser to parse a linker
 // script expression.
 Expr ScriptParser::readExpr() {
   // Our lexer is context-aware. Set the in-expression bit so that
   // they apply different tokenization rules.
-  bool Orig = InExpr;
-  InExpr = true;
+  auto Orig = Context;
+  Context = LexerContext::Expr;
   Expr E = readExpr1(readPrimary(), 0);
-  InExpr = Orig;
+  Context = Orig;
   return E;
 }
 
 static Expr combine(StringRef Op, Expr L, Expr R) {
   if (Op == "+")
     return [=] { return add(L(), R()); };
   if (Op == "-")
     return [=] { return sub(L(), R()); };

 
 void elf::readVersionScript(MemoryBufferRef MB) {
   ScriptParser(MB).readVersionScript();
 }
 
 void elf::readDynamicList(MemoryBufferRef MB) {
   ScriptParser(MB).readDynamicList();
 }
+
+void elf::readDefsym(MemoryBufferRef MB) {
+  ScriptParser(MB).readDefsym();
+}

test/ELF/defsym.s

 # CHECK-NEXT:   Value: 0x123
 # CHECK-NEXT:   Size:
 # CHECK-NEXT:   Binding: Global
 # CHECK-NEXT:   Type:
 # CHECK-NEXT:   Other:
 # CHECK-NEXT:   Section: Absolute
 # CHECK-NEXT: }
 # CHECK-NEXT: Symbol {
-# CHECK-NEXT:   Name: foo1
+# CHECK-NEXT:   Name: foo2
 # CHECK-NEXT:   Value: 0x123
 # CHECK-NEXT:   Size:
 # CHECK-NEXT:   Binding: Global
 # CHECK-NEXT:   Type:
 # CHECK-NEXT:   Other:
 # CHECK-NEXT:   Section: Absolute
 # CHECK-NEXT: }
 
 ## Check we can use foo2 and it that it is an alias for foo1.
 # USE:       Disassembly of section .text:
 # USE-NEXT:  _start:
 # USE-NEXT:    movl $0x123, %edx
 
-# RUN: not ld.lld -o %t %t.o --defsym=foo2=1 2>&1 | FileCheck %s -check-prefix=ERR1
-# ERR1: error: --defsym: symbol name expected, but got 1
+# RUN: ld.lld -o %t %t.o --defsym=foo2=1
+# RUN: llvm-readobj -t -s %t | FileCheck %s --check-prefix=ABS
 
-# RUN: not ld.lld -o %t %t.o --defsym=foo2=und 2>&1 | FileCheck %s -check-prefix=ERR2
-# ERR2: error: -defsym: undefined symbol: und
+# ABS:      Symbol {
+# ABS:        Name: foo2
+# ABS-NEXT:   Value: 0x1
+# ABS-NEXT:   Size:
+# ABS-NEXT:   Binding: Global
+# ABS-NEXT:   Type:
+# ABS-NEXT:   Other:
+# ABS-NEXT:   Section: Absolute
+# ABS-NEXT: }
+
+# RUN: ld.lld -o %t %t.o --defsym=foo2=foo1+5
+# RUN: llvm-readobj -t -s %t | FileCheck %s --check-prefix=EXPR
+
+# EXPR:      Symbol {
+# EXPR:        Name: foo1
+# EXPR-NEXT:   Value: 0x123
+# EXPR-NEXT:   Size:
+# EXPR-NEXT:   Binding: Global
+# EXPR-NEXT:   Type:
+# EXPR-NEXT:   Other:
+# EXPR-NEXT:   Section: Absolute
+# EXPR-NEXT: }
+# EXPR-NEXT: Symbol {
+# EXPR-NEXT:   Name: foo2
+# EXPR-NEXT:   Value: 0x128
+# EXPR-NEXT:   Size:
+# EXPR-NEXT:   Binding: Global
+# EXPR-NEXT:   Type:
+# EXPR-NEXT:   Other:
+# EXPR-NEXT:   Section: Absolute
+# EXPR-NEXT: }
+
+# RUN: not ld.lld -o %t %t.o --defsym=foo2=und 2>&1 | FileCheck %s -check-prefix=ERR1
+# ERR1: error: -defsym:1: symbol not found: und
+
+# RUN: not ld.lld -o %t %t.o --defsym=foo2 2>&1 | FileCheck %s -check-prefix=ERR2
+# ERR2: error: -defsym:1: unexpected EOF
 
 .globl foo1
  foo1 = 0x123
 
 .global _start
 _start:
   movl $foo2, %edx