cmd/asm: add support to shift operands on arm64

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

Change-Id: Icc43ff7383cc06b8eaccabd9ff0aefa61c4ecb88
Reviewed-on: https://go-review.googlesource.com/33595
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/asm/parse.go

@@ -424,7 +424,7 @@ func (p *Parser) atStartOfRegister(name string) bool {
 // We have consumed the register or R prefix.
 func (p *Parser) atRegisterShift() bool {
 	// ARM only.
-	if p.arch.Family != sys.ARM {
+	if !p.arch.InFamily(sys.ARM, sys.ARM64) {
 		return false
 	}
 	// R1<<...
@@ -506,7 +506,7 @@ func (p *Parser) register(name string, prefix rune) (r1, r2 int16, scale int8, o
 	return r1, r2, scale, true
 }
 
-// registerShift parses an ARM shifted register reference and returns the encoded representation.
+// registerShift parses an ARM/ARM64 shifted register reference and returns the encoded representation.
 // There is known to be a register (current token) and a shift operator (peeked token).
 func (p *Parser) registerShift(name string, prefix rune) int64 {
 	if prefix != 0 {
@@ -531,6 +531,8 @@ func (p *Parser) registerShift(name string, prefix rune) int64 {
 	case lex.ARR:
 		op = 2
 	case lex.ROT:
+		// following instructions on ARM64 support rotate right
+		// AND, ANDS, TST, BIC, BICS, EON, EOR, ORR, MVN, ORN
 		op = 3
 	}
 	tok := p.next()
@@ -538,22 +540,37 @@ func (p *Parser) registerShift(name string, prefix rune) int64 {
 	var count int16
 	switch tok.ScanToken {
 	case scanner.Ident:
-		r2, ok := p.registerReference(str)
-		if !ok {
-			p.errorf("rhs of shift must be register or integer: %s", str)
+		if p.arch.Family == sys.ARM64 {
+				p.errorf("rhs of shift must be integer: %s", str)
+		} else {
+			r2, ok := p.registerReference(str)
+			if !ok {
+				p.errorf("rhs of shift must be register or integer: %s", str)
+			}
+			count = (r2&15)<<8 | 1<<4
 		}
-		count = (r2&15)<<8 | 1<<4
 	case scanner.Int, '(':
 		p.back()
 		x := int64(p.expr())
-		if x >= 32 {
-			p.errorf("register shift count too large: %s", str)
+		if p.arch.Family == sys.ARM64 {
+			if x >= 64 {
+				p.errorf("register shift count too large: %s", str)
+			}
+			count = int16((x & 63) << 10)
+		} else {
+			if x >= 32 {
+				p.errorf("register shift count too large: %s", str)
+			}
+			count = int16((x & 31) << 7)
 		}
-		count = int16((x & 31) << 7)
 	default:
 		p.errorf("unexpected %s in register shift", tok.String())
 	}
-	return int64((r1 & 15) | op<<5 | count)
+	if p.arch.Family == sys.ARM64 {
+		return int64(int64(r1 & 31)<<16 | int64(op)<<22 | int64(uint16(count)))
+	} else {
+		return int64((r1 & 15) | op<<5 | count)
+	}
 }
 
 // symbolReference parses a symbol that is known not to be a register.

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

@@ -16,7 +16,6 @@ TEXT	foo(SB), 7, $-8
 //		outcode($1, &$2, $4, &$6);
 //	}
 // imsr comes from the old 7a, we only support immediates and registers
-// at the moment, no shifted registers.
 	ADDW	$1, R2, R3
 	ADDW	R1, R2, R3
 	ADDW	R1, ZR, R3
@@ -24,6 +23,10 @@ TEXT	foo(SB), 7, $-8
 	ADD	R1, R2, R3
 	ADD	R1, ZR, R3
 	ADD	$1, R2, R3
+	ADD	R1>>11, R2, R3
+	ADD	R1<<22, R2, R3
+	ADD	R1->33, R2, R3
+	AND	R1@>33, R2, R3
 
 //	LTYPE1 imsr ',' spreg ','
 //	{
@@ -37,6 +40,10 @@ TEXT	foo(SB), 7, $-8
 	ADDW	R1, R2
 	ADD	$1, R2
 	ADD	R1, R2
+	ADD	R1>>11, R2
+	ADD	R1<<22, R2
+	ADD	R1->33, R2
+	AND	R1@>33, R2
 
 //
 // CLS
@@ -118,7 +125,9 @@ TEXT	foo(SB), 7, $-8
 //	}
 	CMP	$3, R2
 	CMP	R1, R2
-
+	CMP	R1->11, R2
+	CMP	R1>>22, R2
+	CMP	R1<<33, R2
 //
 // CBZ
 //