From a8399305069021ee83a3ab603bfe1da98598e0b6 Mon Sep 17 00:00:00 2001
From: Nathan Baulch <nathan.baulch@gmail.com>
Date: Tue, 21 Oct 2025 21:24:08 +1100
Subject: [PATCH] Search within nested statement blocks

---
 pkg/prealloc.go    |  6 +++---
 prealloc_test.go   | 12 ++++++++++++
 testdata/sample.go | 25 +++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/pkg/prealloc.go b/pkg/prealloc.go
index 72d8b95..e2449fc 100644
--- a/pkg/prealloc.go
+++ b/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:
diff --git a/prealloc_test.go b/prealloc_test.go
index e55cfca..558784a 100644
--- a/prealloc_test.go
+++ b/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) {
diff --git a/testdata/sample.go b/testdata/sample.go
index 8cf72ff..849a823 100644
--- a/testdata/sample.go
+++ b/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 {