Find variable declarations in function closure

Treat these as shadowed variables and ignore them.

1 files changed, 53 insertions, 21 deletions

diff --git a/gocapturedrefrace.go b/gocapturedrefrace.go
index fd2cc94..cab0b4f 100644
--- a/gocapturedrefrace.go
+++ b/gocapturedrefrace.go
@@ -86,7 +86,7 @@ func checkClosure(pass *analysis.Pass, funcLit *ast.FuncLit) {
 
 	// Build a list of assignments local to funcLit. These will be ignored as
 	// shadowed variables.
-	localAssignments := findLocalAssignments(pass, funcLit)
+	localAssignments := findLocalVarDeclarations(pass, funcLit)
 
 	ast.Inspect(
 		funcLit,
@@ -143,41 +143,73 @@ func checkClosure(pass *analysis.Pass, funcLit *ast.FuncLit) {
 	)
 }
 
-// findLocalAssignments returns a list of all variable definitions in funcLit.
-func findLocalAssignments(
+// findLocalVarDeclarations returns a list of all variable declarations in
+// funcLit.
+func findLocalVarDeclarations(
 	pass *analysis.Pass,
 	funcLit *ast.FuncLit,
-) (localAssignments []*ast.Ident) {
-	localAssignments = []*ast.Ident{}
+) (declarations []*ast.Ident) {
+	declarations = []*ast.Ident{}
 
 	ast.Inspect(
 		funcLit,
 		func(node ast.Node) bool {
-			assignStmt, ok := node.(*ast.AssignStmt)
-			if !ok {
-				return true
-			}
-
-			if assignStmt.Tok != token.DEFINE {
-				return true
-			}
-
-			for _, lhs := range assignStmt.Lhs {
-				ident, ok := lhs.(*ast.Ident)
-				if !ok {
+			switch node := node.(type) {
+			case *ast.AssignStmt:
+				// assignStmt, ok := node.(*ast.AssignStmt)
+				// if !ok {
+				// 	return true
+				// }
+				assignStmt := node
+
+				if assignStmt.Tok != token.DEFINE {
 					return true
 				}
 
-				if ident == nil {
-					return true
+				for _, lhs := range assignStmt.Lhs {
+					ident, ok := lhs.(*ast.Ident)
+					if !ok {
+						return true
+					}
+
+					if ident == nil {
+						return true
+					}
+
+					declarations = append(declarations, ident)
 				}
 
-				localAssignments = append(localAssignments, ident)
+			case *ast.GenDecl:
+				decl := varDeclaration(node)
+				if decl != nil {
+					declarations = append(declarations, decl)
+				}
 			}
 
 			return true
 		},
 	)
 
-	return localAssignments
+	return declarations
+}
+
+// varDeclaration returns the identifier corresponding to variable declarations
+// in decl, or nil if decl is not a variable declaration.
+func varDeclaration(decl *ast.GenDecl) *ast.Ident {
+	if decl.Tok != token.VAR {
+		return nil
+	}
+
+	for _, spec := range decl.Specs {
+		valueSpec, ok := spec.(*ast.ValueSpec)
+		if !ok {
+			return nil
+		}
+
+		for _, ident := range valueSpec.Names {
+			return ident
+		}
+	}
+
+	return nil
 }