diff --git a/pkg/math.go b/pkg/math.go index 4958f31..4a49236 100644 --- a/pkg/math.go +++ b/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)} } diff --git a/pkg/prealloc.go b/pkg/prealloc.go index 334f3ca..704c184 100644 --- a/pkg/prealloc.go +++ b/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) { diff --git a/testdata/for.go b/testdata/for.go index 6cf46ff..45497f0 100644 --- a/testdata/for.go +++ b/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) + } +}