[AArch64] Refactor @plt, @gotpcrel, and @AUTH to use parseDataExpr

Following PR llvm#132569, which added `parseDataExpr` for parsing
expressions in data directives (e.g., `.word`), this PR migrates AArch64
`@plt`, `@gotpcrel`, and `@AUTH` from the `parsePrimaryExpr` workaround
to `parseDataExpr`. The goal is to align with the GNU assembler model,
where relocation specifiers apply to the entire operand rather than
individual terms, reducing complexity-especially evident in `@AUTH`
parsing.

Note: AArch64 ELF lacks an official syntax for data directives
(llvm#132570). A prefix notation might be a preferable future direction.

In the test elf-reloc-ptrauth.s, many errors are now reported at parse
time.

Author: MaskRay

llvm/include/llvm/MC/MCParser/MCAsmLexer.h

@@ -148,7 +148,9 @@ class MCAsmLexer {
   void setSkipSpace(bool val) { SkipSpace = val; }
 
   bool getAllowAtInIdentifier() { return AllowAtInIdentifier; }
-  void setAllowAtInIdentifier(bool v) { AllowAtInIdentifier = v; }
+  bool setAllowAtInIdentifier(bool v) {
+    return std::exchange(AllowAtInIdentifier, v);
+  }
 
   void setAllowHashInIdentifier(bool V) { AllowHashInIdentifier = V; }
 

llvm/lib/MC/MCParser/AsmParser.cpp

@@ -1204,7 +1204,7 @@ bool AsmParser::parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc,
 
           Split = std::make_pair(Identifier, VName);
         }
-      } else {
+      } else if (Lexer.getAllowAtInIdentifier()) {
         Split = Identifier.split('@');
       }
     } else if (MAI.useParensForSpecifier() &&
@@ -1463,8 +1463,7 @@ bool AsmParser::parseExpression(const MCExpr *&Res, SMLoc &EndLoc) {
   // As a special case, we support 'a op b @ modifier' by rewriting the
   // expression to include the modifier. This is inefficient, but in general we
   // expect users to use 'a@modifier op b'.
-  if (Ctx.getAsmInfo()->useAtForSpecifier() &&
-      parseOptionalToken(AsmToken::At)) {
+  if (Lexer.getAllowAtInIdentifier() && parseOptionalToken(AsmToken::At)) {
     if (Lexer.isNot(AsmToken::Identifier))
       return TokError("unexpected symbol modifier following '@'");
 

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

@@ -18,13 +18,15 @@
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
@@ -180,6 +182,7 @@ class AArch64AsmParser : public MCTargetAsmParser {
   bool showMatchError(SMLoc Loc, unsigned ErrCode, uint64_t ErrorInfo,
                       OperandVector &Operands);
 
+  bool parseDataExpr(const MCExpr *&Res) override;
   bool parseAuthExpr(const MCExpr *&Res, SMLoc &EndLoc);
 
   bool parseDirectiveArch(SMLoc L);
@@ -335,8 +338,6 @@ class AArch64AsmParser : public MCTargetAsmParser {
   unsigned validateTargetOperandClass(MCParsedAsmOperand &Op,
                                       unsigned Kind) override;
 
-  bool parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc) override;
-
   static bool classifySymbolRef(const MCExpr *Expr,
                                 AArch64MCExpr::Specifier &ELFSpec,
                                 MCSymbolRefExpr::VariantKind &DarwinRefKind,
@@ -4478,6 +4479,18 @@ bool AArch64AsmParser::parseSymbolicImmVal(const MCExpr *&ImmVal) {
   if (HasELFModifier)
     ImmVal = AArch64MCExpr::create(ImmVal, RefKind, getContext());
 
+  if (getContext().getAsmInfo()->hasSubsectionsViaSymbols() &&
+      parseOptionalToken(AsmToken::At)) {
+    StringRef Spec;
+    if (getParser().parseIdentifier(Spec))
+      return Error(getLoc(), "expected specifier after '@'");
+    if (auto *SRE = dyn_cast<MCSymbolRefExpr>(ImmVal)) {
+      auto MaybeSpec = getContext().getAsmInfo()->getSpecifierForName(Spec);
+      ImmVal =
+          MCSymbolRefExpr::create(&SRE->getSymbol(), *MaybeSpec, getContext());
+    }
+  }
+
   return false;
 }
 
@@ -8086,11 +8099,48 @@ bool AArch64AsmParser::parseDirectiveAeabiAArch64Attr(SMLoc L) {
   return false;
 }
 
-bool AArch64AsmParser::parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc) {
-  // Try @AUTH expressions: they're more complex than the usual symbol variants.
-  if (!parseAuthExpr(Res, EndLoc))
+bool AArch64AsmParser::parseDataExpr(const MCExpr *&Res) {
+  SMLoc EndLoc;
+
+  if (getParser().parseExpression(Res))
+    return true;
+  MCAsmParser &Parser = getParser();
+  if (!parseOptionalToken(AsmToken::At))
     return false;
-  return getParser().parsePrimaryExpr(Res, EndLoc, nullptr);
+  if (getLexer().getKind() != AsmToken::Identifier)
+    return Error(getLoc(), "expected relocation specifier");
+
+  std::string Identifier = Parser.getTok().getIdentifier().lower();
+  SMLoc Loc = getLoc();
+  Lex();
+  if (Identifier == "auth")
+    return parseAuthExpr(Res, EndLoc);
+
+  auto Spec = MCSymbolRefExpr::VK_PLT;
+  if (Identifier == "gotpcrel")
+    Spec = MCSymbolRefExpr::VK_GOTPCREL;
+  else if (Identifier != "plt")
+    return Error(Loc, "invalid relocation specifier");
+  if (auto *SRE = dyn_cast<MCSymbolRefExpr>(Res))
+    Res = MCSymbolRefExpr::create(&SRE->getSymbol(), Spec, getContext(),
+                                  SRE->getLoc());
+  else
+    return Error(Loc, "this relocation specifier must follow a symbol");
+
+  for (;;) {
+    std::optional<MCBinaryExpr::Opcode> Opcode;
+    if (parseOptionalToken(AsmToken::Plus))
+      Opcode = MCBinaryExpr::Add;
+    else if (parseOptionalToken(AsmToken::Minus))
+      Opcode = MCBinaryExpr::Sub;
+    else
+      break;
+    const MCExpr *Term;
+    if (getParser().parsePrimaryExpr(Term, EndLoc, nullptr))
+      return true;
+    Res = MCBinaryExpr::create(*Opcode, Res, Term, getContext());
+  }
+  return false;
 }
 
 ///  parseAuthExpr
@@ -8100,54 +8150,8 @@ bool AArch64AsmParser::parsePrimaryExpr(const MCExpr *&Res, SMLoc &EndLoc) {
 bool AArch64AsmParser::parseAuthExpr(const MCExpr *&Res, SMLoc &EndLoc) {
   MCAsmParser &Parser = getParser();
   MCContext &Ctx = getContext();
-
   AsmToken Tok = Parser.getTok();
 
-  // Look for '_sym@AUTH' ...
-  if (Tok.is(AsmToken::Identifier) && Tok.getIdentifier().ends_with("@AUTH")) {
-    StringRef SymName = Tok.getIdentifier().drop_back(strlen("@AUTH"));
-    if (SymName.contains('@'))
-      return TokError(
-          "combination of @AUTH with other modifiers not supported");
-    Res = MCSymbolRefExpr::create(Ctx.getOrCreateSymbol(SymName), Ctx);
-
-    Parser.Lex(); // Eat the identifier.
-  } else {
-    // ... or look for a more complex symbol reference, such as ...
-    SmallVector<AsmToken, 6> Tokens;
-
-    // ... '"_long sym"@AUTH' ...
-    if (Tok.is(AsmToken::String))
-      Tokens.resize(2);
-    // ... or '(_sym + 5)@AUTH'.
-    else if (Tok.is(AsmToken::LParen))
-      Tokens.resize(6);
-    else
-      return true;
-
-    if (Parser.getLexer().peekTokens(Tokens) != Tokens.size())
-      return true;
-
-    // In either case, the expression ends with '@' 'AUTH'.
-    if (Tokens[Tokens.size() - 2].isNot(AsmToken::At) ||
-        Tokens[Tokens.size() - 1].isNot(AsmToken::Identifier) ||
-        Tokens[Tokens.size() - 1].getIdentifier() != "AUTH")
-      return true;
-
-    if (Tok.is(AsmToken::String)) {
-      StringRef SymName;
-      if (Parser.parseIdentifier(SymName))
-        return true;
-      Res = MCSymbolRefExpr::create(Ctx.getOrCreateSymbol(SymName), Ctx);
-    } else {
-      if (Parser.parsePrimaryExpr(Res, EndLoc, nullptr))
-        return true;
-    }
-
-    Parser.Lex(); // '@'
-    Parser.Lex(); // 'AUTH'
-  }
-
   // At this point, we encountered "<id>@AUTH". There is no fallback anymore.
   if (parseToken(AsmToken::LParen, "expected '('"))
     return true;

llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCAsmInfo.cpp

@@ -61,6 +61,7 @@ AArch64MCAsmInfoDarwin::AArch64MCAsmInfoDarwin(bool IsILP32) {
   UsesELFSectionDirectiveForBSS = true;
   SupportsDebugInformation = true;
   UseDataRegionDirectives = true;
+  UseAtForSpecifier = false;
 
   ExceptionsType = ExceptionHandling::DwarfCFI;
 
@@ -105,6 +106,7 @@ AArch64MCAsmInfoELF::AArch64MCAsmInfoELF(const Triple &T) {
   Data64bitsDirective = "\t.xword\t";
 
   UseDataRegionDirectives = false;
+  UseAtForSpecifier = false;
 
   WeakRefDirective = "\t.weak\t";
 

llvm/test/MC/AArch64/data-directive-specifier.s

@@ -1,5 +1,6 @@
 # RUN: llvm-mc -triple=aarch64 -filetype=obj %s | llvm-readobj -r - | FileCheck %s
-# RUN: not llvm-mc -triple=aarch64 -filetype=obj %s --defsym ERR=1 -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR --implicit-check-not=error:
+# RUN: not llvm-mc -triple=aarch64 %s --defsym ERR=1 -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR --implicit-check-not=error:
+# RUN: not llvm-mc -triple=aarch64 -filetype=obj %s --defsym OBJERR=1 -o /dev/null 2>&1 | FileCheck %s --check-prefix=OBJERR --implicit-check-not=error:
 
 .globl g
 g:
@@ -32,13 +33,21 @@ data1:
 .word extern@GOTPCREL-5
 
 .ifdef ERR
-# ERR: [[#@LINE+1]]:7: error: symbol 'und' can not be undefined in a subtraction expression
-.word extern@plt - und
+# ERR: [[#@LINE+1]]:9: error: this relocation specifier must follow a symbol
+.quad 3@plt - .
+
+# ERR: [[#@LINE+1]]:9: error: expected ')'
+.quad (l@plt - .)
+.endif
 
+.ifdef OBJERR
 .quad g@plt - .
 
 .word extern@gotpcrel - .
 
-# ERR: [[#@LINE+1]]:7: error: symbol 'und' can not be undefined in a subtraction expression
+# OBJERR: [[#@LINE+1]]:7: error: symbol 'und' can not be undefined in a subtraction expression
+.word extern@plt - und
+
+# OBJERR: [[#@LINE+1]]:7: error: symbol 'und' can not be undefined in a subtraction expression
 .word extern@gotpcrel - und
 .endif

llvm/test/MC/AArch64/elf-reloc-ptrauth.s

@@ -1,4 +1,4 @@
-// RUN: llvm-mc -triple=aarch64 %s --defsym=ASMONLY=1 | FileCheck %s --check-prefix=ASM
+// RUN: llvm-mc -triple=aarch64 %s | FileCheck %s --check-prefix=ASM
 
 // RUN: llvm-mc -triple=aarch64 -filetype=obj %s | \
 // RUN:   llvm-readelf -S -r -x .test - | FileCheck %s --check-prefix=RELOC
@@ -41,8 +41,6 @@
 // RELOC-NEXT: 70 00000000 10000000
 //                         ^^^^ discriminator
 //                               ^^ 0 no addr diversity 0 reserved 00 ia key 0000 reserved
-// RELOC-NEXT: 80 04000000 00000000
-// Folded to constant 4 bytes difference between _g9 and _g8
 
 .section    .helper
 .local "_g 6"
@@ -63,12 +61,12 @@ _g9:
 .quad _g0@AUTH(ia,42)
 .quad 0
 
-// ASM:          .xword _g1@AUTH(ib,0)
-.quad _g1@AUTH(ib,0)
+// ASM:          .xword (+_g1)@AUTH(ib,0)
+.quad +_g1@AUTH(ib,0)
 .quad 0
 
 // ASM:          .xword _g2@AUTH(da,5,addr)
-.quad _g2@AUTH(da,5,addr)
+.quad _g2 @ AUTH(da,5,addr)
 .quad 0
 
 // ASM:          .xword _g3@AUTH(db,65535,addr)
@@ -91,33 +89,20 @@ _g9:
 .quad ("_g 7" + 7)@AUTH(ia,16)
 .quad 0
 
-// ASM:          .xword _g9@AUTH(ia,42)-_g8@AUTH(ia,42)
-.quad _g9@AUTH(ia,42) - _g8@AUTH(ia,42)
-.quad 0
+// RUN: not llvm-mc -triple=aarch64 --defsym=ERR=1 %s 2>&1 | \
+// RUN:   FileCheck %s --check-prefix=ERR
 
-.ifdef ASMONLY
+.ifdef ERR
 
-// ASM:          .xword _g10@AUTH(ia,42)+1
 .quad _g10@AUTH(ia,42) + 1
 
-// ASM:          .xword 1+_g11@AUTH(ia,42)
 .quad 1 + _g11@AUTH(ia,42)
 
-// ASM:          .xword 1+_g12@AUTH(ia,42)+1
 .quad 1 + _g12@AUTH(ia,42) + 1
 
-// ASM:          .xword _g13@AUTH(ia,42)+_g14@AUTH(ia,42)
 .quad _g13@AUTH(ia,42) + _g14@AUTH(ia,42)
 
-// ASM:          .xword _g9@AUTH(ia,42)-_g8
 .quad _g9@AUTH(ia,42) - _g8
-.quad 0
-
-.endif // ASMONLY
-
-.ifdef ERR
-// RUN: not llvm-mc -triple=aarch64 --defsym=ERR=1 %s 2>&1 | \
-// RUN:   FileCheck %s --check-prefix=ERR
 
 // ERR: :[[#@LINE+1]]:15: error: expected '('
 .quad sym@AUTH)ia,42)
@@ -143,51 +128,28 @@ _g9:
 // ERR: :[[#@LINE+1]]:21: error: expected ')'
 .quad sym@AUTH(ia,42(
 
-// ERR: :[[#@LINE+1]]:7: error: combination of @AUTH with other modifiers not supported
+// ERR: :[[#@LINE+1]]:14: error: unexpected token
 .quad sym@PLT@AUTH(ia,42)
 
-// ERR: :[[#@LINE+1]]:11: error: invalid variant 'AUTH@GOT'
+// ERR: :[[#@LINE+1]]:15: error: expected '('
 .quad sym@AUTH@GOT(ia,42)
 
-// ERR: :[[#@LINE+1]]:18: error: invalid variant 'TLSDESC@AUTH'
-.quad "long sym"@TLSDESC@AUTH(ia,42)
-
-// ERR: :[[#@LINE+1]]:18: error: invalid variant 'AUTH@PLT'
+// ERR: :[[#@LINE+1]]:22: error: expected '('
 .quad "long sym"@AUTH@PLT(ia,42)
 
-// ERR: :[[#@LINE+1]]:17: error: invalid variant 'GOT@AUTH'
+// ERR: :[[#@LINE+1]]:17: error: invalid relocation specifier
 .quad (sym - 5)@GOT@AUTH(ia,42)
 
-// ERR: :[[#@LINE+1]]:17: error: invalid variant 'AUTH@TLSDESC'
-.quad (sym + 5)@AUTH@TLSDESC(ia,42)
-
-// ERR: :[[#@LINE+1]]:12: error: invalid variant 'AUTH'
-.quad +sym@AUTH(ia,42)
-
-.endif // ERR
-
-.ifdef ERROBJ
-// RUN: not llvm-mc -triple=aarch64 -filetype=obj --defsym=ERROBJ=1 %s -o /dev/null 2>&1 | \
-// RUN:   FileCheck %s --check-prefix=ERROBJ
-
-// ERROBJ: :[[#@LINE+1]]:7: error: expected relocatable expression
+// ERR: :[[#@LINE+1]]:23: error: unexpected token
 .quad sym@AUTH(ia,42) + 1
 
-// ERROBJ: :[[#@LINE+1]]:7: error: expected relocatable expression
-.quad 1 + sym@AUTH(ia,42)
-
-// ERROBJ: :[[#@LINE+1]]:7: error: expected relocatable expression
+// ERR: :[[#@LINE+1]]:27: error: unexpected token
 .quad 1 + sym@AUTH(ia,42) + 1
 
-// ERROBJ: :[[#@LINE+1]]:7: error: expected relocatable expression
-.quad sym@AUTH(ia,42) + sym@AUTH(ia,42)
-
-// TODO: do we really want to emit an error here? It might not be important
-// whether a symbol has an AUTH modifier or not since the compile-time computed
-// distance remains the same. Leave it in such state as for now since it
-// makes code simpler: subtraction of a non-AUTH symbol and of a constant
-// are handled identically.
-// ERROBJ: :[[#@LINE+1]]:7: error: expected relocatable expression
+/// @AUTH applies to the whole operand instead of an individual term.
+/// Trailing expression parts are not allowed even if the logical subtraction
+/// result might make sense.
+// ERR: :[[#@LINE+1]]:23: error: unexpected token
 .quad _g9@AUTH(ia,42) - _g8
 
-.endif // ERROBJ
+.endif // ERR