[release-branch.go1.18] cmd/compile: more fix on boolean ops on ARM64

Following CL 421457, the extension rule is also wrong. It is safe
to drop the extension if the value is from a boolean-generating
instruction, but not a boolean-typed Value in general (e.g. a Phi
or a in-register parameter). Fix it.

Updates #52788.
Updates #53397.

Change-Id: Icf3028fe8e90806f9f57fbe2b38d47da27a97e2a
Reviewed-on: https://go-review.googlesource.com/c/go/+/405115
Reviewed-by: David Chase <drchase@google.com>
Run-TryBot: Cherry Mui <cherryyz@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-on: https://go-review.googlesource.com/c/go/+/421458
Reviewed-by: Cherry Mui <cherryyz@google.com>
Run-TryBot: Dmitri Shuralyov <dmitshur@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>

4 files changed, 177 insertions, 5 deletions

diff --git a/src/cmd/compile/internal/ssa/gen/ARM64.rules b/src/cmd/compile/internal/ssa/gen/ARM64.rules
index cc237ad8f3..09e70ad13b 100644
--- a/src/cmd/compile/internal/ssa/gen/ARM64.rules
+++ b/src/cmd/compile/internal/ssa/gen/ARM64.rules
@@ -1578,9 +1578,9 @@
 (GreaterThanF (InvertFlags x)) => (LessThanF x)
 (GreaterEqualF (InvertFlags x)) => (LessEqualF x)
 
-// Boolean-generating instructions always
+// Boolean-generating instructions (NOTE: NOT all boolean Values) always
 // zero upper bit of the register; no need to zero-extend
-(MOVBUreg x) && x.Type.IsBoolean() => (MOVDreg x)
+(MOVBUreg x:((Equal|NotEqual|LessThan|LessThanU|LessThanF|LessEqual|LessEqualU|LessEqualF|GreaterThan|GreaterThanU|GreaterThanF|GreaterEqual|GreaterEqualU|GreaterEqualF) _)) => (MOVDreg x)
 
 // absorb flag constants into conditional instructions
 (CSEL [cc] x _ flag) && ccARM64Eval(cc, flag) > 0 => x
diff --git a/src/cmd/compile/internal/ssa/rewriteARM64.go b/src/cmd/compile/internal/ssa/rewriteARM64.go
index c356d6b97a..1ee25c2eee 100644
--- a/src/cmd/compile/internal/ssa/rewriteARM64.go
+++ b/src/cmd/compile/internal/ssa/rewriteARM64.go
@@ -7321,12 +7321,154 @@ func rewriteValueARM64_OpARM64MOVBUreg(v *Value) bool {
 		v.AuxInt = int64ToAuxInt(int64(uint8(c)))
 		return true
 	}
-	// match: (MOVBUreg x)
-	// cond: x.Type.IsBoolean()
+	// match: (MOVBUreg x:(Equal _))
 	// result: (MOVDreg x)
 	for {
 		x := v_0
-		if !(x.Type.IsBoolean()) {
+		if x.Op != OpARM64Equal {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(NotEqual _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64NotEqual {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(LessThan _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64LessThan {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(LessThanU _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64LessThanU {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(LessThanF _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64LessThanF {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(LessEqual _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64LessEqual {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(LessEqualU _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64LessEqualU {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(LessEqualF _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64LessEqualF {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(GreaterThan _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64GreaterThan {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(GreaterThanU _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64GreaterThanU {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(GreaterThanF _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64GreaterThanF {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(GreaterEqual _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64GreaterEqual {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(GreaterEqualU _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64GreaterEqualU {
+			break
+		}
+		v.reset(OpARM64MOVDreg)
+		v.AddArg(x)
+		return true
+	}
+	// match: (MOVBUreg x:(GreaterEqualF _))
+	// result: (MOVDreg x)
+	for {
+		x := v_0
+		if x.Op != OpARM64GreaterEqualF {
 			break
 		}
 		v.reset(OpARM64MOVDreg)
diff --git a/test/fixedbugs/issue52788a.go b/test/fixedbugs/issue52788a.go
new file mode 100644
index 0000000000..38b8416150
--- /dev/null
+++ b/test/fixedbugs/issue52788a.go
@@ -0,0 +1,29 @@
+// 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 ops on ARM64.
+
+package main
+
+import (
+	"fmt"
+	"reflect"
+	"os"
+)
+
+func f(next func() bool) {
+	for b := next(); b; b = next() {
+		fmt.Printf("%v\n", b)
+		os.Exit(0)
+	}
+}
+
+func main() {
+	next := reflect.MakeFunc(reflect.TypeOf((func() bool)(nil)), func(_ []reflect.Value) []reflect.Value {
+		return []reflect.Value{reflect.ValueOf(true)}
+	})
+	reflect.ValueOf(f).Call([]reflect.Value{next})
+}
diff --git a/test/fixedbugs/issue52788a.out b/test/fixedbugs/issue52788a.out
new file mode 100644
index 0000000000..27ba77ddaf
--- /dev/null
+++ b/test/fixedbugs/issue52788a.out
@@ -0,0 +1 @@
+true