├── .github └── workflows │ └── build-test.yml ├── LICENSE ├── README.md ├── go.mod ├── import.go ├── pkg └── prealloc.go ├── prealloc.go ├── prealloc_test.go └── testdata └── sample.go /.github/workflows/build-test.yml: -------------------------------------------------------------------------------- 1 | name: build-test 2 | on: 3 | push: 4 | pull_request: 5 | 6 | jobs: 7 | test: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/setup-go@v1 11 | with: 12 | go-version: 1.13.x 13 | - uses: actions/checkout@v2 14 | - name: Build 15 | run: go build . 16 | - name: Test 17 | run: go test -v . 18 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Alex Kohler 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # prealloc 2 | 3 | prealloc is a Go static analysis tool to find slice declarations that could potentially be preallocated. 4 | 5 | ## Installation 6 | 7 | go install github.com/alexkohler/prealloc@latest 8 | 9 | ## Usage 10 | 11 | Similar to other Go static analysis tools (such as golint, go vet), prealloc can be invoked with one or more filenames, directories, or packages named by its import path. Prealloc also supports the `...` wildcard. 12 | 13 | prealloc [flags] files/directories/packages 14 | 15 | ### Flags 16 | - **-simple** (default true) - Report preallocation suggestions only on simple loops that have no returns/breaks/continues/gotos in them. Setting this to false may increase false positives. 17 | - **-rangeloops** (default true) - Report preallocation suggestions on range loops. 18 | - **-forloops** (default false) - Report preallocation suggestions on for loops. This is false by default due to there generally being weirder things happening inside for loops (at least from what I've observed in the Standard Library). 19 | - **-set_exit_status** (default false) - Set exit status to 1 if any issues are found. 20 | 21 | ## Purpose 22 | 23 | While the [Go *does* attempt to avoid reallocation by growing the capacity in advance](https://github.com/golang/go/blob/87e48c5afdcf5e01bb2b7f51b7643e8901f4b7f9/src/runtime/slice.go#L100-L112), this sometimes isn't enough for longer slices. If the size of a slice is known at the time of its creation, it should be specified. 24 | 25 | Consider the following benchmark: (this can be found in prealloc_test.go in this repo) 26 | 27 | ```Go 28 | import "testing" 29 | 30 | func BenchmarkNoPreallocate(b *testing.B) { 31 | existing := make([]int64, 10, 10) 32 | b.ResetTimer() 33 | for i := 0; i < b.N; i++ { 34 | // Don't preallocate our initial slice 35 | var init []int64 36 | for _, element := range existing { 37 | init = append(init, element) 38 | } 39 | } 40 | } 41 | 42 | func BenchmarkPreallocate(b *testing.B) { 43 | existing := make([]int64, 10, 10) 44 | b.ResetTimer() 45 | for i := 0; i < b.N; i++ { 46 | // Preallocate our initial slice 47 | init := make([]int64, 0, len(existing)) 48 | for _, element := range existing { 49 | init = append(init, element) 50 | } 51 | } 52 | } 53 | ``` 54 | 55 | ```Bash 56 | $ go test -bench=. -benchmem 57 | goos: linux 58 | goarch: amd64 59 | BenchmarkNoPreallocate-4 3000000 510 ns/op 248 B/op 5 allocs/op 60 | BenchmarkPreallocate-4 20000000 111 ns/op 80 B/op 1 allocs/op 61 | ``` 62 | 63 | As you can see, not preallocating can cause a performance hit, primarily due to Go having to reallocate the underlying array. The pattern benchmarked above is common in Go: declare a slice, then write some sort of range or for loop that appends or indexes into it. The purpose of this tool is to flag slice/loop declarations like the one in `BenchmarkNoPreallocate`. 64 | 65 | ## Example 66 | 67 | Some examples from the Go 1.9.2 source: 68 | 69 | ```Bash 70 | $ prealloc go/src/.... 71 | archive/tar/reader_test.go:854 Consider preallocating ss 72 | archive/zip/zip_test.go:201 Consider preallocating all 73 | cmd/api/goapi.go:301 Consider preallocating missing 74 | cmd/api/goapi.go:476 Consider preallocating files 75 | cmd/asm/internal/asm/endtoend_test.go:345 Consider preallocating extra 76 | cmd/cgo/main.go:60 Consider preallocating ks 77 | cmd/cgo/ast.go:149 Consider preallocating pieces 78 | cmd/compile/internal/ssa/flagalloc.go:64 Consider preallocating oldSched 79 | cmd/compile/internal/ssa/regalloc.go:719 Consider preallocating phis 80 | cmd/compile/internal/ssa/regalloc.go:718 Consider preallocating oldSched 81 | cmd/compile/internal/ssa/regalloc.go:1674 Consider preallocating oldSched 82 | cmd/compile/internal/ssa/gen/rulegen.go:145 Consider preallocating ops 83 | cmd/compile/internal/ssa/gen/rulegen.go:145 Consider preallocating ops 84 | cmd/dist/build.go:893 Consider preallocating all 85 | cmd/dist/build.go:1246 Consider preallocating plats 86 | cmd/dist/build.go:1264 Consider preallocating results 87 | cmd/dist/buildgo.go:59 Consider preallocating list 88 | cmd/doc/pkg.go:363 Consider preallocating names 89 | cmd/fix/typecheck.go:219 Consider preallocating b 90 | cmd/go/internal/base/path.go:34 Consider preallocating out 91 | cmd/go/internal/get/get.go:175 Consider preallocating out 92 | cmd/go/internal/load/pkg.go:1894 Consider preallocating dirent 93 | cmd/go/internal/work/build.go:2402 Consider preallocating absOfiles 94 | cmd/go/internal/work/build.go:2731 Consider preallocating absOfiles 95 | cmd/internal/objfile/pe.go:48 Consider preallocating syms 96 | cmd/internal/objfile/pe.go:38 Consider preallocating addrs 97 | cmd/internal/objfile/goobj.go:43 Consider preallocating syms 98 | cmd/internal/objfile/elf.go:35 Consider preallocating syms 99 | cmd/link/internal/ld/lib.go:1070 Consider preallocating argv 100 | cmd/vet/all/main.go:91 Consider preallocating pp 101 | database/sql/sql.go:66 Consider preallocating list 102 | debug/macho/file.go:506 Consider preallocating all 103 | internal/trace/order.go:55 Consider preallocating batches 104 | mime/quotedprintable/reader_test.go:191 Consider preallocating outcomes 105 | net/dnsclient_unix_test.go:954 Consider preallocating confLines 106 | net/interface_solaris.go:85 Consider preallocating ifat 107 | net/interface_linux_test.go:91 Consider preallocating ifmat4 108 | net/interface_linux_test.go:100 Consider preallocating ifmat6 109 | net/internal/socktest/switch.go:34 Consider preallocating st 110 | os/os_windows_test.go:766 Consider preallocating args 111 | runtime/pprof/internal/profile/filter.go:77 Consider preallocating lines 112 | runtime/pprof/internal/profile/profile.go:554 Consider preallocating names 113 | text/template/parse/node.go:189 Consider preallocating decl 114 | ``` 115 | 116 | ```Go 117 | // cmd/api/goapi.go:301 118 | var missing []string 119 | for feature := range optionalSet { 120 | missing = append(missing, feature) 121 | } 122 | 123 | // cmd/fix/typecheck.go:219 124 | var b []ast.Expr 125 | for _, x := range a { 126 | b = append(b, x) 127 | } 128 | 129 | // net/internal/socktest/switch.go:34 130 | var st []Stat 131 | sw.smu.RLock() 132 | for _, s := range sw.stats { 133 | ns := *s 134 | st = append(st, ns) 135 | } 136 | sw.smu.RUnlock() 137 | 138 | // cmd/api/goapi.go:301 139 | var missing []string 140 | for feature := range optionalSet { 141 | missing = append(missing, feature) 142 | } 143 | ``` 144 | 145 | Even if the size the slice is being preallocated to is small, there's still a performance gain to be had in explicitly specifying the capacity rather than leaving it up to `append` to discover that it needs to preallocate. Of course, preallocation doesn't need to be done *everywhere*. This tool's job is just to help suggest places where one should consider preallocating. 146 | 147 | ## How do I fix prealloc's suggestions? 148 | 149 | During the declaration of your slice, rather than using the zero value of the slice with `var`, initialize it with Go's built-in `make` function, passing the appropriate type and length. This length will generally be whatever you are ranging over. Fixing the examples from above would look like so: 150 | 151 | ```Go 152 | // cmd/api/goapi.go:301 153 | missing := make([]string, 0, len(optionalSet)) 154 | for feature := range optionalSet { 155 | missing = append(missing, feature) 156 | } 157 | 158 | // cmd/fix/typecheck.go:219 159 | b := make([]ast.Expr, 0, len(a)) 160 | for _, x := range a { 161 | b = append(b, x) 162 | } 163 | 164 | // net/internal/socktest/switch.go:34 165 | st := make([]Stat, 0, len(sw.stats)) 166 | sw.smu.RLock() 167 | for _, s := range sw.stats { 168 | ns := *s 169 | st = append(st, ns) 170 | } 171 | sw.smu.RUnlock() 172 | 173 | // cmd/api/goapi.go:301 174 | missing := make ([]string, 0, len(optionalSet)) 175 | for feature := range optionalSet { 176 | missing = append(missing, feature) 177 | } 178 | ``` 179 | 180 | Note: If performance is absolutely critical, it may be more efficient to use `copy` instead of `append` for larger slices. For reference, see the following benchmark: 181 | ```Go 182 | func BenchmarkSize200PreallocateCopy(b *testing.B) { 183 | existing := make([]int64, 200, 200) 184 | b.ResetTimer() 185 | for i := 0; i < b.N; i++ { 186 | // Preallocate our initial slice 187 | init := make([]int64, len(existing)) 188 | copy(init, existing) 189 | } 190 | } 191 | ``` 192 | ``` 193 | $ go test -bench=. -benchmem 194 | goos: linux 195 | goarch: amd64 196 | BenchmarkSize200NoPreallocate-4 500000 3080 ns/op 4088 B/op 9 allocs/op 197 | BenchmarkSize200Preallocate-4 1000000 1163 ns/op 1792 B/op 1 allocs/op 198 | BenchmarkSize200PreallocateCopy-4 2000000 807 ns/op 1792 B/op 1 allocs/op 199 | ``` 200 | 201 | ## TODO 202 | 203 | - Configuration on whether or not to run on test files 204 | - Support for embedded ifs (currently, prealloc will only find breaks/returns/continues/gotos if they are in a single if block, I'd like to expand this to supporting multiple if blocks in the future). 205 | - Globbing support (e.g. prealloc *.go) 206 | 207 | 208 | ## Contributing 209 | 210 | Pull requests welcome! 211 | 212 | 213 | ## Other static analysis tools 214 | 215 | If you've enjoyed prealloc, take a look at my other static analysis tools! 216 | - [nakedret](https://github.com/alexkohler/nakedret) - Finds naked returns. 217 | - [unimport](https://github.com/alexkohler/unimport) - Finds unnecessary import aliases. 218 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/alexkohler/prealloc 2 | 3 | go 1.15 4 | -------------------------------------------------------------------------------- /import.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | /* 4 | 5 | This file holds a direct copy of the import path matching code of 6 | https://github.com/golang/go/blob/master/src/cmd/go/main.go. It can be 7 | replaced when https://golang.org/issue/8768 is resolved. 8 | 9 | It has been updated to follow upstream changes in a few ways. 10 | 11 | */ 12 | 13 | import ( 14 | "fmt" 15 | "go/build" 16 | "log" 17 | "os" 18 | "path" 19 | "path/filepath" 20 | "regexp" 21 | "runtime" 22 | "strings" 23 | ) 24 | 25 | var buildContext = build.Default 26 | 27 | var ( 28 | goroot = filepath.Clean(runtime.GOROOT()) 29 | gorootSrc = filepath.Join(goroot, "src") 30 | ) 31 | 32 | // importPathsNoDotExpansion returns the import paths to use for the given 33 | // command line, but it does no ... expansion. 34 | func importPathsNoDotExpansion(args []string) []string { 35 | if len(args) == 0 { 36 | return []string{"."} 37 | } 38 | var out []string 39 | for _, a := range args { 40 | // Arguments are supposed to be import paths, but 41 | // as a courtesy to Windows developers, rewrite \ to / 42 | // in command-line arguments. Handles .\... and so on. 43 | if filepath.Separator == '\\' { 44 | a = strings.Replace(a, `\`, `/`, -1) 45 | } 46 | 47 | // Put argument in canonical form, but preserve leading ./. 48 | if strings.HasPrefix(a, "./") { 49 | a = "./" + path.Clean(a) 50 | if a == "./." { 51 | a = "." 52 | } 53 | } else { 54 | a = path.Clean(a) 55 | } 56 | if a == "all" || a == "std" { 57 | out = append(out, allPackages(a)...) 58 | continue 59 | } 60 | out = append(out, a) 61 | } 62 | return out 63 | } 64 | 65 | // importPaths returns the import paths to use for the given command line. 66 | func importPaths(args []string) []string { 67 | args = importPathsNoDotExpansion(args) 68 | var out []string 69 | for _, a := range args { 70 | if strings.Contains(a, "...") { 71 | if build.IsLocalImport(a) { 72 | out = append(out, allPackagesInFS(a)...) 73 | } else { 74 | out = append(out, allPackages(a)...) 75 | } 76 | continue 77 | } 78 | out = append(out, a) 79 | } 80 | return out 81 | } 82 | 83 | // matchPattern(pattern)(name) reports whether 84 | // name matches pattern. Pattern is a limited glob 85 | // pattern in which '...' means 'any string' and there 86 | // is no other special syntax. 87 | func matchPattern(pattern string) func(name string) bool { 88 | re := regexp.QuoteMeta(pattern) 89 | re = strings.Replace(re, `\.\.\.`, `.*`, -1) 90 | // Special case: foo/... matches foo too. 91 | if strings.HasSuffix(re, `/.*`) { 92 | re = re[:len(re)-len(`/.*`)] + `(/.*)?` 93 | } 94 | reg := regexp.MustCompile(`^` + re + `$`) 95 | return func(name string) bool { 96 | return reg.MatchString(name) 97 | } 98 | } 99 | 100 | // hasPathPrefix reports whether the path s begins with the 101 | // elements in prefix. 102 | func hasPathPrefix(s, prefix string) bool { 103 | switch { 104 | default: 105 | return false 106 | case len(s) == len(prefix): 107 | return s == prefix 108 | case len(s) > len(prefix): 109 | if prefix != "" && prefix[len(prefix)-1] == '/' { 110 | return strings.HasPrefix(s, prefix) 111 | } 112 | return s[len(prefix)] == '/' && s[:len(prefix)] == prefix 113 | } 114 | } 115 | 116 | // treeCanMatchPattern(pattern)(name) reports whether 117 | // name or children of name can possibly match pattern. 118 | // Pattern is the same limited glob accepted by matchPattern. 119 | func treeCanMatchPattern(pattern string) func(name string) bool { 120 | wildCard := false 121 | if i := strings.Index(pattern, "..."); i >= 0 { 122 | wildCard = true 123 | pattern = pattern[:i] 124 | } 125 | return func(name string) bool { 126 | return len(name) <= len(pattern) && hasPathPrefix(pattern, name) || 127 | wildCard && strings.HasPrefix(name, pattern) 128 | } 129 | } 130 | 131 | // allPackages returns all the packages that can be found 132 | // under the $GOPATH directories and $GOROOT matching pattern. 133 | // The pattern is either "all" (all packages), "std" (standard packages) 134 | // or a path including "...". 135 | func allPackages(pattern string) []string { 136 | pkgs := matchPackages(pattern) 137 | if len(pkgs) == 0 { 138 | fmt.Fprintf(os.Stderr, "warning: %q matched no packages\n", pattern) 139 | } 140 | return pkgs 141 | } 142 | 143 | func matchPackages(pattern string) []string { 144 | match := func(string) bool { return true } 145 | treeCanMatch := func(string) bool { return true } 146 | if pattern != "all" && pattern != "std" { 147 | match = matchPattern(pattern) 148 | treeCanMatch = treeCanMatchPattern(pattern) 149 | } 150 | 151 | have := map[string]bool{ 152 | "builtin": true, // ignore pseudo-package that exists only for documentation 153 | } 154 | if !buildContext.CgoEnabled { 155 | have["runtime/cgo"] = true // ignore during walk 156 | } 157 | var pkgs []string 158 | 159 | // Commands 160 | cmd := filepath.Join(goroot, "src/cmd") + string(filepath.Separator) 161 | filepath.Walk(cmd, func(path string, fi os.FileInfo, err error) error { 162 | if err != nil || !fi.IsDir() || path == cmd { 163 | return nil 164 | } 165 | name := path[len(cmd):] 166 | if !treeCanMatch(name) { 167 | return filepath.SkipDir 168 | } 169 | // Commands are all in cmd/, not in subdirectories. 170 | if strings.Contains(name, string(filepath.Separator)) { 171 | return filepath.SkipDir 172 | } 173 | 174 | // We use, e.g., cmd/gofmt as the pseudo import path for gofmt. 175 | name = "cmd/" + name 176 | if have[name] { 177 | return nil 178 | } 179 | have[name] = true 180 | if !match(name) { 181 | return nil 182 | } 183 | _, err = buildContext.ImportDir(path, 0) 184 | if err != nil { 185 | if _, noGo := err.(*build.NoGoError); !noGo { 186 | log.Print(err) 187 | } 188 | return nil 189 | } 190 | pkgs = append(pkgs, name) 191 | return nil 192 | }) 193 | 194 | for _, src := range buildContext.SrcDirs() { 195 | if (pattern == "std" || pattern == "cmd") && src != gorootSrc { 196 | continue 197 | } 198 | src = filepath.Clean(src) + string(filepath.Separator) 199 | root := src 200 | if pattern == "cmd" { 201 | root += "cmd" + string(filepath.Separator) 202 | } 203 | filepath.Walk(root, func(path string, fi os.FileInfo, err error) error { 204 | if err != nil || !fi.IsDir() || path == src { 205 | return nil 206 | } 207 | 208 | // Avoid .foo, _foo, testdata and vendor directory trees. 209 | _, elem := filepath.Split(path) 210 | if strings.HasPrefix(elem, ".") || strings.HasPrefix(elem, "_") || elem == "testdata" || elem == "vendor" { 211 | return filepath.SkipDir 212 | } 213 | 214 | name := filepath.ToSlash(path[len(src):]) 215 | if pattern == "std" && (strings.Contains(name, ".") || name == "cmd") { 216 | // The name "std" is only the standard library. 217 | // If the name is cmd, it's the root of the command tree. 218 | return filepath.SkipDir 219 | } 220 | if !treeCanMatch(name) { 221 | return filepath.SkipDir 222 | } 223 | if have[name] { 224 | return nil 225 | } 226 | have[name] = true 227 | if !match(name) { 228 | return nil 229 | } 230 | _, err = buildContext.ImportDir(path, 0) 231 | if err != nil { 232 | if _, noGo := err.(*build.NoGoError); noGo { 233 | return nil 234 | } 235 | } 236 | pkgs = append(pkgs, name) 237 | return nil 238 | }) 239 | } 240 | return pkgs 241 | } 242 | 243 | // allPackagesInFS is like allPackages but is passed a pattern 244 | // beginning ./ or ../, meaning it should scan the tree rooted 245 | // at the given directory. There are ... in the pattern too. 246 | func allPackagesInFS(pattern string) []string { 247 | pkgs := matchPackagesInFS(pattern) 248 | if len(pkgs) == 0 { 249 | fmt.Fprintf(os.Stderr, "warning: %q matched no packages\n", pattern) 250 | } 251 | return pkgs 252 | } 253 | 254 | func matchPackagesInFS(pattern string) []string { 255 | // Find directory to begin the scan. 256 | // Could be smarter but this one optimization 257 | // is enough for now, since ... is usually at the 258 | // end of a path. 259 | i := strings.Index(pattern, "...") 260 | dir, _ := path.Split(pattern[:i]) 261 | 262 | // pattern begins with ./ or ../. 263 | // path.Clean will discard the ./ but not the ../. 264 | // We need to preserve the ./ for pattern matching 265 | // and in the returned import paths. 266 | prefix := "" 267 | if strings.HasPrefix(pattern, "./") { 268 | prefix = "./" 269 | } 270 | match := matchPattern(pattern) 271 | 272 | var pkgs []string 273 | filepath.Walk(dir, func(path string, fi os.FileInfo, err error) error { 274 | if err != nil || !fi.IsDir() { 275 | return nil 276 | } 277 | if path == dir { 278 | // filepath.Walk starts at dir and recurses. For the recursive case, 279 | // the path is the result of filepath.Join, which calls filepath.Clean. 280 | // The initial case is not Cleaned, though, so we do this explicitly. 281 | // 282 | // This converts a path like "./io/" to "io". Without this step, running 283 | // "cd $GOROOT/src/pkg; go list ./io/..." would incorrectly skip the io 284 | // package, because prepending the prefix "./" to the unclean path would 285 | // result in "././io", and match("././io") returns false. 286 | path = filepath.Clean(path) 287 | } 288 | 289 | // Avoid .foo, _foo, testdata and vendor directory trees, but do not avoid "." or "..". 290 | _, elem := filepath.Split(path) 291 | dot := strings.HasPrefix(elem, ".") && elem != "." && elem != ".." 292 | if dot || strings.HasPrefix(elem, "_") || elem == "testdata" || elem == "vendor" { 293 | return filepath.SkipDir 294 | } 295 | 296 | name := prefix + filepath.ToSlash(path) 297 | if !match(name) { 298 | return nil 299 | } 300 | if _, err = build.ImportDir(path, 0); err != nil { 301 | if _, noGo := err.(*build.NoGoError); !noGo { 302 | log.Print(err) 303 | } 304 | return nil 305 | } 306 | pkgs = append(pkgs, name) 307 | return nil 308 | }) 309 | return pkgs 310 | } 311 | -------------------------------------------------------------------------------- /pkg/prealloc.go: -------------------------------------------------------------------------------- 1 | package pkg 2 | 3 | import ( 4 | "fmt" 5 | "go/ast" 6 | "go/token" 7 | ) 8 | 9 | type sliceDeclaration struct { 10 | name string 11 | // sType string 12 | genD *ast.GenDecl 13 | } 14 | 15 | type returnsVisitor struct { 16 | // flags 17 | simple bool 18 | includeRangeLoops bool 19 | includeForLoops bool 20 | // visitor fields 21 | sliceDeclarations []*sliceDeclaration 22 | preallocHints []Hint 23 | returnsInsideOfLoop bool 24 | arrayTypes []string 25 | } 26 | 27 | func Check(files []*ast.File, simple, includeRangeLoops, includeForLoops bool) []Hint { 28 | hints := []Hint{} 29 | for _, f := range files { 30 | retVis := &returnsVisitor{ 31 | simple: simple, 32 | includeRangeLoops: includeRangeLoops, 33 | includeForLoops: includeForLoops, 34 | } 35 | ast.Walk(retVis, f) 36 | // if simple is true, then we actually have to check if we had returns 37 | // inside of our loop. Otherwise, we can just report all messages. 38 | if !retVis.simple || !retVis.returnsInsideOfLoop { 39 | hints = append(hints, retVis.preallocHints...) 40 | } 41 | } 42 | 43 | return hints 44 | } 45 | 46 | func contains(slice []string, item string) bool { 47 | for _, s := range slice { 48 | if s == item { 49 | return true 50 | } 51 | } 52 | 53 | return false 54 | } 55 | 56 | func (v *returnsVisitor) Visit(node ast.Node) ast.Visitor { 57 | 58 | v.sliceDeclarations = nil 59 | v.returnsInsideOfLoop = false 60 | 61 | switch n := node.(type) { 62 | case *ast.TypeSpec: 63 | if _, ok := n.Type.(*ast.ArrayType); ok { 64 | if n.Name != nil { 65 | v.arrayTypes = append(v.arrayTypes, n.Name.Name) 66 | } 67 | } 68 | case *ast.FuncDecl: 69 | if n.Body != nil { 70 | for _, stmt := range n.Body.List { 71 | switch s := stmt.(type) { 72 | // Find non pre-allocated slices 73 | case *ast.DeclStmt: 74 | genD, ok := s.Decl.(*ast.GenDecl) 75 | if !ok { 76 | continue 77 | } 78 | if genD.Tok == token.TYPE { 79 | for _, spec := range genD.Specs { 80 | tSpec, ok := spec.(*ast.TypeSpec) 81 | if !ok { 82 | continue 83 | } 84 | 85 | if _, ok := tSpec.Type.(*ast.ArrayType); ok { 86 | if tSpec.Name != nil { 87 | v.arrayTypes = append(v.arrayTypes, tSpec.Name.Name) 88 | } 89 | } 90 | } 91 | } else if genD.Tok == token.VAR { 92 | for _, spec := range genD.Specs { 93 | vSpec, ok := spec.(*ast.ValueSpec) 94 | if !ok { 95 | continue 96 | } 97 | var isArrType bool 98 | switch val := vSpec.Type.(type) { 99 | case *ast.ArrayType: 100 | isArrType = true 101 | case *ast.Ident: 102 | isArrType = contains(v.arrayTypes, val.Name) 103 | } 104 | if isArrType { 105 | if vSpec.Names != nil { 106 | /*atID, ok := arrayType.Elt.(*ast.Ident) 107 | if !ok { 108 | continue 109 | }*/ 110 | 111 | // We should handle multiple slices declared on same line e.g. var mySlice1, mySlice2 []uint32 112 | for _, vName := range vSpec.Names { 113 | v.sliceDeclarations = append(v.sliceDeclarations, &sliceDeclaration{name: vName.Name /*sType: atID.Name,*/, genD: genD}) 114 | } 115 | } 116 | } 117 | } 118 | } 119 | 120 | case *ast.RangeStmt: 121 | if v.includeRangeLoops { 122 | if len(v.sliceDeclarations) == 0 { 123 | continue 124 | } 125 | // Check the value being ranged over and ensure it's not a channel (we cannot offer any recommendations on channel ranges). 126 | rangeIdent, ok := s.X.(*ast.Ident) 127 | if ok && rangeIdent.Obj != nil { 128 | valueSpec, ok := rangeIdent.Obj.Decl.(*ast.ValueSpec) 129 | if ok { 130 | if _, rangeTargetIsChannel := valueSpec.Type.(*ast.ChanType); rangeTargetIsChannel { 131 | continue 132 | } 133 | } 134 | } 135 | if s.Body != nil { 136 | v.handleLoops(s.Body) 137 | } 138 | } 139 | 140 | case *ast.ForStmt: 141 | if v.includeForLoops { 142 | if len(v.sliceDeclarations) == 0 { 143 | continue 144 | } 145 | if s.Body != nil { 146 | v.handleLoops(s.Body) 147 | } 148 | } 149 | 150 | default: 151 | } 152 | } 153 | } 154 | } 155 | return v 156 | } 157 | 158 | // handleLoops is a helper function to share the logic required for both *ast.RangeLoops and *ast.ForLoops 159 | func (v *returnsVisitor) handleLoops(blockStmt *ast.BlockStmt) { 160 | 161 | for _, stmt := range blockStmt.List { 162 | switch bodyStmt := stmt.(type) { 163 | case *ast.AssignStmt: 164 | asgnStmt := bodyStmt 165 | for index, expr := range asgnStmt.Rhs { 166 | if index >= len(asgnStmt.Lhs) { 167 | continue 168 | } 169 | 170 | lhsIdent, ok := asgnStmt.Lhs[index].(*ast.Ident) 171 | if !ok { 172 | continue 173 | } 174 | 175 | callExpr, ok := expr.(*ast.CallExpr) 176 | if !ok { 177 | continue 178 | } 179 | 180 | rhsFuncIdent, ok := callExpr.Fun.(*ast.Ident) 181 | if !ok { 182 | continue 183 | } 184 | 185 | if rhsFuncIdent.Name != "append" { 186 | continue 187 | } 188 | 189 | // e.g., `x = append(x)` 190 | // Pointless, but pre-allocation will not help. 191 | if len(callExpr.Args) < 2 { 192 | continue 193 | } 194 | 195 | rhsIdent, ok := callExpr.Args[0].(*ast.Ident) 196 | if !ok { 197 | continue 198 | } 199 | 200 | // e.g., `x = append(y, a)` 201 | // This is weird (and maybe a logic error), 202 | // but we cannot recommend pre-allocation. 203 | if lhsIdent.Name != rhsIdent.Name { 204 | continue 205 | } 206 | 207 | // e.g., `x = append(x, y...)` 208 | // we should ignore this. Pre-allocating in this case 209 | // is confusing, and is not possible in general. 210 | if callExpr.Ellipsis.IsValid() { 211 | continue 212 | } 213 | 214 | for _, sliceDecl := range v.sliceDeclarations { 215 | if sliceDecl.name == lhsIdent.Name { 216 | // This is a potential mark, we just need to make sure there are no returns/continues in the 217 | // range loop. 218 | // now we just need to grab whatever we're ranging over 219 | /*sxIdent, ok := s.X.(*ast.Ident) 220 | if !ok { 221 | continue 222 | }*/ 223 | 224 | v.preallocHints = append(v.preallocHints, Hint{ 225 | Pos: sliceDecl.genD.Pos(), 226 | DeclaredSliceName: sliceDecl.name, 227 | }) 228 | } 229 | } 230 | } 231 | case *ast.IfStmt: 232 | ifStmt := bodyStmt 233 | if ifStmt.Body != nil { 234 | for _, ifBodyStmt := range ifStmt.Body.List { 235 | // TODO should probably handle embedded ifs here 236 | switch /*ift :=*/ ifBodyStmt.(type) { 237 | case *ast.BranchStmt, *ast.ReturnStmt: 238 | v.returnsInsideOfLoop = true 239 | default: 240 | } 241 | } 242 | } 243 | 244 | default: 245 | 246 | } 247 | } 248 | 249 | } 250 | 251 | // Hint stores the information about an occurrence of a slice that could be 252 | // preallocated. 253 | type Hint struct { 254 | Pos token.Pos 255 | DeclaredSliceName string 256 | } 257 | 258 | func (h Hint) String() string { 259 | return fmt.Sprintf("%v: Consider preallocating %v", h.Pos, h.DeclaredSliceName) 260 | } 261 | 262 | func (h Hint) StringFromFS(f *token.FileSet) string { 263 | file := f.File(h.Pos) 264 | lineNumber := file.Position(h.Pos).Line 265 | 266 | return fmt.Sprintf("%v:%v Consider preallocating %v", file.Name(), lineNumber, h.DeclaredSliceName) 267 | } 268 | -------------------------------------------------------------------------------- /prealloc.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "fmt" 6 | "go/ast" 7 | "go/build" 8 | "go/parser" 9 | "go/token" 10 | "log" 11 | "os" 12 | "path/filepath" 13 | "strings" 14 | 15 | "github.com/alexkohler/prealloc/pkg" 16 | ) 17 | 18 | // Support: (in order of priority) 19 | // * Full make suggestion with type? 20 | // * Test flag 21 | // * Embedded ifs? 22 | // * Use an import rather than the duplcated import.go 23 | 24 | const ( 25 | pwd = "./" 26 | ) 27 | 28 | func init() { 29 | // Ignore build flags 30 | build.Default.UseAllFiles = true 31 | } 32 | 33 | func usage() { 34 | log.Printf("Usage of %s:\n", os.Args[0]) 35 | log.Printf("\nprealloc [flags] # runs on package in current directory\n") 36 | log.Printf("\nprealloc [flags] [packages]\n") 37 | log.Printf("Flags:\n") 38 | flag.PrintDefaults() 39 | } 40 | 41 | func main() { 42 | 43 | // Remove log timestamp 44 | log.SetFlags(0) 45 | 46 | simple := flag.Bool("simple", true, "Report preallocation suggestions only on simple loops that have no returns/breaks/continues/gotos in them") 47 | includeRangeLoops := flag.Bool("rangeloops", true, "Report preallocation suggestions on range loops") 48 | includeForLoops := flag.Bool("forloops", false, "Report preallocation suggestions on for loops") 49 | setExitStatus := flag.Bool("set_exit_status", false, "Set exit status to 1 if any issues are found") 50 | flag.Usage = usage 51 | flag.Parse() 52 | 53 | fset := token.NewFileSet() 54 | 55 | hints, err := checkForPreallocations( 56 | flag.Args(), 57 | fset, 58 | *simple, 59 | *includeRangeLoops, 60 | *includeForLoops, 61 | ) 62 | if err != nil { 63 | log.Println(err) 64 | } 65 | 66 | for _, hint := range hints { 67 | log.Println(hint.StringFromFS(fset)) 68 | } 69 | if *setExitStatus && len(hints) > 0 { 70 | os.Exit(1) 71 | } 72 | } 73 | 74 | func checkForPreallocations( 75 | args []string, 76 | fset *token.FileSet, 77 | simple, includeRangeLoops, includeForLoops bool, 78 | ) ([]pkg.Hint, error) { 79 | 80 | files, err := parseInput(args, fset) 81 | if err != nil { 82 | return nil, fmt.Errorf("could not parse input %v", err) 83 | } 84 | 85 | hints := pkg.Check(files, simple, includeRangeLoops, includeForLoops) 86 | 87 | return hints, nil 88 | } 89 | 90 | func parseInput(args []string, fset *token.FileSet) ([]*ast.File, error) { 91 | var directoryList []string 92 | var fileMode bool 93 | files := make([]*ast.File, 0) 94 | 95 | if len(args) == 0 { 96 | directoryList = append(directoryList, pwd) 97 | } else { 98 | for _, arg := range args { 99 | if strings.HasSuffix(arg, "/...") && isDir(arg[:len(arg)-len("/...")]) { 100 | 101 | for _, dirname := range allPackagesInFS(arg) { 102 | directoryList = append(directoryList, dirname) 103 | } 104 | 105 | } else if isDir(arg) { 106 | directoryList = append(directoryList, arg) 107 | 108 | } else if exists(arg) { 109 | if strings.HasSuffix(arg, ".go") { 110 | fileMode = true 111 | f, err := parser.ParseFile(fset, arg, nil, 0) 112 | if err != nil { 113 | return nil, err 114 | } 115 | files = append(files, f) 116 | } else { 117 | return nil, fmt.Errorf("invalid file %v specified", arg) 118 | } 119 | } else { 120 | 121 | //TODO clean this up a bit 122 | imPaths := importPaths([]string{arg}) 123 | for _, importPath := range imPaths { 124 | pkg, err := build.Import(importPath, ".", 0) 125 | if err != nil { 126 | return nil, err 127 | } 128 | var stringFiles []string 129 | stringFiles = append(stringFiles, pkg.GoFiles...) 130 | // files = append(files, pkg.CgoFiles...) 131 | stringFiles = append(stringFiles, pkg.TestGoFiles...) 132 | if pkg.Dir != "." { 133 | for i, f := range stringFiles { 134 | stringFiles[i] = filepath.Join(pkg.Dir, f) 135 | } 136 | } 137 | 138 | fileMode = true 139 | for _, stringFile := range stringFiles { 140 | f, err := parser.ParseFile(fset, stringFile, nil, 0) 141 | if err != nil { 142 | return nil, err 143 | } 144 | files = append(files, f) 145 | } 146 | 147 | } 148 | } 149 | } 150 | } 151 | 152 | // if we're not in file mode, then we need to grab each and every package in each directory 153 | // we can to grab all the files 154 | if !fileMode { 155 | for _, fpath := range directoryList { 156 | pkgs, err := parser.ParseDir(fset, fpath, nil, 0) 157 | if err != nil { 158 | return nil, err 159 | } 160 | 161 | for _, pkg := range pkgs { 162 | for _, f := range pkg.Files { 163 | files = append(files, f) 164 | } 165 | } 166 | } 167 | } 168 | 169 | return files, nil 170 | } 171 | 172 | func isDir(filename string) bool { 173 | fi, err := os.Stat(filename) 174 | return err == nil && fi.IsDir() 175 | } 176 | 177 | func exists(filename string) bool { 178 | _, err := os.Stat(filename) 179 | return err == nil 180 | } 181 | -------------------------------------------------------------------------------- /prealloc_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "go/token" 6 | "testing" 7 | 8 | "github.com/alexkohler/prealloc/pkg" 9 | ) 10 | 11 | func Test_checkForPreallocations(t *testing.T) { 12 | const filename = "testdata/sample.go" 13 | 14 | fset := token.NewFileSet() 15 | 16 | got, err := checkForPreallocations([]string{filename}, fset, true, true, true) 17 | if err != nil { 18 | t.Fatal(err) 19 | } 20 | 21 | want := []pkg.Hint{ 22 | pkg.Hint{ 23 | Pos: 63, 24 | DeclaredSliceName: "y", 25 | }, 26 | pkg.Hint{ 27 | Pos: 77, 28 | DeclaredSliceName: "z", 29 | }, 30 | pkg.Hint{ 31 | Pos: 102, 32 | DeclaredSliceName: "t", 33 | }, 34 | } 35 | 36 | if len(got) != len(want) { 37 | t.Fatalf("expected %d hints, but got %d: %+v", len(want), len(got), got) 38 | } 39 | 40 | for i := range got { 41 | act, exp := got[i], want[i] 42 | 43 | file := fset.File(act.Pos) 44 | 45 | if file.Name() != filename { 46 | t.Errorf("wrong hints[%d].Filename: %q (expected: %q)", i, file.Name(), filename) 47 | } 48 | 49 | actLineNumber := file.Position(act.Pos).Line 50 | expLineNumber := file.Position(exp.Pos).Line 51 | 52 | if actLineNumber != expLineNumber { 53 | t.Errorf("wrong hints[%d].LineNumber: %d (expected: %d)", i, actLineNumber, expLineNumber) 54 | } 55 | 56 | if act.DeclaredSliceName != exp.DeclaredSliceName { 57 | t.Errorf("wrong hints[%d].DeclaredSliceName: %q (expected: %q)", i, act.DeclaredSliceName, exp.DeclaredSliceName) 58 | } 59 | } 60 | } 61 | 62 | func BenchmarkSize10NoPreallocate(b *testing.B) { 63 | existing := make([]int64, 10, 10) 64 | b.ResetTimer() 65 | for i := 0; i < b.N; i++ { 66 | // Don't preallocate our initial slice 67 | var init []int64 68 | for _, element := range existing { 69 | init = append(init, element) 70 | } 71 | } 72 | } 73 | 74 | func BenchmarkSize10Preallocate(b *testing.B) { 75 | existing := make([]int64, 10, 10) 76 | b.ResetTimer() 77 | for i := 0; i < b.N; i++ { 78 | // Preallocate our initial slice 79 | init := make([]int64, 0, len(existing)) 80 | for _, element := range existing { 81 | init = append(init, element) 82 | } 83 | } 84 | } 85 | 86 | func BenchmarkSize10PreallocateCopy(b *testing.B) { 87 | existing := make([]int64, 10, 10) 88 | b.ResetTimer() 89 | for i := 0; i < b.N; i++ { 90 | // Preallocate our initial slice 91 | init := make([]int64, len(existing)) 92 | copy(init, existing) 93 | } 94 | } 95 | 96 | func BenchmarkSize200NoPreallocate(b *testing.B) { 97 | existing := make([]int64, 200, 200) 98 | b.ResetTimer() 99 | for i := 0; i < b.N; i++ { 100 | // Don't preallocate our initial slice 101 | var init []int64 102 | for _, element := range existing { 103 | init = append(init, element) 104 | } 105 | } 106 | } 107 | 108 | func BenchmarkSize200Preallocate(b *testing.B) { 109 | existing := make([]int64, 200, 200) 110 | b.ResetTimer() 111 | for i := 0; i < b.N; i++ { 112 | // Preallocate our initial slice 113 | init := make([]int64, 0, len(existing)) 114 | for _, element := range existing { 115 | init = append(init, element) 116 | } 117 | } 118 | } 119 | 120 | func BenchmarkSize200PreallocateCopy(b *testing.B) { 121 | existing := make([]int64, 200, 200) 122 | b.ResetTimer() 123 | for i := 0; i < b.N; i++ { 124 | // Preallocate our initial slice 125 | init := make([]int64, len(existing)) 126 | copy(init, existing) 127 | } 128 | } 129 | 130 | func BenchmarkMap(b *testing.B) { 131 | benchmarks := []struct { 132 | size int 133 | preallocate bool 134 | }{ 135 | {10, false}, 136 | {10, true}, 137 | {200, false}, 138 | {200, true}, 139 | } 140 | var m map[int]int 141 | for _, bm := range benchmarks { 142 | no := "" 143 | if !bm.preallocate { 144 | no = "No" 145 | } 146 | b.Run(fmt.Sprintf("Size%d%sPreallocate", bm.size, no), func(b *testing.B) { 147 | for i := 0; i < b.N; i++ { 148 | if bm.preallocate { 149 | m = make(map[int]int, bm.size) 150 | } else { 151 | m = make(map[int]int) 152 | } 153 | for j := 0; j < bm.size; j++ { 154 | m[j] = j 155 | } 156 | } 157 | }) 158 | } 159 | } 160 | -------------------------------------------------------------------------------- /testdata/sample.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | func main() { 4 | x := make([]rune, len("Hello")) 5 | var y []rune 6 | var z, w, v, u, s []int 7 | var t [][]int 8 | var intChan chan int 9 | 10 | for i, r := range "Hello" { 11 | // x is already pre-allocated 12 | // y is a candidate for pre-allocation 13 | x[i], y = r, append(y, r) 14 | 15 | // w is not a candidate for pre-allocation due to `...` 16 | w = append(w, foo(i)...) 17 | 18 | // v is not a candidate for pre-allocation since this appends to u 19 | v = append(u, i) 20 | 21 | // u is not a candidate for pre-allocation since nothing was actually appended 22 | u = append(u) 23 | 24 | // z is a candidate for pre-allocation 25 | z = append(z, i) 26 | 27 | // t is a candidate for pre-allocation 28 | t = append(t, foo(i)) 29 | } 30 | 31 | for i := range intChan { 32 | // s is not a candidate for pre-allocation since the range target is a channel 33 | s = append(s, i) 34 | } 35 | 36 | _ = v 37 | } 38 | 39 | func foo(n int) []int { 40 | return make([]int, n) 41 | } 42 | --------------------------------------------------------------------------------