From ad4129706486ea2c7f71f7d9f8f7465b14ddc538 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sat, 3 Feb 2024 16:40:47 -0800
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4
---
 .../AArch64/AsmParser/AArch64AsmParser.cpp    | 33 +++++++++++++++----
 llvm/test/MC/AArch64/arm64-adr.s              | 10 ++++++
 llvm/test/MC/AArch64/arm64-branch-encoding.s  | 10 ++++++
 llvm/test/MC/AArch64/basic-a64-diagnostics.s  |  8 +++++
 4 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
index e9d96f3b838d4..4e7c8f66d9d54 100644
--- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
+++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
@@ -4809,20 +4809,30 @@ bool AArch64AsmParser::parseOperand(OperandVector &Operands, bool isCondCode,
       return parseCondCode(Operands, invertCondCode);
 
     // If it's a register name, parse it.
-    if (!parseRegister(Operands))
+    if (!parseRegister(Operands)) {
+      // Parse an optional shift/extend modifier.
+      AsmToken SavedTok = getTok();
+      if (parseOptionalToken(AsmToken::Comma)) {
+        // The operand after the register may be a label (e.g. ADR/ADRP).  Check
+        // such cases and don't report an error when <label> happens to match a
+        // shift/extend modifier.
+        ParseStatus Res = MatchOperandParserImpl(Operands, Mnemonic,
+                                                 /*ParseForAllFeatures=*/true);
+        if (!Res.isNoMatch())
+          return Res.isFailure();
+        Res = tryParseOptionalShiftExtend(Operands);
+        if (!Res.isNoMatch())
+          return Res.isFailure();
+        getLexer().UnLex(SavedTok);
+      }
       return false;
+    }
 
     // See if this is a "mul vl" decoration or "mul #<int>" operand used
     // by SVE instructions.
     if (!parseOptionalMulOperand(Operands))
       return false;
 
-    // This could be an optional "shift" or "extend" operand.
-    ParseStatus GotShift = tryParseOptionalShiftExtend(Operands);
-    // We can only continue if no tokens were eaten.
-    if (!GotShift.isNoMatch())
-      return GotShift.isFailure();
-
     // If this is a two-word mnemonic, parse its special keyword
     // operand as an identifier.
     if (Mnemonic == "brb" || Mnemonic == "smstart" || Mnemonic == "smstop" ||
@@ -4883,6 +4893,15 @@ bool AArch64AsmParser::parseOperand(OperandVector &Operands, bool isCondCode,
 
     E = SMLoc::getFromPointer(getLoc().getPointer() - 1);
     Operands.push_back(AArch64Operand::CreateImm(ImmVal, S, E, getContext()));
+
+    // Parse an optional shift/extend modifier.
+    AsmToken SavedTok = Tok;
+    if (parseOptionalToken(AsmToken::Comma)) {
+      ParseStatus Res = tryParseOptionalShiftExtend(Operands);
+      if (!Res.isNoMatch())
+        return Res.isFailure();
+      getLexer().UnLex(SavedTok);
+    }
     return false;
   }
   case AsmToken::Equal: {
diff --git a/llvm/test/MC/AArch64/arm64-adr.s b/llvm/test/MC/AArch64/arm64-adr.s
index 131e545d3bb5c..d1a91b6acf39a 100644
--- a/llvm/test/MC/AArch64/arm64-adr.s
+++ b/llvm/test/MC/AArch64/arm64-adr.s
@@ -23,6 +23,16 @@ adrp x0, foo
 // CHECK: adrp    x0, foo     // encoding: [A,A,A,0x90'A']
 // CHECK-NEXT:                //   fixup A - offset: 0, value: foo, kind: fixup_aarch64_pcrel_adrp_imm21
 
+// CHECK:      adrp    x0, lsl     // encoding: [A,A,A,0x90'A']
+// CHECK-NEXT:                //   fixup A - offset: 0, value: lsl, kind: fixup_aarch64_pcrel_adrp_imm21
+// CHECK-NEXT: adrp    x0, ror     // encoding: [A,A,A,0x90'A']
+// CHECK-NEXT:                //   fixup A - offset: 0, value: ror, kind: fixup_aarch64_pcrel_adrp_imm21
+// CHECK-NEXT: adr    x0, uxtb    // encoding: [A,A,A,0x10'A']
+// CHECK-NEXT:                //   fixup A - offset: 0, value: uxtb, kind: fixup_aarch64_pcrel_adr_imm21
+adrp x0, lsl
+adrp x0, ror
+adr x0, uxtb
+
 adr x0, #0xffffffff
 adrp x0, #0xffffffff
 adrp x0, #1
diff --git a/llvm/test/MC/AArch64/arm64-branch-encoding.s b/llvm/test/MC/AArch64/arm64-branch-encoding.s
index 48c2099012f68..17fdbd9e43fbc 100644
--- a/llvm/test/MC/AArch64/arm64-branch-encoding.s
+++ b/llvm/test/MC/AArch64/arm64-branch-encoding.s
@@ -157,3 +157,13 @@ L1:
 ; CHECK: dcps2                     ; encoding: [0x02,0x00,0xa0,0xd4]
 ; CHECK: dcps3                     ; encoding: [0x03,0x00,0xa0,0xd4]
 
+;; Test "bad" names
+  bl lsl
+  b.eq lsr
+  b.ne uxth
+; CHECK:      bl lsl     ; encoding: [A,A,A,0b100101AA]
+; CHECK-NEXT:   fixup A - offset: 0, value: lsl, kind: fixup_aarch64_pcrel_call26
+; CHECK-NEXT: b.eq lsr   ; encoding: [0bAAA00000,A,A,0x54]
+; CHECK-NEXT:   fixup A - offset: 0, value: lsr, kind: fixup_aarch64_pcrel_branch19
+; CHECK-NEXT: b.ne uxth  ; encoding: [0bAAA00001,A,A,0x54]
+; CHECK-NEXT:   fixup A - offset: 0, value: uxth, kind: fixup_aarch64_pcrel_branch19
diff --git a/llvm/test/MC/AArch64/basic-a64-diagnostics.s b/llvm/test/MC/AArch64/basic-a64-diagnostics.s
index a59861e13472a..028700f4d11df 100644
--- a/llvm/test/MC/AArch64/basic-a64-diagnostics.s
+++ b/llvm/test/MC/AArch64/basic-a64-diagnostics.s
@@ -1149,6 +1149,10 @@
 // CHECK-ERROR-NEXT:           cbz x29, #1
 // CHECK-ERROR-NEXT:                    ^
 
+/// Test "bad" names
+cbz w1, lsl
+// CHECK-ERROR: [[#@LINE-1]]:12: error: expected #imm after shift specifier
+
 //------------------------------------------------------------------------------
 // Conditional branch (immediate)
 //------------------------------------------------------------------------------
@@ -1343,6 +1347,7 @@
         csel sp, x2, x3, ne
         csel x10, x11, sp, ge
         csel x1, x2, x3, #3
+        csel x1, x2, x3, lsl #1, eq
 // CHECK-ERROR: error: invalid operand for instruction
 // CHECK-ERROR-NEXT:        csel w4, wsp, w9, eq
 // CHECK-ERROR-NEXT:                 ^
@@ -1366,6 +1371,9 @@
 // CHECK-ERROR-NEXT:                       ^
 // CHECK-ERROR-NEXT: error: expected AArch64 condition code
 // CHECK-ERROR-NEXT:        csel x1, x2, x3, #3
+// CHECK-ERROR-NEXT:                         ^
+// CHECK-ERROR-NEXT: error: expected AArch64 condition code
+// CHECK-ERROR-NEXT:        csel x1, x2, x3, lsl #1, eq
 // CHECK-ERROR-NEXT:                         ^
 
         csinc w20, w21, wsp, mi