diff --git a/pkg/prealloc.go b/pkg/prealloc.go index 72d8b95..197890c 100644 --- a/pkg/prealloc.go +++ b/pkg/prealloc.go @@ -94,13 +94,15 @@ func (v *returnsVisitor) Visit(node ast.Node) ast.Visitor { if !ok { continue } - var isArrType bool - switch val := vSpec.Type.(type) { - case *ast.ArrayType: - isArrType = true - case *ast.Ident: - isArrType = contains(v.arrayTypes, val.Name) + + var isArrType = v.isArrayType(vSpec.Type) + if len(vSpec.Values) > 0 { + callExpr, ok := vSpec.Values[0].(*ast.CallExpr) + if ok && isMakeFunc(callExpr) && len(callExpr.Args) > 0 { + isArrType = v.isArrayType(callExpr.Args[0]) + } } + if isArrType { if vSpec.Names != nil { /*atID, ok := arrayType.Elt.(*ast.Ident) @@ -109,7 +111,11 @@ func (v *returnsVisitor) Visit(node ast.Node) ast.Visitor { }*/ // We should handle multiple slices declared on same line e.g. var mySlice1, mySlice2 []uint32 - for _, vName := range vSpec.Names { + for index, vName := range vSpec.Names { + // Skip pre-allocated slices + if index < len(vSpec.Values) && preAllocated(vSpec.Values[index]) { + continue + } v.sliceDeclarations = append(v.sliceDeclarations, &sliceDeclaration{name: vName.Name /*sType: atID.Name,*/, genD: genD}) } } @@ -248,6 +254,17 @@ func (v *returnsVisitor) handleLoops(blockStmt *ast.BlockStmt) { } +// handleLoops is a helper function to check whether a expr is array type +func (v *returnsVisitor) isArrayType(expr ast.Expr) bool { + switch val := expr.(type) { + case *ast.ArrayType: + return true + case *ast.Ident: + return contains(v.arrayTypes, val.Name) + } + return false +} + // Hint stores the information about an occurrence of a slice that could be // preallocated. type Hint struct { @@ -265,3 +282,23 @@ func (h Hint) StringFromFS(f *token.FileSet) string { return fmt.Sprintf("%v:%v Consider preallocating %v", file.Name(), lineNumber, h.DeclaredSliceName) } + +func preAllocated(expr ast.Expr) bool { + callExpr, ok := expr.(*ast.CallExpr) + if !ok || !isMakeFunc(callExpr) { + return false + } + if len(callExpr.Args) > 2 { + literal, ok := callExpr.Args[len(callExpr.Args)-1].(*ast.BasicLit) + return ok && literal.Value != "0" + } + return false +} + +func isMakeFunc(callExpr *ast.CallExpr) bool { + funcIdent, ok := callExpr.Fun.(*ast.Ident) + if !ok { + return false + } + return funcIdent.Name == "make" +} diff --git a/testdata/sample.go b/testdata/sample.go index 8cf72ff..5b2e78a 100644 --- a/testdata/sample.go +++ b/testdata/sample.go @@ -6,6 +6,7 @@ func main() { var z, w, v, u, s []int var t [][]int var intChan chan int + var p, q = make([]int, 0, 10), make([]int, 0) for i, r := range "Hello" { // x is already pre-allocated @@ -26,6 +27,12 @@ func main() { // t is a candidate for pre-allocation t = append(t, foo(i)) + + // p is not a candidate for pre-allocation since it has been pre-allocated + p = append(p, i) + + // q is a candidate for pre-allocation + q = append(q, i) } for i := range intChan {