From c1373fe44cf6d1fc56e78171cdc8be7afa7c031c Mon Sep 17 00:00:00 2001
From: JanLJL <jan.laukemann@fau.de>
Date: Fri, 3 Mar 2023 14:41:48 +0100
Subject: [PATCH] enabled indexing without shape and lane for vector regs

---
 osaca/parser/parser_AArch64.py | 17 +++++++----------
 tests/test_parser_AArch64.py   | 10 ++++++++++
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/osaca/parser/parser_AArch64.py b/osaca/parser/parser_AArch64.py
index 4cf8c78..a572c94 100755
--- a/osaca/parser/parser_AArch64.py
+++ b/osaca/parser/parser_AArch64.py
@@ -132,8 +132,8 @@ class ParserAArch64(BaseParser):
                 pp.Literal(".")
                 + pp.Optional(pp.Word("12468")).setResultsName("lanes")
                 + pp.Word(pp.alphas, exact=1).setResultsName("shape")
-                + pp.Optional(index)
             )
+            + pp.Optional(index)
         )
         predicate = (
             pp.CaselessLiteral("p").setResultsName("prefix")
@@ -510,15 +510,12 @@ class ParserAArch64(BaseParser):
 
     def get_full_reg_name(self, register):
         """Return one register name string including all attributes"""
-        if "lanes" in register:
-            return (
-                register["prefix"]
-                + str(register["name"])
-                + "."
-                + str(register["lanes"])
-                + register["shape"]
-            )
-        return register["prefix"] + str(register["name"])
+        name = register["prefix"] + str(register["name"])
+        if "shape" in register:
+            name += "." + str(register.get("lanes", "")) + register["shape"]
+        if "index" in register:
+            name += "[" + register["index"] + "]"
+        return name
 
     def normalize_imd(self, imd):
         """Normalize immediate to decimal based representation"""
diff --git a/tests/test_parser_AArch64.py b/tests/test_parser_AArch64.py
index fdcf7f1..053aca1 100755
--- a/tests/test_parser_AArch64.py
+++ b/tests/test_parser_AArch64.py
@@ -81,6 +81,7 @@ class TestParserAArch64(unittest.TestCase):
         instr5 = "ldr x0, [x0, #:got_lo12:q2c]"
         instr6 = "adrp    x0, :got:visited"
         instr7 = "fadd    v17.2d, v16.2d, v1.2d"
+        instr8 = "mov.d   x0, v16.d[1]"
 
         parsed_1 = self.parser.parse_instruction(instr1)
         parsed_2 = self.parser.parse_instruction(instr2)
@@ -89,6 +90,7 @@ class TestParserAArch64(unittest.TestCase):
         parsed_5 = self.parser.parse_instruction(instr5)
         parsed_6 = self.parser.parse_instruction(instr6)
         parsed_7 = self.parser.parse_instruction(instr7)
+        parsed_8 = self.parser.parse_instruction(instr8)
 
         self.assertEqual(parsed_1.instruction, "vcvt.F32.S32")
         self.assertEqual(parsed_1.operands[0].register.name, "1")
@@ -142,6 +144,14 @@ class TestParserAArch64(unittest.TestCase):
         self.assertEqual(parsed_7.operands[0].register.shape, "d")
         self.assertEqual(self.parser.get_full_reg_name(parsed_7.operands[2].register), "v1.2d")
 
+        self.assertEqual(parsed_8.instruction, "mov.d")
+        self.assertEqual(parsed_8.operands[0].register.name, "0")
+        self.assertEqual(parsed_8.operands[0].register.prefix, "x")
+        self.assertEqual(parsed_8.operands[1].register.name, "16")
+        self.assertEqual(parsed_8.operands[1].register.prefix, "v")
+        self.assertEqual(parsed_8.operands[1].register.index, "1")
+        self.assertEqual(self.parser.get_full_reg_name(parsed_8.operands[1].register), "v16.d[1]")
+
     def test_parse_line(self):
         line_comment = "// -- Begin  main"
         line_label = ".LBB0_1:              // =>This Inner Loop Header: Depth=1"