cmd/compile: enhance induction variable detection for unrolled loops

Would suggest extending capabilities (32-bit, unsigned, etc)
in separate CLs because prove bugs are so mystifying.

This implements the suggestion in this comment
https://go-review.googlesource.com/c/go/+/104041/10/src/cmd/compile/internal/ssa/loopbce.go#164
for inferring properly bounded iteration for loops of the form

for i := K0; i < KNN-(K-1); i += K
for i := K0; i <= KNN-K;    i += K

Where KNN is "known non negative" (i.e., len or cap) and K
is also not negative.  Because i <= KNN-K, i+K <= KNN and
no overflow occurs.

Also handles decreasing case (K1 > 0)
for i := KNN; i >= K0; i -= K1
which works when MININT+K1 < K0
(i.e. MININT < K0-K1, no overflow)

Signed only, also only 64 bit for now.

Change-Id: I5da6015aba2f781ec76c4ad59c9c48d952325fdc
Reviewed-on: https://go-review.googlesource.com/c/go/+/136375
Run-TryBot: David Chase <drchase@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Alexandru Moșoi <alexandru@mosoi.ro>

2 files changed, 211 insertions, 4 deletions

diff --git a/src/cmd/compile/internal/ssa/loopbce.go b/src/cmd/compile/internal/ssa/loopbce.go
index 8ab1a0c695..5f02643ccd 100644
--- a/src/cmd/compile/internal/ssa/loopbce.go
+++ b/src/cmd/compile/internal/ssa/loopbce.go
@@ -4,7 +4,10 @@
 
 package ssa
 
-import "fmt"
+import (
+	"fmt"
+	"math"
+)
 
 type indVarFlags uint8
 
@@ -59,6 +62,7 @@ func findIndVar(f *Func) []indVar {
 
 		// Check thet the control if it either ind </<= max or max >/>= ind.
 		// TODO: Handle 32-bit comparisons.
+		// TODO: Handle unsigned comparisons?
 		switch b.Control.Op {
 		case OpLeq64:
 			flags |= indVarMaxInc
@@ -165,15 +169,101 @@ func findIndVar(f *Func) []indVar {
 			continue
 		}
 
-		// We can only guarantee that the loops runs within limits of induction variable
-		// if the increment is ±1 or when the limits are constants.
-		if step != 1 {
+		// We can only guarantee that the loop runs within limits of induction variable
+		// if (one of)
+		// (1) the increment is ±1
+		// (2) the limits are constants
+		// (3) loop is of the form k0 upto Known_not_negative-k inclusive, step <= k
+		// (4) loop is of the form k0 upto Known_not_negative-k exclusive, step <= k+1
+		// (5) loop is of the form Known_not_negative downto k0, minint+step < k0
+		if step > 1 {
 			ok := false
 			if min.Op == OpConst64 && max.Op == OpConst64 {
 				if max.AuxInt > min.AuxInt && max.AuxInt%step == min.AuxInt%step { // handle overflow
 					ok = true
 				}
 			}
+			// Handle induction variables of these forms.
+			// KNN is known-not-negative.
+			// SIGNED ARITHMETIC ONLY. (see switch on b.Control.Op above)
+			// Possibilitis for KNN are len and cap; perhaps we can infer others.
+			// for i := 0; i <= KNN-k    ; i += k
+			// for i := 0; i <  KNN-(k-1); i += k
+			// Also handle decreasing.
+
+			// "Proof" copied from https://go-review.googlesource.com/c/go/+/104041/10/src/cmd/compile/internal/ssa/loopbce.go#164
+			//
+			//	In the case of
+			//	// PC is Positive Constant
+			//	L := len(A)-PC
+			//	for i := 0; i < L; i = i+PC
+			//
+			//	we know:
+			//
+			//	0 + PC does not over/underflow.
+			//	len(A)-PC does not over/underflow
+			//	maximum value for L is MaxInt-PC
+			//	i < L <= MaxInt-PC means i + PC < MaxInt hence no overflow.
+
+			// To match in SSA:
+			// if  (a) min.Op == OpConst64(k0)
+			// and (b) k0 >= MININT + step
+			// and (c) max.Op == OpSubtract(Op{StringLen,SliceLen,SliceCap}, k)
+			// or  (c) max.Op == OpAdd(Op{StringLen,SliceLen,SliceCap}, -k)
+			// or  (c) max.Op == Op{StringLen,SliceLen,SliceCap}
+			// and (d) if upto loop, require indVarMaxInc && step <= k or !indVarMaxInc && step-1 <= k
+
+			if min.Op == OpConst64 && min.AuxInt >= step+math.MinInt64 {
+				knn := max
+				k := int64(0)
+				var kArg *Value
+
+				switch max.Op {
+				case OpSub64:
+					knn = max.Args[0]
+					kArg = max.Args[1]
+
+				case OpAdd64:
+					knn = max.Args[0]
+					kArg = max.Args[1]
+					if knn.Op == OpConst64 {
+						knn, kArg = kArg, knn
+					}
+				}
+				switch knn.Op {
+				case OpSliceLen, OpStringLen, OpSliceCap:
+				default:
+					knn = nil
+				}
+
+				if kArg != nil && kArg.Op == OpConst64 {
+					k = kArg.AuxInt
+					if max.Op == OpAdd64 {
+						k = -k
+					}
+				}
+				if k >= 0 && knn != nil {
+					if inc.AuxInt > 0 { // increasing iteration
+						// The concern for the relation between step and k is to ensure that iv never exceeds knn
+						// i.e., iv < knn-(K-1) ==> iv + K <= knn; iv <= knn-K ==> iv +K < knn
+						if step <= k || flags&indVarMaxInc == 0 && step-1 == k {
+							ok = true
+						}
+					} else { // decreasing iteration
+						// Will be decrementing from max towards min; max is knn-k; will only attempt decrement if
+						// knn-k >[=] min; underflow is only a concern if min-step is not smaller than min.
+						// This all assumes signed integer arithmetic
+						// This is already assured by the test above: min.AuxInt >= step+math.MinInt64
+						ok = true
+					}
+				}
+			}
+
+			// TODO: other unrolling idioms
+			// for i := 0; i < KNN - KNN % k ; i += k
+			// for i := 0; i < KNN&^(k-1) ; i += k // k a power of 2
+			// for i := 0; i < KNN&(-k) ; i += k // k a power of 2
+
 			if !ok {
 				continue
 			}
diff --git a/test/prove.go b/test/prove.go
index 275528dde7..39b23c5e0a 100644
--- a/test/prove.go
+++ b/test/prove.go
@@ -726,6 +726,123 @@ func signHint2(b []byte, n int) {
 	}
 }
 
+// Induction variable in unrolled loop.
+func unrollUpExcl(a []int) int {
+	var i, x int
+	for i = 0; i < len(a)-1; i += 2 { // ERROR "Induction variable: limits \[0,\?\), increment 2$"
+		x += a[i] // ERROR "Proved IsInBounds$"
+		x += a[i+1]
+	}
+	if i == len(a)-1 {
+		x += a[i]
+	}
+	return x
+}
+
+// Induction variable in unrolled loop.
+func unrollUpIncl(a []int) int {
+	var i, x int
+	for i = 0; i <= len(a)-2; i += 2 { // ERROR "Induction variable: limits \[0,\?\], increment 2$"
+		x += a[i]
+		x += a[i+1]
+	}
+	if i == len(a)-1 {
+		x += a[i]
+	}
+	return x
+}
+
+// Induction variable in unrolled loop.
+func unrollDownExcl0(a []int) int {
+	var i, x int
+	for i = len(a) - 1; i > 0; i -= 2 { // ERROR "Induction variable: limits \(0,\?\], increment 2$"
+		x += a[i]   // ERROR "Proved IsInBounds$"
+		x += a[i-1] // ERROR "Proved IsInBounds$"
+	}
+	if i == 0 {
+		x += a[i]
+	}
+	return x
+}
+
+// Induction variable in unrolled loop.
+func unrollDownExcl1(a []int) int {
+	var i, x int
+	for i = len(a) - 1; i >= 1; i -= 2 { // ERROR "Induction variable: limits \[1,\?\], increment 2$"
+		x += a[i]   // ERROR "Proved IsInBounds$"
+		x += a[i-1] // ERROR "Proved IsInBounds$"
+	}
+	if i == 0 {
+		x += a[i]
+	}
+	return x
+}
+
+// Induction variable in unrolled loop.
+func unrollDownInclStep(a []int) int {
+	var i, x int
+	for i = len(a); i >= 2; i -= 2 { // ERROR "Induction variable: limits \[2,\?\], increment 2$"
+		x += a[i-1] // ERROR "Proved IsInBounds$"
+		x += a[i-2]
+	}
+	if i == 1 {
+		x += a[i-1]
+	}
+	return x
+}
+
+// Not an induction variable (step too large)
+func unrollExclStepTooLarge(a []int) int {
+	var i, x int
+	for i = 0; i < len(a)-1; i += 3 {
+		x += a[i]
+		x += a[i+1]
+	}
+	if i == len(a)-1 {
+		x += a[i]
+	}
+	return x
+}
+
+// Not an induction variable (step too large)
+func unrollInclStepTooLarge(a []int) int {
+	var i, x int
+	for i = 0; i <= len(a)-2; i += 3 {
+		x += a[i]
+		x += a[i+1]
+	}
+	if i == len(a)-1 {
+		x += a[i]
+	}
+	return x
+}
+
+// Not an induction variable (min too small, iterating down)
+func unrollDecMin(a []int) int {
+	var i, x int
+	for i = len(a); i >= math.MinInt64; i -= 2 {
+		x += a[i-1]
+		x += a[i-2]
+	}
+	if i == 1 { // ERROR "Disproved Eq64$"
+		x += a[i-1]
+	}
+	return x
+}
+
+// Not an induction variable (min too small, iterating up -- perhaps could allow, but why bother?)
+func unrollIncMin(a []int) int {
+	var i, x int
+	for i = len(a); i >= math.MinInt64; i += 2 {
+		x += a[i-1]
+		x += a[i-2]
+	}
+	if i == 1 { // ERROR "Disproved Eq64$"
+		x += a[i-1]
+	}
+	return x
+}
+
 //go:noinline
 func useInt(a int) {
 }