[AArch64][AsmParser] Add support for parsing shift/extend operands with symbols.

[mcrosier]

Summary

This commit allows symbols defined with the .equ/.set directive to be used as shift/extend operands for AArch64. E.g.,

.set IMM2, 2
.equ IMM4, 4

add w1, w2, w3, uxtb #IMM2
add w4, w5, w6, uxth #IMM4

This addresses PR31431.

PTAL,

Chad

[mcrosier]

-Add a diagnostic check for when the symbol is not defined.

[rovka]

LGTM.

test/MC/AArch64/shift_extend_op_w_symbol.s

3 ↗

(On Diff #82110)

Why are you running the same thing twice? Can't you just 1> %t1 2> %t2 on the first run?

[mcrosier]

Thanks for the review, Diana. I addressed your test comment in the final commit.

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

if (!Hash && getLexer().isNot(AsmToken::Integer)) {
     if (ShOp == AArch64_AM::LSL || ShOp == AArch64_AM::LSR ||
         ShOp == AArch64_AM::ASR || ShOp == AArch64_AM::ROR ||
         ShOp == AArch64_AM::MSL) {
       // We expect a number here.
       TokError("expected #imm after shift specifier");
       return MatchOperand_ParseFail;
     }
 
-    // "extend" type operatoins don't need an immediate, #0 is implicit.
+    // "extend" type operations don't need an immediate, #0 is implicit.
     SMLoc E = SMLoc::getFromPointer(getLoc().getPointer() - 1);
     Operands.push_back(
         AArch64Operand::CreateShiftExtend(ShOp, 0, false, S, E, getContext()));
     return MatchOperand_Success;
   }
 
-  // Make sure we do actually have a number or a parenthesized expression.
+  // Make sure we do actually have a number, identifier or a parenthesized
+  // expression.
   SMLoc E = Parser.getTok().getLoc();
   if (!Parser.getTok().is(AsmToken::Integer) &&
-      !Parser.getTok().is(AsmToken::LParen)) {
+      !Parser.getTok().is(AsmToken::LParen) &&
+      !Parser.getTok().is(AsmToken::Identifier)) {
     Error(E, "expected integer shift amount");
     return MatchOperand_ParseFail;
   }
 
   const MCExpr *ImmVal;
   if (getParser().parseExpression(ImmVal))
     return MatchOperand_ParseFail;
 

llvm/trunk/test/MC/AArch64/shift_extend_op_w_symbol.s

+// RUN: not llvm-mc -triple aarch64-none-linux-gnu < %s 1> %t1 2> %t2
+// RUN: FileCheck < %t1 %s
+// RUN: FileCheck --check-prefix=CHECK-ERROR < %t2 %s
+
+        .globl _func
+_func:
+// CHECK-LABEL: _func
+
+        .set IMM2, 2
+        .equ IMM4, 4
+
+// Make sure we can use a symbol with the optionally shift left operand.
+
+        add w1, w2, w3, uxtb #IMM2
+        add w4, w5, w6, uxth #IMM4
+        add x7, x8, x9, lsl #IMM2
+        add w7, w8, w9, uxtw #IMM4
+        add x1, x2, x3, uxtx #IMM4
+
+// CHECK: add w1, w2, w3, uxtb #2
+// CHECK: add w4, w5, w6, uxth #4
+// CHECK: add x7, x8, x9, lsl #2
+// CHECK: add w7, w8, w9, uxtw #4
+// CHECK: add x1, x2, x3, uxtx #4
+
+        add w1, w2, w3, sxtb #IMM2
+        add w4, w5, w6, sxth #IMM4
+        add x7, x8, x9, lsl #IMM2
+        add w7, w8, w9, sxtw #IMM2
+        add x1, x2, x3, sxtx #IMM4
+
+// CHECK: add w1, w2, w3, sxtb #2
+// CHECK: add w4, w5, w6, sxth #4
+// CHECK: add x7, x8, x9, lsl #2
+// CHECK: add w7, w8, w9, sxtw #2
+// CHECK: add x1, x2, x3, sxtx #4
+
+        add w1, w2, w3, lsl #IMM3
+
+// CHECK-ERROR: error: expected constant '#imm' after shift specifier
+// CHECK-ERROR:        add w1, w2, w3, lsl #IMM3
+// CHECK-ERROR:                             ^