[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/include/llvm/MC/MCParser/MCAsmParser.h

@@ -136,8 +136,10 @@ class MCAsmParser {
   MCTargetAsmParser *TargetParser = nullptr;
 
 protected: // Can only create subclasses.
-  MCAsmParser();
+  MCAsmParser(MCContext &, const MCAsmInfo &);
 
+  MCContext &Ctx;
+  const MCAsmInfo &MAI;
   SmallVector<MCPendingError, 0> PendingErrors;
 
   /// Flag tracking whether any errors have been encountered.
@@ -333,6 +335,9 @@ class MCAsmParser {
 
   /// Parse a .gnu_attribute.
   bool parseGNUAttribute(SMLoc L, int64_t &Tag, int64_t &IntegerValue);
+
+  bool parseAtSpecifier(const MCExpr *&Res, SMLoc &EndLoc);
+  const MCExpr *applySpecifier(const MCExpr *E, uint32_t Variant);
 };
 
 /// Create an MCAsmParser instance for parsing assembly similar to gas syntax

llvm/lib/MC/MCParser/AsmParser.cpp

@@ -14,6 +14,7 @@
 #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/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
@@ -118,9 +119,7 @@ struct ParseStatementInfo {
 class AsmParser : public MCAsmParser {
 private:
   AsmLexer Lexer;
-  MCContext &Ctx;
   MCStreamer &Out;
-  const MCAsmInfo &MAI;
   SourceMgr &SrcMgr;
   SourceMgr::DiagHandlerTy SavedDiagHandler;
   void *SavedDiagContext;
@@ -680,8 +679,6 @@ class AsmParser : public MCAsmParser {
   bool parseEscapedString(std::string &Data) override;
   bool parseAngleBracketString(std::string &Data) override;
 
-  const MCExpr *applySpecifier(const MCExpr *E, uint32_t Variant);
-
   // Macro-like directives
   MCAsmMacro *parseMacroLikeBody(SMLoc DirectiveLoc);
   void instantiateMacroLikeBody(MCAsmMacro *M, SMLoc DirectiveLoc,
@@ -773,7 +770,7 @@ enum { DEFAULT_ADDRSPACE = 0 };
 
 AsmParser::AsmParser(SourceMgr &SM, MCContext &Ctx, MCStreamer &Out,
                      const MCAsmInfo &MAI, unsigned CB = 0)
-    : Lexer(MAI), Ctx(Ctx), Out(Out), MAI(MAI), SrcMgr(SM),
+    : MCAsmParser(Ctx, MAI), Lexer(MAI), Out(Out), SrcMgr(SM),
       CurBuffer(CB ? CB : SM.getMainFileID()), MacrosEnabledFlag(true) {
   HadError = false;
   // Save the old handler.
@@ -1204,7 +1201,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() &&
@@ -1352,7 +1349,7 @@ bool AsmParser::parseExpression(const MCExpr *&Res) {
   return parseExpression(Res, EndLoc);
 }
 
-const MCExpr *AsmParser::applySpecifier(const MCExpr *E, uint32_t Spec) {
+const MCExpr *MCAsmParser::applySpecifier(const MCExpr *E, uint32_t Spec) {
   // Ask the target implementation about this expression first.
   const MCExpr *NewE = getTargetParser().applySpecifier(E, Spec, Ctx);
   if (NewE)
@@ -1443,6 +1440,27 @@ static std::string angleBracketString(StringRef AltMacroStr) {
   return Res;
 }
 
+bool MCAsmParser::parseAtSpecifier(const MCExpr *&Res, SMLoc &EndLoc) {
+  bool SavedAllowAt = getLexer().getAllowAtInIdentifier();
+  getLexer().setAllowAtInIdentifier(true);
+  auto _ = make_scope_exit(
+      [&]() { getLexer().setAllowAtInIdentifier(SavedAllowAt); });
+  if (parseOptionalToken(AsmToken::At)) {
+    if (getLexer().isNot(AsmToken::Identifier))
+      return TokError("expected specifier following '@'");
+
+    auto Spec = MAI.getSpecifierForName(getTok().getIdentifier());
+    if (!Spec)
+      return TokError("invalid specifier '@" + getTok().getIdentifier() + "'");
+
+    const MCExpr *ModifiedRes = applySpecifier(Res, *Spec);
+    if (ModifiedRes)
+      Res = ModifiedRes;
+    Lex();
+  }
+  return false;
+}
+
 /// Parse an expression and return it.
 ///
 ///  expr ::= expr &&,|| expr               -> lowest.
@@ -1463,8 +1481,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/MC/MCParser/MCAsmParser.cpp

@@ -27,7 +27,8 @@ cl::opt<unsigned> AsmMacroMaxNestingDepth(
     cl::desc("The maximum nesting depth allowed for assembly macros."));
 }
 
-MCAsmParser::MCAsmParser() = default;
+MCAsmParser::MCAsmParser(MCContext &Ctx, const MCAsmInfo &MAI)
+    : Ctx(Ctx), MAI(MAI) {}
 
 MCAsmParser::~MCAsmParser() = default;
 

llvm/lib/MC/MCParser/MasmParser.cpp

@@ -376,9 +376,7 @@ FieldInitializer &FieldInitializer::operator=(FieldInitializer &&Initializer) {
 class MasmParser : public MCAsmParser {
 private:
   AsmLexer Lexer;
-  MCContext &Ctx;
   MCStreamer &Out;
-  const MCAsmInfo &MAI;
   SourceMgr &SrcMgr;
   SourceMgr::DiagHandlerTy SavedDiagHandler;
   void *SavedDiagContext;
@@ -973,7 +971,7 @@ enum { DEFAULT_ADDRSPACE = 0 };
 
 MasmParser::MasmParser(SourceMgr &SM, MCContext &Ctx, MCStreamer &Out,
                        const MCAsmInfo &MAI, struct tm TM, unsigned CB)
-    : Lexer(MAI), Ctx(Ctx), Out(Out), MAI(MAI), SrcMgr(SM),
+    : MCAsmParser(Ctx, MAI), Lexer(MAI), Out(Out), SrcMgr(SM),
       CurBuffer(CB ? CB : SM.getMainFileID()), TM(TM) {
   HadError = false;
   // Save the old handler.

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,19 @@ bool AArch64AsmParser::parseSymbolicImmVal(const MCExpr *&ImmVal) {
   if (HasELFModifier)
     ImmVal = AArch64MCExpr::create(ImmVal, RefKind, getContext());
 
+  SMLoc EndLoc;
+  if (getContext().getAsmInfo()->hasSubsectionsViaSymbols()) {
+    if (getParser().parseAtSpecifier(ImmVal, EndLoc))
+      return true;
+    const MCExpr *Term;
+    if (parseOptionalToken(AsmToken::Plus)) {
+      if (getParser().parseExpression(Term, EndLoc))
+        return true;
+      ImmVal =
+          MCBinaryExpr::create(MCBinaryExpr::Add, ImmVal, Term, getContext());
+    }
+  }
+
   return false;
 }
 
@@ -5007,11 +5021,18 @@ bool AArch64AsmParser::parseOperand(OperandVector &Operands, bool isCondCode,
 
     // This was not a register so parse other operands that start with an
     // identifier (like labels) as expressions and create them as immediates.
-    const MCExpr *IdVal;
+    const MCExpr *IdVal, *Term;
     S = getLoc();
     if (getParser().parseExpression(IdVal))
       return true;
-    E = SMLoc::getFromPointer(getLoc().getPointer() - 1);
+    if (getParser().parseAtSpecifier(IdVal, E))
+      return true;
+    if (parseOptionalToken(AsmToken::Plus)) {
+      if (getParser().parseExpression(Term, E))
+        return true;
+      IdVal =
+          MCBinaryExpr::create(MCBinaryExpr::Add, IdVal, Term, getContext());
+    }
     Operands.push_back(AArch64Operand::CreateImm(IdVal, S, E, getContext()));
 
     // Parse an optional shift/extend modifier.
@@ -8086,11 +8107,52 @@ 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")
+    Spec = MCSymbolRefExpr::VK_PLT;
+  else if (Identifier == "got") // Mach-O specific
+    Spec = MCSymbolRefExpr::VK_GOT;
+  else
+    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 +8162,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