From 0a0effdd5b654ed90f3c3cfae5d9d9dd0a0db8de Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sat, 25 Jun 2022 22:22:59 -0700
Subject: [PATCH] [ELF] Support -= *= /= <<= >>= &= |= in symbol assignments

---
 lld/ELF/ScriptLexer.cpp                  | 12 ++++++----
 lld/ELF/ScriptParser.cpp                 | 40 ++++++++++++++++++++++++++++----
 lld/test/ELF/linkerscript/operators.test | 25 ++++++++++++++++++++
 3 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/lld/ELF/ScriptLexer.cpp b/lld/ELF/ScriptLexer.cpp
index 6f9ffb7..e0a8ed7 100644
--- a/lld/ELF/ScriptLexer.cpp
+++ b/lld/ELF/ScriptLexer.cpp
@@ -134,10 +134,14 @@ void ScriptLexer::tokenize(MemoryBufferRef mb) {
       continue;
     }
 
-    // ">foo" is parsed to ">" and "foo", but ">>" is parsed to ">>".
-    // "|", "||", "&" and "&&" are different operators.
-    if (s.startswith("<<") || s.startswith("<=") || s.startswith(">>") ||
-        s.startswith(">=") || s.startswith("||") || s.startswith("&&")) {
+    // Some operators form separate tokens.
+    if (s.startswith("<<=") || s.startswith(">>=")) {
+      vec.push_back(s.substr(0, 3));
+      s = s.substr(3);
+      continue;
+    }
+    if (s.size() > 1 && ((s[1] == '=' && strchr("*/+-<>&|", s[0])) ||
+                         (s[0] == s[1] && strchr("<>&|", s[0])))) {
       vec.push_back(s.substr(0, 2));
       s = s.substr(2);
       continue;
diff --git a/lld/ELF/ScriptParser.cpp b/lld/ELF/ScriptParser.cpp
index 1f06005..b3cc479 100644
--- a/lld/ELF/ScriptParser.cpp
+++ b/lld/ELF/ScriptParser.cpp
@@ -1038,11 +1038,14 @@ SymbolAssignment *ScriptParser::readAssignment(StringRef tok) {
 
   size_t oldPos = pos;
   SymbolAssignment *cmd = nullptr;
-  if (peek().startswith("=")) {
+  const StringRef op = peek();
+  if (op.startswith("=")) {
     // Support = followed by an expression without whitespace.
     SaveAndRestore<bool> saved(inExpr, true);
     cmd = readSymbolAssignment(tok);
-  } else if (peek() == "+=") {
+  } else if ((op.size() == 2 && op[1] == '=' &&
+              is_contained("*/+-&|", op[0])) ||
+             op == "<<=" || op == ">>=") {
     cmd = readSymbolAssignment(tok);
   } else if (tok == "PROVIDE") {
     SaveAndRestore<bool> saved(inExpr, true);
@@ -1067,11 +1070,38 @@ SymbolAssignment *ScriptParser::readAssignment(StringRef tok) {
 SymbolAssignment *ScriptParser::readSymbolAssignment(StringRef name) {
   name = unquote(name);
   StringRef op = next();
-  assert(op == "=" || op == "+=");
+  assert(op == "=" || op == "*=" || op == "/=" || op == "+=" || op == "-=" ||
+         op == "&=" || op == "|=" || op == "<<=" || op == ">>=");
+  // Note: GNU ld does not support %= or ^=.
   Expr e = readExpr();
-  if (op == "+=") {
+  if (op != "=") {
     std::string loc = getCurrentLocation();
-    e = [=] { return add(script->getSymbolValue(name, loc), e()); };
+    e = [=, c = op[0]]() -> ExprValue {
+      ExprValue lhs = script->getSymbolValue(name, loc);
+      switch (c) {
+      case '*':
+        return lhs.getValue() * e().getValue();
+      case '/':
+        if (uint64_t rv = e().getValue())
+          return lhs.getValue() / rv;
+        error(loc + ": division by zero");
+        return 0;
+      case '+':
+        return add(lhs, e());
+      case '-':
+        return sub(lhs, e());
+      case '<':
+        return lhs.getValue() << e().getValue();
+      case '>':
+        return lhs.getValue() >> e().getValue();
+      case '&':
+        return lhs.getValue() & e().getValue();
+      case '|':
+        return lhs.getValue() | e().getValue();
+      default:
+        llvm_unreachable("");
+      }
+    };
   }
   return make<SymbolAssignment>(name, e, getCurrentLocation());
 }
diff --git a/lld/test/ELF/linkerscript/operators.test b/lld/test/ELF/linkerscript/operators.test
index 2480dad..9ca13c6 100644
--- a/lld/test/ELF/linkerscript/operators.test
+++ b/lld/test/ELF/linkerscript/operators.test
@@ -27,8 +27,22 @@ SECTIONS {
   ternary1 = 0 ? 1 : 2 & 6;
   ternary2 = 1 ? 2?3:4 : 5?6 :7;
 
+  mulassign =2;
+  mulassign *= 2;
+  divassign = 8;
+  divassign /= 2;
   plusassign =1;
   plusassign += 2;
+  minusassign = 3;
+  minusassign -= 1;
+  lshiftassign = 1;
+  lshiftassign <<= 2;
+  rshiftassign = 5;
+  rshiftassign >>= 1;
+  andassign = 6;
+  andassign &= 4;
+  orassign = 4;
+  orassign |= 1;
   braces = 1 + (2 + 3) * 4;
   precedence1 = 1|0xff&1/1<<1+1*2;
   precedence2 = (1 | (0xff & (1 << (1 + (1 * 2)))));
@@ -73,7 +87,14 @@ SECTIONS {
 # CHECK-NEXT: 0000000000000001 A logicalor2
 # CHECK-NEXT: 0000000000000002 A ternary1
 # CHECK-NEXT: 0000000000000003 A ternary2
+# CHECK-NEXT: 0000000000000004 A mulassign
+# CHECK-NEXT: 0000000000000004 A divassign
 # CHECK-NEXT: 0000000000000003 A plusassign
+# CHECK-NEXT: 0000000000000002 A minusassign
+# CHECK-NEXT: 0000000000000004 A lshiftassign
+# CHECK-NEXT: 0000000000000002 A rshiftassign
+# CHECK-NEXT: 0000000000000004 A andassign
+# CHECK-NEXT: 0000000000000005 A orassign
 # CHECK-NEXT: 0000000000000015 A braces
 # CHECK-NEXT: 0000000000000009 A precedence1
 # CHECK-NEXT: 0000000000000009 A precedence2
@@ -135,3 +156,7 @@ SECTIONS {
 # RUN: not ld.lld %t.o -T %t.script -o /dev/null 2>&1 | \
 # RUN:  FileCheck --check-prefix=TERNERR %s
 # TERNERR: : expected, but got ;
+
+## Div by zero error.
+# RUN: echo 'a = 1; a /= 0;' > %t.script
+# RUN: not ld.lld %t.o -T %t.script -o /dev/null 2>&1 | FileCheck --check-prefix=DIVZERO %s
-- 
2.7.4