difflib: optimize SplitLines

Author: rdwilliamson

difflib/difflib.go

@@ -752,9 +752,7 @@ func GetContextDiffString(diff ContextDiff) (string, error) {
 // Split a string on "\n" while preserving them. The output can be used
 // as input for UnifiedDiff and ContextDiff structures.
 func SplitLines(s string) []string {
-	lines := []string{}
-	for _, line := range strings.Split(s, "\n") {
-		lines = append(lines, line+"\n")
-	}
+	lines := strings.SplitAfter(s, "\n")
+	lines[len(lines)-1] += "\n"
 	return lines
 }

difflib/difflib_test.go

@@ -317,3 +317,36 @@ func TestOutputFormatNoTrailingTabOnEmptyFiledate(t *testing.T) {
 	assertEqual(t, err, nil)
 	assertEqual(t, SplitLines(cd)[:2], []string{"*** Original\n", "--- Current\n"})
 }
+
+func TestSplitLines(t *testing.T) {
+	allTests := []struct {
+		input string
+		want  []string
+	}{
+		{"foo", []string{"foo\n"}},
+		{"foo\nbar", []string{"foo\n", "bar\n"}},
+		{"foo\nbar\n", []string{"foo\n", "bar\n", "\n"}},
+	}
+	for _, test := range allTests {
+		assertEqual(t, SplitLines(test.input), test.want)
+	}
+}
+
+func benchmarkSplitLines(b *testing.B, count int) {
+	str := strings.Repeat("foo\n", count)
+
+	b.ResetTimer()
+
+	n := 0
+	for i := 0; i < b.N; i++ {
+		n += len(SplitLines(str))
+	}
+}
+
+func BenchmarkSplitLines100(b *testing.B) {
+	benchmarkSplitLines(b, 100)
+}
+
+func BenchmarkSplitLines10000(b *testing.B) {
+	benchmarkSplitLines(b, 10000)
+}