cmd/compile/internal/ssa: prove support induction variable pair

We have two induction variables i and j in the following loop:

for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
    // loop body
}

This CL enables the prove pass to handle cases where one if block
uses two induction variables.

Updates #45078

Change-Id: I8b8dc8b7b2d160a796dab1d1e29a00ef4e8e8157
Reviewed-on: https://go-review.googlesource.com/c/go/+/757700
Auto-Submit: Keith Randall <khr@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

2 files changed, 192 insertions, 208 deletions

diff --git a/src/cmd/compile/internal/ssa/loopbce.go b/src/cmd/compile/internal/ssa/loopbce.go
index 913a7faf7f..fc96834b88 100644
--- a/src/cmd/compile/internal/ssa/loopbce.go
+++ b/src/cmd/compile/internal/ssa/loopbce.go
@@ -118,246 +118,230 @@ func findIndVar(f *Func) []indVar {
 	var iv []indVar
 	sdom := f.Sdom()
 
+nextblock:
 	for _, b := range f.Blocks {
 		if b.Kind != BlockIf {
 			continue
 		}
-
-		var ind *Value   // induction variable
-		var init *Value  // starting value
-		var limit *Value // ending value
-
-		// Check that the control if it either ind </<= limit or limit </<= ind.
-		// TODO: Handle unsigned comparisons?
 		c := b.Controls[0]
-		inclusive := false
-		switch c.Op {
-		case OpLeq64, OpLeq32, OpLeq16, OpLeq8:
-			inclusive = true
-			fallthrough
-		case OpLess64, OpLess32, OpLess16, OpLess8:
-			ind, limit = c.Args[0], c.Args[1]
-		default:
-			continue
-		}
+		for idx := range 2 {
+			// Check that the control if it either ind </<= limit or limit </<= ind.
+			// TODO: Handle unsigned comparisons?
+			inclusive := false
+			switch c.Op {
+			case OpLeq64, OpLeq32, OpLeq16, OpLeq8:
+				inclusive = true
+			case OpLess64, OpLess32, OpLess16, OpLess8:
+			default:
+				continue nextblock
+			}
 
-		// See if this is really an induction variable
-		less := true
-		init, inc, nxt, loopReturn := parseIndVar(ind)
-		if init == nil {
-			// We failed to parse the induction variable. Before punting, we want to check
-			// whether the control op was written with the induction variable on the RHS
-			// instead of the LHS. This happens for the downwards case, like:
-			//     for i := len(n)-1; i >= 0; i--
-			init, inc, nxt, loopReturn = parseIndVar(limit)
+			less := idx == 0
+			// induction variable, ending value
+			ind, limit := c.Args[idx], c.Args[1-idx]
+			// starting value, increment value, next value, loop return edge
+			init, inc, nxt, loopReturn := parseIndVar(ind)
 			if init == nil {
-				// No recognized induction variable on either operand
-				continue
+				continue // this is not an induction variable
 			}
 
-			// Ok, the arguments were reversed. Swap them, and remember that we're
-			// looking at an ind >/>= loop (so the induction must be decrementing).
-			ind, limit = limit, ind
-			less = false
-		}
-
-		// This is ind.Block.Preds, not b.Preds. That's a restriction on the loop header,
-		// not the comparison block.
-		if len(ind.Block.Preds) != 2 {
-			continue
-		}
+			// This is ind.Block.Preds, not b.Preds. That's a restriction on the loop header,
+			// not the comparison block.
+			if len(ind.Block.Preds) != 2 {
+				continue
+			}
 
-		// Expect the increment to be a nonzero constant.
-		if !inc.isGenericIntConst() {
-			continue
-		}
-		step := inc.AuxInt
-		if step == 0 {
-			continue
-		}
+			// Expect the increment to be a nonzero constant.
+			if !inc.isGenericIntConst() {
+				continue
+			}
+			step := inc.AuxInt
+			if step == 0 {
+				continue
+			}
 
-		// startBody is the edge that eventually returns to the loop header.
-		var startBody Edge
-		switch {
-		case sdom.IsAncestorEq(b.Succs[0].b, loopReturn.b):
-			startBody = b.Succs[0]
-		case sdom.IsAncestorEq(b.Succs[1].b, loopReturn.b):
-			// if x { goto exit } else { goto entry } is identical to if !x { goto entry } else { goto exit }
-			startBody = b.Succs[1]
-			less = !less
-			inclusive = !inclusive
-		default:
-			continue
-		}
+			// startBody is the edge that eventually returns to the loop header.
+			var startBody Edge
+			switch {
+			case sdom.IsAncestorEq(b.Succs[0].b, loopReturn.b):
+				startBody = b.Succs[0]
+			case sdom.IsAncestorEq(b.Succs[1].b, loopReturn.b):
+				// if x { goto exit } else { goto entry } is identical to if !x { goto entry } else { goto exit }
+				startBody = b.Succs[1]
+				less = !less
+				inclusive = !inclusive
+			default:
+				continue
+			}
 
-		// Increment sign must match comparison direction.
-		// When incrementing, the termination comparison must be ind </<= limit.
-		// When decrementing, the termination comparison must be ind >/>= limit.
-		// See issue 26116.
-		if step > 0 && !less {
-			continue
-		}
-		if step < 0 && less {
-			continue
-		}
+			// Increment sign must match comparison direction.
+			// When incrementing, the termination comparison must be ind </<= limit.
+			// When decrementing, the termination comparison must be ind >/>= limit.
+			// See issue 26116.
+			if step > 0 && !less {
+				continue
+			}
+			if step < 0 && less {
+				continue
+			}
 
-		// Up to now we extracted the induction variable (ind),
-		// the increment delta (inc), the temporary sum (nxt),
-		// the initial value (init) and the limiting value (limit).
-		//
-		// We also know that ind has the form (Phi init nxt) where
-		// nxt is (Add inc nxt) which means: 1) inc dominates nxt
-		// and 2) there is a loop starting at inc and containing nxt.
-		//
-		// We need to prove that the induction variable is incremented
-		// only when it's smaller than the limiting value.
-		// Two conditions must happen listed below to accept ind
-		// as an induction variable.
+			// Up to now we extracted the induction variable (ind),
+			// the increment delta (inc), the temporary sum (nxt),
+			// the initial value (init) and the limiting value (limit).
+			//
+			// We also know that ind has the form (Phi init nxt) where
+			// nxt is (Add inc nxt) which means: 1) inc dominates nxt
+			// and 2) there is a loop starting at inc and containing nxt.
+			//
+			// We need to prove that the induction variable is incremented
+			// only when it's smaller than the limiting value.
+			// Two conditions must happen listed below to accept ind
+			// as an induction variable.
 
-		// First condition: the entry block has a single predecessor.
-		// The entry now means the in-loop edge where the induction variable
-		// comparison succeeded. Its predecessor is not necessarily the header
-		// block. This implies that b.Succs[0] is reached iff ind < limit.
-		if len(startBody.b.Preds) != 1 {
-			// the other successor must exit the loop.
-			continue
-		}
+			// First condition: the entry block has a single predecessor.
+			// The entry now means the in-loop edge where the induction variable
+			// comparison succeeded. Its predecessor is not necessarily the header
+			// block. This implies that b.Succs[0] is reached iff ind < limit.
+			if len(startBody.b.Preds) != 1 {
+				// the other successor must exit the loop.
+				continue
+			}
 
-		// Second condition: startBody.b dominates nxt so that
-		// nxt is computed when inc < limit.
-		if !sdom.IsAncestorEq(startBody.b, nxt.Block) {
-			// inc+ind can only be reached through the branch that confirmed the
-			// induction variable is in bounds.
-			continue
-		}
+			// Second condition: startBody.b dominates nxt so that
+			// nxt is computed when inc < limit.
+			if !sdom.IsAncestorEq(startBody.b, nxt.Block) {
+				// inc+ind can only be reached through the branch that confirmed the
+				// induction variable is in bounds.
+				continue
+			}
 
-		// Check for overflow/underflow. We need to make sure that inc never causes
-		// the induction variable to wrap around.
-		// We use a function wrapper here for easy return true / return false / keep going logic.
-		// This function returns true if the increment will never overflow/underflow.
-		ok := func() bool {
-			if step > 0 {
-				if limit.isGenericIntConst() {
-					// Figure out the actual largest value.
-					v := limit.AuxInt
-					if !inclusive {
-						if v == minSignedValue(limit.Type) {
-							return false // < minint is never satisfiable.
+			// Check for overflow/underflow. We need to make sure that inc never causes
+			// the induction variable to wrap around.
+			// We use a function wrapper here for easy return true / return false / keep going logic.
+			// This function returns true if the increment will never overflow/underflow.
+			ok := func() bool {
+				if step > 0 {
+					if limit.isGenericIntConst() {
+						// Figure out the actual largest value.
+						v := limit.AuxInt
+						if !inclusive {
+							if v == minSignedValue(limit.Type) {
+								return false // < minint is never satisfiable.
+							}
+							v--
 						}
-						v--
-					}
-					if init.isGenericIntConst() {
-						// Use stride to compute a better lower limit.
-						if init.AuxInt > v {
+						if init.isGenericIntConst() {
+							// Use stride to compute a better lower limit.
+							if init.AuxInt > v {
+								return false
+							}
+							v = addU(init.AuxInt, diff(v, init.AuxInt)/uint64(step)*uint64(step))
+						}
+						if addWillOverflow(v, step) {
 							return false
 						}
-						v = addU(init.AuxInt, diff(v, init.AuxInt)/uint64(step)*uint64(step))
+						if inclusive && v != limit.AuxInt || !inclusive && v+1 != limit.AuxInt {
+							// We know a better limit than the programmer did. Use our limit instead.
+							limit = f.constVal(limit.Op, limit.Type, v, true)
+							inclusive = true
+						}
+						return true
 					}
-					if addWillOverflow(v, step) {
+					if step == 1 && !inclusive {
+						// Can't overflow because maxint is never a possible value.
+						return true
+					}
+					// If the limit is not a constant, check to see if it is a
+					// negative offset from a known non-negative value.
+					knn, k := findKNN(limit)
+					if knn == nil || k < 0 {
 						return false
 					}
-					if inclusive && v != limit.AuxInt || !inclusive && v+1 != limit.AuxInt {
-						// We know a better limit than the programmer did. Use our limit instead.
-						limit = f.constVal(limit.Op, limit.Type, v, true)
-						inclusive = true
+					// limit == (something nonnegative) - k. That subtraction can't underflow, so
+					// we can trust it.
+					if inclusive {
+						// ind <= knn - k cannot overflow if step is at most k
+						return step <= k
 					}
-					return true
-				}
-				if step == 1 && !inclusive {
-					// Can't overflow because maxint is never a possible value.
-					return true
-				}
-				// If the limit is not a constant, check to see if it is a
-				// negative offset from a known non-negative value.
-				knn, k := findKNN(limit)
-				if knn == nil || k < 0 {
-					return false
-				}
-				// limit == (something nonnegative) - k. That subtraction can't underflow, so
-				// we can trust it.
-				if inclusive {
-					// ind <= knn - k cannot overflow if step is at most k
-					return step <= k
-				}
-				// ind < knn - k cannot overflow if step is at most k+1
-				return step <= k+1 && k != maxSignedValue(limit.Type)
-			} else { // step < 0
-				if limit.Op == OpConst64 {
-					// Figure out the actual smallest value.
-					v := limit.AuxInt
-					if !inclusive {
-						if v == maxSignedValue(limit.Type) {
-							return false // > maxint is never satisfiable.
+					// ind < knn - k cannot overflow if step is at most k+1
+					return step <= k+1 && k != maxSignedValue(limit.Type)
+
+					// 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
+				} else { // step < 0
+					if limit.Op == OpConst64 {
+						// Figure out the actual smallest value.
+						v := limit.AuxInt
+						if !inclusive {
+							if v == maxSignedValue(limit.Type) {
+								return false // > maxint is never satisfiable.
+							}
+							v++
 						}
-						v++
-					}
-					if init.isGenericIntConst() {
-						// Use stride to compute a better lower limit.
-						if init.AuxInt < v {
+						if init.isGenericIntConst() {
+							// Use stride to compute a better lower limit.
+							if init.AuxInt < v {
+								return false
+							}
+							v = subU(init.AuxInt, diff(init.AuxInt, v)/uint64(-step)*uint64(-step))
+						}
+						if subWillUnderflow(v, -step) {
 							return false
 						}
-						v = subU(init.AuxInt, diff(init.AuxInt, v)/uint64(-step)*uint64(-step))
-					}
-					if subWillUnderflow(v, -step) {
-						return false
+						if inclusive && v != limit.AuxInt || !inclusive && v-1 != limit.AuxInt {
+							// We know a better limit than the programmer did. Use our limit instead.
+							limit = f.constVal(limit.Op, limit.Type, v, true)
+							inclusive = true
+						}
+						return true
 					}
-					if inclusive && v != limit.AuxInt || !inclusive && v-1 != limit.AuxInt {
-						// We know a better limit than the programmer did. Use our limit instead.
-						limit = f.constVal(limit.Op, limit.Type, v, true)
-						inclusive = true
+					if step == -1 && !inclusive {
+						// Can't underflow because minint is never a possible value.
+						return true
 					}
-					return true
-				}
-				if step == -1 && !inclusive {
-					// Can't underflow because minint is never a possible value.
-					return true
 				}
+				return false
 			}
-			return false
-
-		}
 
-		if ok() {
-			flags := indVarFlags(0)
-			var min, max *Value
-			if step > 0 {
-				min = init
-				max = limit
-				if inclusive {
+			if ok() {
+				flags := indVarFlags(0)
+				var min, max *Value
+				if step > 0 {
+					min = init
+					max = limit
+					if inclusive {
+						flags |= indVarMaxInc
+					}
+				} else {
+					min = limit
+					max = init
 					flags |= indVarMaxInc
+					if !inclusive {
+						flags |= indVarMinExc
+					}
+					step = -step
 				}
-			} else {
-				min = limit
-				max = init
-				flags |= indVarMaxInc
-				if !inclusive {
-					flags |= indVarMinExc
+				if f.pass.debug >= 1 {
+					printIndVar(b, ind, min, max, step, flags)
 				}
-				step = -step
-			}
-			if f.pass.debug >= 1 {
-				printIndVar(b, ind, min, max, step, flags)
-			}
 
-			iv = append(iv, indVar{
-				ind: ind,
-				nxt: nxt,
-				min: min,
-				max: max,
-				// This is startBody.b, where startBody is the edge from the comparison for the
-				// induction variable, not necessarily the in-loop edge from the loop header.
-				// Induction variable bounds are not valid in the loop before this edge.
-				entry: startBody.b,
-				step:  step,
-				flags: flags,
-			})
-			b.Logf("found induction variable %v (inc = %v, min = %v, max = %v)\n", ind, inc, min, max)
+				iv = append(iv, indVar{
+					ind: ind,
+					nxt: nxt,
+					min: min,
+					max: max,
+					// This is startBody.b, where startBody is the edge from the comparison for the
+					// induction variable, not necessarily the in-loop edge from the loop header.
+					// Induction variable bounds are not valid in the loop before this edge.
+					entry: startBody.b,
+					step:  step,
+					flags: flags,
+				})
+				b.Logf("found induction variable %v (inc = %v, min = %v, max = %v)\n", ind, inc, min, max)
+			}
 		}
-
-		// 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
 	}
 
 	return iv
diff --git a/src/cmd/compile/internal/ssa/prove.go b/src/cmd/compile/internal/ssa/prove.go
index 960d10d98f..f4f2e3c6c7 100644
--- a/src/cmd/compile/internal/ssa/prove.go
+++ b/src/cmd/compile/internal/ssa/prove.go
@@ -1572,7 +1572,7 @@ func getSliceInfo(vp *Value) (inf sliceInfo) {
 // successor.
 func prove(f *Func) {
 	// Find induction variables.
-	var indVars map[*Block]indVar
+	var indVars map[*Block][]indVar
 	for _, v := range findIndVar(f) {
 		ind := v.ind
 		if len(ind.Args) != 2 {
@@ -1585,9 +1585,9 @@ func prove(f *Func) {
 			nxt.Uses == 1) { // 1 used by induction
 			// ind or nxt is used inside the loop, add it for the facts table
 			if indVars == nil {
-				indVars = make(map[*Block]indVar)
+				indVars = make(map[*Block][]indVar)
 			}
-			indVars[v.entry] = v
+			indVars[v.entry] = append(indVars[v.entry], v)
 			continue
 		} else {
 			// Since this induction variable is not used for anything but counting the iterations,
@@ -1677,7 +1677,7 @@ func prove(f *Func) {
 
 			// Entering the block, add facts about the induction variable
 			// that is bound to this block.
-			if iv, ok := indVars[node.block]; ok {
+			for _, iv := range indVars[node.block] {
 				addIndVarRestrictions(ft, parent, iv)
 			}