cmd/compile: rewrite Rsh to RshU if arguments are proved positive

Fixes #76332

Change-Id: I9044025d5dc599531c7f88ed2870bcf3d8b0acbd
Reviewed-on: https://go-review.googlesource.com/c/go/+/721206
Reviewed-by: Mark Freeman <markfreeman@google.com>
Reviewed-by: Keith Randall <khr@google.com>
Auto-Submit: Jorropo <jorropo.pgm@gmail.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@golang.org>

3 files changed, 43 insertions, 20 deletions

diff --git a/src/cmd/compile/internal/ssa/prove.go b/src/cmd/compile/internal/ssa/prove.go
index f6175a3bd7..5581da445d 100644
--- a/src/cmd/compile/internal/ssa/prove.go
+++ b/src/cmd/compile/internal/ssa/prove.go
@@ -2664,14 +2664,30 @@ var mostNegativeDividend = map[Op]int64{
 	OpMod64: -1 << 63,
 }
 var unsignedOp = map[Op]Op{
-	OpDiv8:  OpDiv8u,
-	OpDiv16: OpDiv16u,
-	OpDiv32: OpDiv32u,
-	OpDiv64: OpDiv64u,
-	OpMod8:  OpMod8u,
-	OpMod16: OpMod16u,
-	OpMod32: OpMod32u,
-	OpMod64: OpMod64u,
+	OpDiv8:     OpDiv8u,
+	OpDiv16:    OpDiv16u,
+	OpDiv32:    OpDiv32u,
+	OpDiv64:    OpDiv64u,
+	OpMod8:     OpMod8u,
+	OpMod16:    OpMod16u,
+	OpMod32:    OpMod32u,
+	OpMod64:    OpMod64u,
+	OpRsh8x8:   OpRsh8Ux8,
+	OpRsh8x16:  OpRsh8Ux16,
+	OpRsh8x32:  OpRsh8Ux32,
+	OpRsh8x64:  OpRsh8Ux64,
+	OpRsh16x8:  OpRsh16Ux8,
+	OpRsh16x16: OpRsh16Ux16,
+	OpRsh16x32: OpRsh16Ux32,
+	OpRsh16x64: OpRsh16Ux64,
+	OpRsh32x8:  OpRsh32Ux8,
+	OpRsh32x16: OpRsh32Ux16,
+	OpRsh32x32: OpRsh32Ux32,
+	OpRsh32x64: OpRsh32Ux64,
+	OpRsh64x8:  OpRsh64Ux8,
+	OpRsh64x16: OpRsh64Ux16,
+	OpRsh64x32: OpRsh64Ux32,
+	OpRsh64x64: OpRsh64Ux64,
 }
 
 var bytesizeToConst = [...]Op{
@@ -2758,8 +2774,15 @@ func simplifyBlock(sdom SparseTree, ft *factsTable, b *Block) {
 		case OpRsh8x8, OpRsh8x16, OpRsh8x32, OpRsh8x64,
 			OpRsh16x8, OpRsh16x16, OpRsh16x32, OpRsh16x64,
 			OpRsh32x8, OpRsh32x16, OpRsh32x32, OpRsh32x64,
-			OpRsh64x8, OpRsh64x16, OpRsh64x32, OpRsh64x64,
-			OpLsh8x8, OpLsh8x16, OpLsh8x32, OpLsh8x64,
+			OpRsh64x8, OpRsh64x16, OpRsh64x32, OpRsh64x64:
+			if ft.isNonNegative(v.Args[0]) {
+				if b.Func.pass.debug > 0 {
+					b.Func.Warnl(v.Pos, "Proved %v is unsigned", v.Op)
+				}
+				v.Op = unsignedOp[v.Op]
+			}
+			fallthrough
+		case OpLsh8x8, OpLsh8x16, OpLsh8x32, OpLsh8x64,
 			OpLsh16x8, OpLsh16x16, OpLsh16x32, OpLsh16x64,
 			OpLsh32x8, OpLsh32x16, OpLsh32x32, OpLsh32x64,
 			OpLsh64x8, OpLsh64x16, OpLsh64x32, OpLsh64x64,
diff --git a/test/prove.go b/test/prove.go
index 38181b5fa5..e440634813 100644
--- a/test/prove.go
+++ b/test/prove.go
@@ -253,7 +253,7 @@ func f9(a, b bool) int {
 
 func f10(a string) int {
 	n := len(a)
-	b := a[:n>>1] // ERROR "Proved IsSliceInBounds$"
+	b := a[:n>>1] // ERROR "(Proved IsSliceInBounds|Proved Rsh64x64 is unsigned)$"
 	// We optimize comparisons with small constant strings (see cmd/compile/internal/gc/walk.go),
 	// so this string literal must be long.
 	if b == "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" {
@@ -1086,7 +1086,7 @@ func issue57077(s []int) (left, right []int) {
 }
 
 func issue76332(s []int) (left, right []int) {
-	middle := len(s) >> 1
+	middle := len(s) >> 1 // ERROR "Proved Rsh64x64 is unsigned$"
 	left = s[:middle]     // ERROR "Proved IsSliceInBounds$"
 	right = s[middle:]    // ERROR "Proved IsSliceInBounds$"
 	return
@@ -2564,7 +2564,7 @@ func rightshift(v *[256]int) int {
 		}
 	}
 	for i := range 1024 { // ERROR "Induction"
-		if v[i>>2] == 0 { // ERROR "Proved IsInBounds"
+		if v[i>>2] == 0 { // ERROR "(Proved IsInBounds|Proved Rsh64x64 is unsigned)"
 			return i
 		}
 	}
diff --git a/test/prove_constant_folding.go b/test/prove_constant_folding.go
index 46764f9b9d..831a7e6403 100644
--- a/test/prove_constant_folding.go
+++ b/test/prove_constant_folding.go
@@ -30,10 +30,10 @@ func f0u(x uint) int {
 	}
 
 	if x < 1000 {
-		return int(x)>>31 // ERROR "Proved.+is constant 0$"
+		return int(x) >> 31 // ERROR "(Proved.+is constant 0|Proved Rsh[0-9]+x[0-9]+ is unsigned)$"
 	}
 	if x := int32(x); x < -1000 {
-		return int(x>>31) // ERROR "Proved.+is constant -1$"
+		return int(x >> 31) // ERROR "Proved.+is constant -1$"
 	}
 
 	return int(x) + 1
@@ -45,35 +45,35 @@ func sh64(n int64) int64 {
 	if n < 0 {
 		return n
 	}
-	return n >> 63 // ERROR "Proved .+ is constant 0$"
+	return n >> 63 // ERROR "(Proved .+ is constant 0|Proved Rsh[0-9]+x[0-9]+ is unsigned)$"
 }
 
 func sh32(n int32) int32 {
 	if n < 0 {
 		return n
 	}
-	return n >> 31 // ERROR "Proved .+ is constant 0$"
+	return n >> 31 // ERROR "(Proved .+ is constant 0|Proved Rsh[0-9]+x[0-9]+ is unsigned)$"
 }
 
 func sh32x64(n int32) int32 {
 	if n < 0 {
 		return n
 	}
-	return n >> uint64(31) // ERROR "Proved .+ is constant 0$"
+	return n >> uint64(31) // ERROR "(Proved .+ is constant 0|Proved Rsh[0-9]+x[0-9]+ is unsigned)$"
 }
 
 func sh32x64n(n int32) int32 {
 	if n >= 0 {
 		return 0
 	}
-	return n >> 31// ERROR "Proved .+ is constant -1$"
+	return n >> 31 // ERROR "Proved .+ is constant -1$"
 }
 
 func sh16(n int16) int16 {
 	if n < 0 {
 		return n
 	}
-	return n >> 15 // ERROR "Proved .+ is constant 0$"
+	return n >> 15 // ERROR "(Proved .+ is constant 0|Proved Rsh[0-9]+x[0-9]+ is unsigned)$"
 }
 
 func sh64noopt(n int64) int64 {