Merge pull request #9 from MichaelUrman/contained

Test and fix cases with nested contexts

Author: Crocmagnon

pkg/analyzer/analyzer.go

@@ -7,6 +7,7 @@ import (
 	"go/ast"
 	"go/printer"
 	"go/token"
+
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/inspect"
 	"golang.org/x/tools/go/ast/inspector"
@@ -54,6 +55,15 @@ func run(pass *analysis.Pass) (interface{}, error) {
 				break
 			}
 
+			// 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 obj.Pos() >= body.Pos() && obj.Pos() < body.End() {
+						continue // definition is within the loop
+					}
+				}
+			}
+
 			suggestedStmt := ast.AssignStmt{
 				Lhs:    assignStmt.Lhs,
 				TokPos: assignStmt.TokPos,
@@ -103,6 +113,24 @@ func getBody(node ast.Node) (*ast.BlockStmt, error) {
 	return nil, errUnknown
 }
 
+func getRootIdent(pass *analysis.Pass, node ast.Node) *ast.Ident {
+	for {
+		switch n := node.(type) {
+		case *ast.Ident:
+			return n
+		case *ast.IndexExpr:
+			node = n.X
+		case *ast.SelectorExpr:
+			if sel, ok := pass.TypesInfo.Selections[n]; ok && sel.Indirect() {
+				return nil // indirected (pointer) roots don't imply a (safe) copy
+			}
+			node = n.X
+		default:
+			return nil
+		}
+	}
+}
+
 // render returns the pretty-print of the given node
 func render(fset *token.FileSet, x interface{}) (string, error) {
 	var buf bytes.Buffer

testdata/src/example.go

@@ -30,3 +30,44 @@ func example() {
 func wrapContext(ctx context.Context) context.Context {
 	return context.WithoutCancel(ctx)
 }
+
+// storing contexts in a struct isn't recommended, but local copies of a non-pointer struct should act like local copies of a context.
+func inStructs(ctx context.Context) {
+	for i := 0; i < 10; i++ {
+		c := struct{ Ctx context.Context }{ctx}
+		c.Ctx = context.WithValue(c.Ctx, "key", i)
+		c.Ctx = context.WithValue(c.Ctx, "other", "val")
+	}
+
+	for i := 0; i < 10; i++ {
+		c := []struct{ Ctx context.Context }{{ctx}}
+		c[0].Ctx = context.WithValue(c[0].Ctx, "key", i)
+		c[0].Ctx = context.WithValue(c[0].Ctx, "other", "val")
+	}
+
+	c := struct{ Ctx context.Context }{ctx}
+	for i := 0; i < 10; i++ {
+		c := c
+		c.Ctx = context.WithValue(c.Ctx, "key", i)
+		c.Ctx = context.WithValue(c.Ctx, "other", "val")
+	}
+
+	pc := &struct{ Ctx context.Context }{ctx}
+	for i := 0; i < 10; i++ {
+		c := pc
+		c.Ctx = context.WithValue(c.Ctx, "key", i) // want "nested context in loop"
+		c.Ctx = context.WithValue(c.Ctx, "other", "val")
+	}
+
+	r := []struct{ Ctx context.Context }{{ctx}}
+	for i := 0; i < 10; i++ {
+		r[0].Ctx = context.WithValue(r[0].Ctx, "key", i) // want "nested context in loop"
+		r[0].Ctx = context.WithValue(r[0].Ctx, "other", "val")
+	}
+
+	rp := []*struct{ Ctx context.Context }{{ctx}}
+	for i := 0; i < 10; i++ {
+		rp[0].Ctx = context.WithValue(rp[0].Ctx, "key", i) // want "nested context in loop"
+		rp[0].Ctx = context.WithValue(rp[0].Ctx, "other", "val")
+	}
+}