Search within nested statement blocks

pkg/prealloc.go

@@ -65,9 +65,9 @@ func (v *returnsVisitor) Visit(node ast.Node) ast.Visitor {
 				v.arrayTypes = append(v.arrayTypes, n.Name.Name)
 			}
 		}
-	case *ast.FuncDecl:
-		if n.Body != nil {
-			for _, stmt := range n.Body.List {
+	case *ast.BlockStmt:
+		if n.List != nil {
+			for _, stmt := range n.List {
 				switch s := stmt.(type) {
 				// Find non pre-allocated slices
 				case *ast.DeclStmt:

prealloc_test.go

@@ -31,6 +31,18 @@ func Test_checkForPreallocations(t *testing.T) {
 			Pos:               102,
 			DeclaredSliceName: "t",
 		},
+		pkg.Hint{
+			Pos:               820,
+			DeclaredSliceName: "m",
+		},
+		pkg.Hint{
+			Pos:               936,
+			DeclaredSliceName: "n",
+		},
+		pkg.Hint{
+			Pos:               1062,
+			DeclaredSliceName: "o",
+		},
 	}
 
 	if len(got) != len(want) {

testdata/sample.go

@@ -34,6 +34,31 @@ func main() {
 	}
 
 	_ = v
+
+	{
+		var m []int
+		for i := range "Hello" {
+			// m is a candidate for preallocation
+			m = append(m, i)
+		}
+
+		if true {
+			var n []int
+			for i := range "Hello" {
+				// n is a candidate for preallocation
+				n = append(n, i)
+			}
+
+			for {
+				var o []int
+				for i := range "Hello" {
+					// o is a candidate for preallocation
+					o = append(o, i)
+				}
+				break
+			}
+		}
+	}
 }
 
 func foo(n int) []int {