cmd/cover: cover funcs in if, for, switch clauses

This peculiar case arose in range statements but there are other contexts
and one turned up in the auto-generated translation of the compiler.
Take care of it always.

	for i := 0; i < 0; func() {i++; q=q.Link}() { ... }

That code has been given the obvious rewrite but we should still handle it.

Odd but easy to fix (tricky to test).

Fixes #10269.

Change-Id: I66e1404eb24da15a24be7f67403e19ed66fba0a7
Reviewed-on: https://go-review.googlesource.com/8284
Reviewed-by: Robert Griesemer <gri@golang.org>
diff --git a/cmd/cover/cover.go b/cmd/cover/cover.go
index e6e7b46..31ec434 100644
--- a/cmd/cover/cover.go
+++ b/cmd/cover/cover.go
@@ -503,6 +503,9 @@
 // Therefore we draw a line at the start of the body of the first function literal we find.
 // TODO: what if there's more than one? Probably doesn't matter much.
 func hasFuncLiteral(n ast.Node) (bool, token.Pos) {
+	if n == nil {
+		return false, 0
+	}
 	var literal funcLitFinder
 	ast.Walk(&literal, n)
 	return literal.found(), token.Pos(literal)
@@ -517,24 +520,54 @@
 		// Treat blocks like basic blocks to avoid overlapping counters.
 		return s.Lbrace
 	case *ast.IfStmt:
+		found, pos := hasFuncLiteral(s.Init)
+		if found {
+			return pos
+		}
+		found, pos = hasFuncLiteral(s.Cond)
+		if found {
+			return pos
+		}
 		return s.Body.Lbrace
 	case *ast.ForStmt:
+		found, pos := hasFuncLiteral(s.Init)
+		if found {
+			return pos
+		}
+		found, pos = hasFuncLiteral(s.Cond)
+		if found {
+			return pos
+		}
+		found, pos = hasFuncLiteral(s.Post)
+		if found {
+			return pos
+		}
 		return s.Body.Lbrace
 	case *ast.LabeledStmt:
 		return f.statementBoundary(s.Stmt)
 	case *ast.RangeStmt:
-		// Ranges might loop over things with function literals.: for _ = range []func(){ ... } {.
-		// TODO: There are a few other such possibilities, but they're extremely unlikely.
 		found, pos := hasFuncLiteral(s.X)
 		if found {
 			return pos
 		}
 		return s.Body.Lbrace
 	case *ast.SwitchStmt:
+		found, pos := hasFuncLiteral(s.Init)
+		if found {
+			return pos
+		}
+		found, pos = hasFuncLiteral(s.Tag)
+		if found {
+			return pos
+		}
 		return s.Body.Lbrace
 	case *ast.SelectStmt:
 		return s.Body.Lbrace
 	case *ast.TypeSwitchStmt:
+		found, pos := hasFuncLiteral(s.Init)
+		if found {
+			return pos
+		}
 		return s.Body.Lbrace
 	}
 	// If not a control flow statement, it is a declaration, expression, call, etc. and it may have a function literal.
diff --git a/cmd/cover/testdata/test.go b/cmd/cover/testdata/test.go
index 16c349c..9013950 100644
--- a/cmd/cover/testdata/test.go
+++ b/cmd/cover/testdata/test.go
@@ -86,10 +86,13 @@
 			check(LINE, 0)
 		}
 	}
+	if func(a, b int) bool { return a < b }(3, 4) {
+		check(LINE, 1)
+	}
 }
 
 func testFor() {
-	for i := 0; i < 10; i++ {
+	for i := 0; i < 10; func() { i++; check(LINE, 10) }() {
 		check(LINE, 10)
 	}
 }
@@ -122,7 +125,7 @@
 }
 
 func testSwitch() {
-	for i := 0; i < 5; i++ {
+	for i := 0; i < 5; func() { i++; check(LINE, 5) }() {
 		switch i {
 		case 0:
 			check(LINE, 1)
@@ -139,7 +142,7 @@
 func testTypeSwitch() {
 	var x = []interface{}{1, 2.0, "hi"}
 	for _, v := range x {
-		switch v.(type) {
+		switch func() { check(LINE, 3) }(); v.(type) {
 		case int:
 			check(LINE, 1)
 		case float64: