feat: also added support for multiple contexts

.github/workflows/go.yml

@@ -14,8 +14,9 @@ permissions:
 jobs:
   build:
     strategy:
+      fail-fast: false
       matrix:
-        go: ['1.21', '1.22', '1.23.0-rc.2']
+        go: ['1.22', '1.23']
         os: [macos-latest, windows-latest, ubuntu-latest]
     name: build
     runs-on: ${{ matrix.os }}

go.mod

@@ -1,6 +1,6 @@
 module github.com/Crocmagnon/fatcontext
 
-go 1.21
+go 1.22.6
 
 require golang.org/x/tools v0.23.0
 

pkg/analyzer/analyzer.go

@@ -38,7 +38,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 			return
 		}
 
-		assignStmt := findNestedContext(pass, body, body.List)
+		assignStmt := findNestedContext(pass, node, body.List)
 		if assignStmt == nil {
 			return
 		}
@@ -78,9 +78,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 
 func getReportMessage(node ast.Node) string {
 	switch node.(type) {
-	case *ast.ForStmt:
-		return "nested context in loop"
-	case *ast.RangeStmt:
+	case *ast.ForStmt, *ast.RangeStmt:
 		return "nested context in loop"
 	case *ast.FuncLit:
 		return "nested context in function literal"
@@ -108,46 +106,46 @@ func getBody(node ast.Node) (*ast.BlockStmt, error) {
 	return nil, errUnknown
 }
 
-func findNestedContext(pass *analysis.Pass, block *ast.BlockStmt, stmts []ast.Stmt) *ast.AssignStmt {
+func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *ast.AssignStmt {
 	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.Body, 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.Body, inner.Body.List)
+			found := findNestedContext(pass, node, inner.Body.List)
 			if found != nil {
 				return found
 			}
 		}
 
 		if inner, ok := stmt.(*ast.CaseClause); ok {
-			found := findNestedContext(pass, block, inner.Body)
+			found := findNestedContext(pass, node, inner.Body)
 			if found != nil {
 				return found
 			}
 		}
 
 		if inner, ok := stmt.(*ast.SelectStmt); ok {
-			found := findNestedContext(pass, inner.Body, inner.Body.List)
+			found := findNestedContext(pass, node, inner.Body.List)
 			if found != nil {
 				return found
 			}
 		}
 
 		if inner, ok := stmt.(*ast.CommClause); ok {
-			found := findNestedContext(pass, block, inner.Body)
+			found := findNestedContext(pass, node, inner.Body)
 			if found != nil {
 				return found
 			}
@@ -169,13 +167,13 @@ func findNestedContext(pass *analysis.Pass, block *ast.BlockStmt, stmts []ast.St
 		}
 
 		if assignStmt.Tok == token.DEFINE {
-			break
+			continue
 		}
 
 		// allow assignment to non-pointer children of values defined within the loop
 		if lhs := getRootIdent(pass, assignStmt.Lhs[0]); lhs != nil {
 			if obj := pass.TypesInfo.ObjectOf(lhs); obj != nil {
-				if checkObjectScopeWithinBlock(obj.Parent(), block) {
+				if checkObjectScopeWithinNode(obj.Parent(), node) {
 					continue // definition is within the loop
 				}
 			}
@@ -187,12 +185,12 @@ func findNestedContext(pass *analysis.Pass, block *ast.BlockStmt, stmts []ast.St
 	return nil
 }
 
-func checkObjectScopeWithinBlock(scope *types.Scope, block *ast.BlockStmt) bool {
+func checkObjectScopeWithinNode(scope *types.Scope, node ast.Node) bool {
 	if scope == nil {
 		return false
 	}
 
-	if scope.Pos() >= block.Pos() && scope.End() <= block.End() {
+	if scope.Pos() >= node.Pos() && scope.End() <= node.End() {
 		return true
 	}
 

testdata/src/example.go

@@ -65,6 +65,20 @@ func example() {
 		ctx = wrapContext(ctx) // want "nested context in function literal"
 	}
 
+	// this is fine because the context is created in the loop
+	for {
+		if ctx := context.Background(); doSomething() != nil {
+			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 {