cmd/asm: fix TBZ/TBNZ instructions on arm64

Fixes #18069
Also added a test in: cmd/asm/internal/asm/testdata/arm64.s

Change-Id: Iee400bda4f30503ea3c1dc5bb8301568f19c92d1
Signed-off-by: Wei Xiao <wei.xiao@arm.com>
Reviewed-on: https://go-review.googlesource.com/33594
Run-TryBot: Cherry Zhang <cherryyz@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Cherry Zhang <cherryyz@google.com>

src/cmd/asm/internal/arch/arm64.go

@@ -43,6 +43,8 @@ var arm64Jump = map[string]bool{
 	"CBNZ":  true,
 	"CBNZW": true,
 	"JMP":   true,
+	"TBNZ":  true,
+	"TBZ":   true,
 }
 
 func jumpArm64(word string) bool {

src/cmd/asm/internal/asm/asm.go

@@ -390,6 +390,18 @@ func (p *Parser) asmJump(op obj.As, cond string, a []obj.Addr) {
 			}
 			break
 		}
+		if p.arch.Family == sys.ARM64 {
+			// Special 3-operand jumps.
+			// a[0] must be immediate constant; a[1] is a register.
+			if a[0].Type != obj.TYPE_CONST {
+				p.errorf("%s: expected immediate constant; found %s", op, obj.Dconv(prog, &a[0]))
+				return
+			}
+			prog.From = a[0]
+			prog.Reg = p.getRegister(prog, op, &a[1])
+			target = &a[2]
+			break
+		}
 
 		fallthrough
 	default:

src/cmd/asm/internal/asm/testdata/arm64.s

@@ -257,6 +257,8 @@ again:
 	B	foo(SB) // JMP foo(SB)
 	BL	foo(SB) // CALL foo(SB)
 	BEQ	2(PC)
+	TBZ	$1, R1, 2(PC)
+	TBNZ	$2, R2, 2(PC)
 	JMP	foo(SB)
 	CALL	foo(SB)