cmd/compile: fix If lowering on ARM64

On ARM64, an If block is lowered to (NZ cond yes no). This is
incorrect because cond is a boolean value and therefore only the
last byte is meaningful (same as AMD64, see ARM64Ops.go). But here
we are comparing a full register width with 0. Correct it by
comparing only the last bit.

Fixes #52788.

Change-Id: I2cacf9f3d2f45e149c361a290f511b2d4ed845c4
Reviewed-on: https://go-review.googlesource.com/c/go/+/405114
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: David Chase <[email protected]>
Run-TryBot: Cherry Mui <[email protected]>

Author: cherrymui

src/cmd/compile/internal/ssa/gen/ARM64.rules

@@ -532,7 +532,7 @@
 (If (GreaterThanF cc) yes no) => (FGT cc yes no)
 (If (GreaterEqualF cc) yes no) => (FGE cc yes no)
 
-(If cond yes no) => (NZ cond yes no)
+(If cond yes no) => (TBNZ [0] cond yes no)
 
 // atomic intrinsics
 // Note: these ops do not accept offset.
@@ -593,6 +593,21 @@
 (NZ (GreaterThanF cc) yes no) => (FGT cc yes no)
 (NZ (GreaterEqualF cc) yes no) => (FGE cc yes no)
 
+(TBNZ [0] (Equal cc) yes no) => (EQ cc yes no)
+(TBNZ [0] (NotEqual cc) yes no) => (NE cc yes no)
+(TBNZ [0] (LessThan cc) yes no) => (LT cc yes no)
+(TBNZ [0] (LessThanU cc) yes no) => (ULT cc yes no)
+(TBNZ [0] (LessEqual cc) yes no) => (LE cc yes no)
+(TBNZ [0] (LessEqualU cc) yes no) => (ULE cc yes no)
+(TBNZ [0] (GreaterThan cc) yes no) => (GT cc yes no)
+(TBNZ [0] (GreaterThanU cc) yes no) => (UGT cc yes no)
+(TBNZ [0] (GreaterEqual cc) yes no) => (GE cc yes no)
+(TBNZ [0] (GreaterEqualU cc) yes no) => (UGE cc yes no)
+(TBNZ [0] (LessThanF cc) yes no) => (FLT cc yes no)
+(TBNZ [0] (LessEqualF cc) yes no) => (FLE cc yes no)
+(TBNZ [0] (GreaterThanF cc) yes no) => (FGT cc yes no)
+(TBNZ [0] (GreaterEqualF cc) yes no) => (FGE cc yes no)
+
 (EQ (CMPWconst [0] x:(ANDconst [c] y)) yes no) && x.Uses == 1 => (EQ (TSTWconst [int32(c)] y) yes no)
 (NE (CMPWconst [0] x:(ANDconst [c] y)) yes no) && x.Uses == 1 => (NE (TSTWconst [int32(c)] y) yes no)
 (LT (CMPWconst [0] x:(ANDconst [c] y)) yes no) && x.Uses == 1 => (LT (TSTWconst [int32(c)] y) yes no)

src/cmd/compile/internal/ssa/rewriteARM64.go

@@ -0,0 +1,27 @@
+// run
+
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Issue 52788: miscompilation for boolean comparison on ARM64.
+
+package main
+
+import (
+	"fmt"
+	"reflect"
+)
+
+func f(next func() bool) {
+	for b := next(); b; b = next() {
+		fmt.Printf("next() returned %v\n", b)
+	}
+}
+
+func main() {
+	next := reflect.MakeFunc(reflect.TypeOf((func() bool)(nil)), func(_ []reflect.Value) []reflect.Value {
+		return []reflect.Value{reflect.ValueOf(false)}
+	})
+	reflect.ValueOf(f).Call([]reflect.Value{next})
+}