[llvm] Fix movk on aarch64 with an .equ operand (PR #124428)

[via llvm-commits]

https://github.com/DaGenix created https://github.com/llvm/llvm-project/pull/124428

Prior to 5da801386c2b820a4596fc6d8da6b5f4a6da94b4, this code worked:

    .equ    p4_low_b0, 0x0000
    movk    x1, p4_low_b0, lsl 16

That commit fixed a different bug, but accidentally broke the case where the second operand to movk is not a literal.

In 442f066fc464e953b7783230e95ccf2a67ebfb38, a fix was applied to handle the case where the second operand is a value like "(Val) >> 16". However, that didn't appear to fix the test case in this commit. In this commit, we extend the change to handle the case where the second operand is a identifier defined by .equ.

Fixes #124427

I do not understand the code super well - I extended the existing fix and it appears to work / passes all tests that I ran.

I'm not 100% sure where to add a test case. I added it to the best file that seemed to make sense to me, but it could be the wrong place. Even if its the wrong place, it does appear to test the fix - it fails without the fix and passes with it.

>From 6b772fb84c6aba5ca170dad99966516ba776d09e Mon Sep 17 00:00:00 2001
From: Palmer Cox <p at lmercox.com>
Date: Sat, 25 Jan 2025 15:09:19 -0500
Subject: [PATCH] Fix movk on aarch64 with an .equ operand

Prior to 5da801386c2b820a4596fc6d8da6b5f4a6da94b4, this code worked:

    .equ    p4_low_b0, 0x0000
    movk    x1, p4_low_b0, lsl 16

That commit fixed a different bug, but accidentally broke the case where
the second operand to movk is not a literal.

In 442f066fc464e953b7783230e95ccf2a67ebfb38, a fix was applied to handle
the case where the second operand is a value like "(Val) >> 16".
However, that didn't appear to fix the test case in this commit. In this
commit, we extend the change to handle the case where the second operand
is a identifier defined by .equ.
---
 llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp | 4 +++-
 llvm/test/MC/AArch64/arm64-basic-a64-instructions.s    | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
index d3eda48f3276e9..27b052825d2133 100644
--- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
+++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
@@ -5017,7 +5017,9 @@ bool AArch64AsmParser::parseOperand(OperandVector &Operands, bool isCondCode,
       return true;
     E = SMLoc::getFromPointer(getLoc().getPointer() - 1);
     Operands.push_back(AArch64Operand::CreateImm(IdVal, S, E, getContext()));
-    return false;
+
+    // Parse an optional shift/extend modifier.
+    return parseOptionalShiftExtend(getTok());
   }
   case AsmToken::Integer:
   case AsmToken::Real:
diff --git a/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s b/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
index 2f58eadfc8462c..5af4140b3e4101 100644
--- a/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
+++ b/llvm/test/MC/AArch64/arm64-basic-a64-instructions.s
@@ -16,3 +16,6 @@
 // CHECK: crc32ch  w13, w17, w25           // encoding: [0x2d,0x56,0xd9,0x1a]
 // CHECK: crc32cw  wzr, w3, w5             // encoding: [0x7f,0x58,0xc5,0x1a]
 // CHECK: crc32cx  w18, w16, xzr           // encoding: [0x12,0x5e,0xdf,0x9a]
+
+        .equ	p4_low_b0, 0x0000
+        movk x1, p4_low_b0, lsl 16