From c2c0e62d59d83c811c9ac863693006cd3e09c7e9 Mon Sep 17 00:00:00 2001
From: Venkatesh Kotwade <venkateshkotwade@gmail.com>
Date: Mon, 26 Aug 2024 01:41:16 +0530
Subject: [PATCH] feat: also added support for multiple contexts

---
 pkg/analyzer/analyzer.go | 10 +++++-----
 testdata/src/example.go  |  8 ++++++++
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/pkg/analyzer/analyzer.go b/pkg/analyzer/analyzer.go
index 611f303..5cd7108 100644
--- a/pkg/analyzer/analyzer.go
+++ b/pkg/analyzer/analyzer.go
@@ -110,21 +110,21 @@ func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *as
 	for _, stmt := range stmts {
 		// Recurse if necessary
 		if inner, ok := stmt.(*ast.BlockStmt); ok {
-			found := findNestedContext(pass, inner, inner.List)
+			found := findNestedContext(pass, node, inner.List)
 			if found != nil {
 				return found
 			}
 		}
 
 		if inner, ok := stmt.(*ast.IfStmt); ok {
-			found := findNestedContext(pass, inner, inner.Body.List)
+			found := findNestedContext(pass, node, inner.Body.List)
 			if found != nil {
 				return found
 			}
 		}
 
 		if inner, ok := stmt.(*ast.SwitchStmt); ok {
-			found := findNestedContext(pass, inner, inner.Body.List)
+			found := findNestedContext(pass, node, inner.Body.List)
 			if found != nil {
 				return found
 			}
@@ -138,7 +138,7 @@ func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *as
 		}
 
 		if inner, ok := stmt.(*ast.SelectStmt); ok {
-			found := findNestedContext(pass, inner, inner.Body.List)
+			found := findNestedContext(pass, node, inner.Body.List)
 			if found != nil {
 				return found
 			}
@@ -167,7 +167,7 @@ func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *as
 		}
 
 		if assignStmt.Tok == token.DEFINE {
-			break
+			continue
 		}
 
 		// allow assignment to non-pointer children of values defined within the loop
diff --git a/testdata/src/example.go b/testdata/src/example.go
index e72813f..565ee49 100644
--- a/testdata/src/example.go
+++ b/testdata/src/example.go
@@ -71,6 +71,14 @@ func example() {
 			ctx = wrapContext(ctx)
 		}
 	}
+
+	for {
+		ctx2 := context.Background()
+		ctx = wrapContext(ctx) // want "nested context in loop"
+		if doSomething() != nil {
+			ctx2 = wrapContext(ctx2)
+		}
+	}
 }
 
 func wrapContext(ctx context.Context) context.Context {