Support for stride for loops with non-unit step increments

pkg/math.go

@@ -39,7 +39,7 @@ func addIntExpr(x, y ast.Expr) ast.Expr {
 	return &ast.BinaryExpr{X: x, Op: token.ADD, Y: y}
 }
 
-func incrementIntExpr(x ast.Expr) ast.Expr {
+func incIntExpr(x ast.Expr) ast.Expr {
 	if x == nil {
 		return nil
 	}
@@ -104,6 +104,26 @@ func mulIntExpr(x, y ast.Expr) ast.Expr {
 	return &ast.BinaryExpr{X: x, Op: token.MUL, Y: y}
 }
 
+func divIntExpr(x, y ast.Expr) (ast.Expr, bool) {
+	if x == nil || y == nil {
+		return nil, false
+	}
+
+	xInt, xOK := intValue(x)
+	yInt, yOK := intValue(y)
+
+	if xOK && yOK {
+		return intExpr(xInt / yInt), xInt%yInt != 0
+	}
+	if yOK && yInt == 0 {
+		return nil, false
+	}
+	if (xOK && xInt == 0) || (yOK && yInt == 1) {
+		return x, false
+	}
+	return &ast.BinaryExpr{X: x, Op: token.QUO, Y: y}, true
+}
+
 func intExpr(n int) *ast.BasicLit {
 	return &ast.BasicLit{Kind: token.INT, Value: strconv.Itoa(n)}
 }

pkg/prealloc.go

@@ -617,22 +617,60 @@ func (v *returnsVisitor) forLoopCount(stmt *ast.ForStmt) (ast.Expr, bool) {
 		return nil, false
 	}
 
-	postStmt, ok := stmt.Post.(*ast.IncDecStmt)
-	if !ok {
-		if assign, ok := stmt.Post.(*ast.AssignStmt); ok {
-			switch assign.Tok {
-			case token.ADD_ASSIGN, token.SUB_ASSIGN, token.MUL_ASSIGN, token.QUO_ASSIGN, token.REM_ASSIGN, token.SHL_ASSIGN, token.SHR_ASSIGN:
-				return nil, true
+	var postIdent *ast.Ident
+	var reverse bool
+	var step ast.Expr
+	switch s := stmt.Post.(type) {
+	case *ast.IncDecStmt:
+		var ok bool
+		if postIdent, ok = s.X.(*ast.Ident); !ok {
+			return nil, true
+		}
+		reverse = s.Tok == token.DEC
+		step = intExpr(1)
+	case *ast.AssignStmt:
+		if len(s.Lhs) != 1 || len(s.Rhs) != 1 {
+			return nil, true
+		}
+		var ok bool
+		if postIdent, ok = s.Lhs[0].(*ast.Ident); !ok {
+			return nil, true
+		}
+		step = s.Rhs[0]
+		switch s.Tok {
+		case token.ADD_ASSIGN:
+		case token.SUB_ASSIGN:
+			reverse = true
+		case token.MUL_ASSIGN, token.QUO_ASSIGN, token.REM_ASSIGN, token.SHL_ASSIGN, token.SHR_ASSIGN:
+			return nil, true
+		case token.ASSIGN:
+			if binary, ok := s.Rhs[0].(*ast.BinaryExpr); ok {
+				switch binary.Op {
+				case token.ADD:
+				case token.SUB:
+					reverse = true
+				default:
+					return nil, false
+				}
+
+				switch {
+				case exprEqual(binary.X, postIdent):
+					step = binary.Y
+				case exprEqual(binary.Y, postIdent):
+					step = binary.X
+				default:
+					return nil, false
+				}
+			} else {
+				return nil, false
 			}
+		default:
+			return nil, false
 		}
+	default:
 		return nil, false
 	}
 
-	postIdent, ok := postStmt.X.(*ast.Ident)
-	if !ok {
-		return nil, true
-	}
-
 	initStmt, ok := stmt.Init.(*ast.AssignStmt)
 	if !ok {
 		return nil, true
@@ -656,7 +694,7 @@ func (v *returnsVisitor) forLoopCount(stmt *ast.ForStmt) (ast.Expr, bool) {
 
 	upper, op := forLoopUpperBound(stmt.Cond, postIdent.Name)
 
-	if postStmt.Tok == token.INC {
+	if !reverse {
 		if op == token.GTR || op == token.GEQ {
 			return nil, false
 		}
@@ -668,10 +706,15 @@ func (v *returnsVisitor) forLoopCount(stmt *ast.ForStmt) (ast.Expr, bool) {
 	}
 
 	if op == token.LEQ || op == token.GEQ {
-		upper = incrementIntExpr(upper)
+		upper = incIntExpr(upper)
 	}
 
-	return subIntExpr(upper, lower), true
+	countExpr, rounded := divIntExpr(subIntExpr(upper, lower), step)
+	if rounded {
+		// extra capacity in case non-unary step increment is rounded down
+		countExpr = incIntExpr(countExpr)
+	}
+	return countExpr, true
 }
 
 func forLoopUpperBound(expr ast.Expr, name string) (ast.Expr, token.Token) {

testdata/for.go

@@ -302,3 +302,46 @@ func forLinkedListTraversal() {
 		x = append(x, n.id)
 	}
 }
+
+func forIncSkipLit() {
+	var x []int // want "Consider preallocating x with capacity 5$"
+	for i := 0; i < 10; i += 2 {
+		x = append(x, i)
+	}
+}
+
+func forIncSkipLitRemainder() {
+	var x []int // want "Consider preallocating x with capacity 4$"
+	for i := 0; i < 10; i += 3 {
+		x = append(x, i)
+	}
+}
+
+func forDecSkipLit() {
+	var x []int // want "Consider preallocating x with capacity 5$"
+	for i := 10; i > 0; i -= 2 {
+		x = append(x, i)
+	}
+}
+
+func forIncSkipVar() {
+	n := 10
+	var x []int // want "Consider preallocating x with capacity n/2 \\+ 1$"
+	for i := 0; i < n; i += 2 {
+		x = append(x, i)
+	}
+}
+
+func forIncSkipBinaryLit() {
+	var x []int // want "Consider preallocating x with capacity 5$"
+	for i := 0; i < 10; i = i + 2 {
+		x = append(x, i)
+	}
+}
+
+func forIncSkipBinaryLitBackwards() {
+	var x []int // want "Consider preallocating x with capacity 5$"
+	for i := 0; i < 10; i = 2 + i {
+		x = append(x, i)
+	}
+}