Merge pull request #13 from heyitsanthony/labels

code: support rewriting labels

code/rewrite.go

@@ -51,13 +51,14 @@ func ToFailpoints(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error)
 				fps = append(fps, curfp)
 				curfp = nil
 			}
-		} else {
+		} else if label := gofailLabel(l); label != "" {
+			// expose gofail label
+			l = label
+		} else if curfp, err = newFailpoint(l); err != nil {
+			return
+		} else if curfp != nil {
 			// found a new failpoint
-			if curfp, err = newFailpoint(l); err != nil {
-				return
-			} else if curfp != nil {
-				continue
-			}
+			continue
 		}
 		if _, err = dst.WriteString(l); err != nil {
 			return
@@ -107,6 +108,10 @@ func ToComments(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error) {
 			continue
 		}
 
+		if isLabel := strings.Contains(l, "\t/* gofail-label */"); isLabel {
+			l = strings.Replace(l, "/* gofail-label */", "// gofail:", 1)
+		}
+
 		if _, werr := dst.WriteString(l); werr != nil {
 			return fps, werr
 		}
@@ -118,6 +123,17 @@ func ToComments(wdst io.Writer, rsrc io.Reader) (fps []*Failpoint, err error) {
 	return
 }
 
+func gofailLabel(l string) string {
+	if !strings.HasPrefix(strings.TrimSpace(l), "// gofail:") {
+		return ""
+	}
+	label := strings.SplitAfter(l, "// gofail:")[1]
+	if len(label) == 0 || !strings.Contains(label, ":") {
+		return ""
+	}
+	return strings.Replace(l, "// gofail:", "/* gofail-label */", 1)
+}
+
 func numBraces(l string) (opening int, closing int) {
 	for i := 0; i < len(l); i++ {
 		switch l[i] {

code/rewrite_test.go

@@ -34,6 +34,18 @@ var examples = []struct {
 	{"func f() {\n\t// gofail: var NoTypeTest struct{}\n\t// fmt.Println(`hi`)\n}\n", 1},
 	{"func f() {\n\t// gofail: var NoTypeTest struct{}\n}\n", 1},
 	{"func f() {\n\t// gofail: var NoTypeTest struct{}\n\t// fmt.Println(`hi`)\n\t// fmt.Println(`bye`)\n}\n", 1},
+	{`
+func f() {
+	// gofail: labelTest:
+	for {
+		if g() {
+			// gofail: var testLabel struct{}
+			// continue labelTest
+			return
+		}
+	}
+}
+`, 1},
 }
 
 func TestToFailpoint(t *testing.T) {

examples/examples.go

@@ -24,3 +24,18 @@ func ExampleOneLineFunc() string {
 	// gofail: var ExampleOneLine struct{}
 	return "abc"
 }
+
+func ExampleLabelsFunc() (s string) {
+	i := 0
+	// gofail: myLabel:
+	for i < 5 {
+		s = s + "i"
+		i++
+		for j := 0; j < 5; j++ {
+			s = s + "j"
+			// gofail: var ExampleLabels struct{}
+			// continue myLabel
+		}
+	}
+	return s
+}