From 315b70ebfe5c6cd7b916de11e990151e39462129 Mon Sep 17 00:00:00 2001
From: Nicolas Barbotin <nicolas@barbot.in>
Date: Mon, 12 Dec 2022 18:11:18 +1030
Subject: [PATCH] Added patching for tbnz and tbx instructions.

---
 gas-preprocessor.pl | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/gas-preprocessor.pl b/gas-preprocessor.pl
index 066db33..9dca991 100755
--- a/gas-preprocessor.pl
+++ b/gas-preprocessor.pl
@@ -904,7 +904,7 @@ sub handle_serialized_line {
                      ($arch eq "aarch64" and !is_aarch64_register($target))) {
                 $call_targets{$target}++;
             }
-        } elsif ($line =~ /(?:^|\n)\s*(\w+\s*:\s*)?(cbn?z|adr|tbz)\s+(\w+)\s*,(\s*#\d+\s*,)?\s*(\w+)/) {
+        } elsif ($line =~ /(?:^|\n)\s*(\w+\s*:\s*)?(cbn?z|adr|tbn?z)\s+(\w+)\s*,(\s*#\d+\s*,)?\s*(\w+)/) {
             my $instr = $2;
             my $reg = $3;
             my $bit = $4;
@@ -922,6 +922,13 @@ sub handle_serialized_line {
                 $xreg =~ s/w/x/;
                 $line =~ s/\b$reg\b/$xreg/;
             }
+            
+            # Same with tbnz
+            if ($instr eq "tbnz" and $reg =~ /w\d+/) {
+                my $xreg = $reg;
+                $xreg =~ s/w/x/;
+                $line =~ s/\b$reg\b/$xreg/;
+            }
         } elsif ($line =~ /^\s*.h?word.*\b\d+[bf]\b/) {
             while ($line =~ /\b(\d+)([bf])\b/g) {
                 $line = handle_local_label($line, $1, $2);
@@ -1063,7 +1070,25 @@ sub handle_serialized_line {
             }
 
             # Convert "ld1 {v0.4h-v3.4h}" into "ld1 {v0.4h,v1.4h,v2.4h,v3.4h}"
-            if ($line =~ /(?:ld|st)\d\s+({\s*v(\d+)\.(\d[bhsdBHSD])\s*-\s*v(\d+)\.(\d[bhsdBHSD])\s*})/) {
+            if ($line =~ /(?:ld|st)\d\s+({\s*v(\d+)\.(\d+[bhsdBHSD])\s*-\s*v(\d+)\.(\d+[bhsdBHSD])\s*})/) {
+                my $regspec = $1;
+                my $reg1 = $2;
+                my $layout1 = $3;
+                my $reg2 = $4;
+                my $layout2 = $5;
+                if ($layout1 eq $layout2) {
+                    my $new_regspec = "{";
+                    foreach my $i ($reg1 .. $reg2) {
+                        $new_regspec .= "," if ($i > $reg1);
+                        $new_regspec .= "v$i.$layout1";
+                    }
+                    $new_regspec .= "}";
+                    $line =~ s/$regspec/$new_regspec/;
+                }
+            }
+
+            # Convert "tbx vXX, {v0.4h-v3.4h}" into "tbx vXX, {v0.4h,v1.4h,v2.4h,v3.4h}"
+            if ($line =~ /tbx\s+v\d+\.\d+[bhsdBHSD]\s*,\s*({\s*v(\d+)\.(\d+[bhsdBHSD])\s*-\s*v(\d+)\.(\d+[bhsdBHSD])\s*})/) {
                 my $regspec = $1;
                 my $reg1 = $2;
                 my $layout1 = $3;