Actually fix instrumentation of nested switch statements

instrumenter.go

@@ -298,6 +298,7 @@ func (i *instrumenter) prepareSwitchStmt(n *ast.SwitchStmt) {
 		newBody = append(newBody, gen.use(gen.ident(tagExprName)))
 	}
 	newBody = append(newBody, gen.switchStmt(nil, n.Body))
+	i.fixStmtRefs(newBody)
 
 	// The initialization statements are executed in a new scope.
 	// Use the same scope for storing the tag expression in a variable
@@ -375,6 +376,7 @@ func (i *instrumenter) prepareTypeSwitchStmt(ts *ast.TypeSwitchStmt) {
 			newBody = append(newBody, gen.use(gen.ident(tagExprName)))
 		}
 		newBody = append(newBody, clause.Body...)
+		i.fixStmtRefs(newBody)
 
 		for range clause.List {
 			test := tests[0]
@@ -401,10 +403,17 @@ func (i *instrumenter) prepareTypeSwitchStmt(ts *ast.TypeSwitchStmt) {
 	newBody = append(newBody, gen.define(tag, tagExpr.X))
 	newBody = append(newBody, assignments...)
 	newBody = append(newBody, gen.switchStmt(nil, gen.block(newClauses)))
+	i.fixStmtRefs(newBody)
 
 	i.stmtSubst[ts] = gen.block(newBody)
 }
 
+func (i *instrumenter) fixStmtRefs(stmts []ast.Stmt) {
+	for si, stmt := range stmts {
+		i.stmtRef[stmt] = &stmts[si]
+	}
+}
+
 // replace replaces each prepared node with the instrumentation code,
 // in declaration order.
 func (i *instrumenter) replace(n ast.Node) bool {

testdata/instrumenter/TypeSwitchStmt.branch

@@ -252,9 +252,12 @@ func typeSwitchStmtNested() {
 		switch {
 		case GobcoCover(22, gobco1):
 		default:
-			switch 1 + 1 {
-			case GobcoCover(23, gobco2 == 2):
-				break
+			{
+				gobco2 := 1 + 1
+				switch {
+				case GobcoCover(23, gobco2 == 2):
+					break
+				}
 			}
 		}
 	}

testdata/instrumenter/TypeSwitchStmt.cond

@@ -252,9 +252,12 @@ func typeSwitchStmtNested() {
 		switch {
 		case GobcoCover(22, gobco1):
 		default:
-			switch 1 + 1 {
-			case GobcoCover(28, gobco2 == 2):
-				break
+			{
+				gobco2 := 1 + 1
+				switch {
+				case GobcoCover(28, gobco2 == 2):
+					break
+				}
 			}
 		}
 	}