bytes, strings: reduce Split{,After}Seq heap allocations

This CL slightly changes flow of splitSeq to help compiler to inline the iterator closure.

goos: linux
goarch: amd64
pkg: strings
cpu: AMD Ryzen 9 5950X 16-Core Processor
                                    │   sec/op    │   sec/op     vs base                │
SplitSeqEmptySeparator-32             3.590m ± 0%   3.430m ± 2%   -4.46% (p=0.000 n=30)
SplitSeqSingleByteSeparator-32        647.0µ ± 0%   656.1µ ± 0%   +1.41% (p=0.000 n=30)
SplitSeqMultiByteSeparator-32         423.9µ ± 1%   384.5µ ± 0%   -9.31% (p=0.000 n=30)
SplitAfterSeqEmptySeparator-32        3.372m ± 4%   3.514m ± 0%   +4.20% (p=0.000 n=30)
SplitAfterSeqSingleByteSeparator-32   648.5µ ± 2%   537.6µ ± 0%  -17.10% (p=0.000 n=30)
SplitAfterSeqMultiByteSeparator-32    423.3µ ± 2%   364.4µ ± 2%  -13.91% (p=0.000 n=30)
geomean                               984.7µ        917.3µ        -6.85%

                                    │    B/op    │   B/op     vs base                     │
SplitSeqEmptySeparator-32             24.00 ± 0%   0.00 ± 0%  -100.00% (p=0.000 n=30)
SplitSeqSingleByteSeparator-32        24.00 ± 0%   0.00 ± 0%  -100.00% (p=0.000 n=30)
SplitSeqMultiByteSeparator-32         24.00 ± 0%   0.00 ± 0%  -100.00% (p=0.000 n=30)
SplitAfterSeqEmptySeparator-32        24.00 ± 0%   0.00 ± 0%  -100.00% (p=0.000 n=30)
SplitAfterSeqSingleByteSeparator-32   24.00 ± 0%   0.00 ± 0%  -100.00% (p=0.000 n=30)
SplitAfterSeqMultiByteSeparator-32    24.00 ± 0%   0.00 ± 0%  -100.00% (p=0.000 n=30)
geomean                               24.00                   ?

For #73524

Change-Id: Ic83c5751a41c65030356a208e4ad1f500723e695
Reviewed-on: https://go-review.googlesource.com/c/go/+/669735
Auto-Submit: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Reviewed-by: Cherry Mui <cherryyz@google.com>
Reviewed-by: qiu laidongfeng2 <2645477756@qq.com>
Commit-Queue: Alan Donovan <adonovan@google.com>

src/bytes/iter.go

@@ -32,8 +32,7 @@ func Lines(s []byte) iter.Seq[[]byte] {
 }
 
 // explodeSeq returns an iterator over the runes in s.
-func explodeSeq(s []byte) iter.Seq[[]byte] {
-	return func(yield func([]byte) bool) {
+func explodeSeq(s []byte, yield func([]byte) bool) {
 	for len(s) > 0 {
 		_, size := utf8.DecodeRune(s)
 		if !yield(s[:size:size]) {
@@ -42,15 +41,15 @@ func explodeSeq(s []byte) iter.Seq[[]byte] {
 		s = s[size:]
 	}
 }
-}
 
 // splitSeq is SplitSeq or SplitAfterSeq, configured by how many
 // bytes of sep to include in the results (none or all).
 func splitSeq(s, sep []byte, sepSave int) iter.Seq[[]byte] {
-	if len(sep) == 0 {
-		return explodeSeq(s)
-	}
 	return func(yield func([]byte) bool) {
+		if len(sep) == 0 {
+			explodeSeq(s, yield)
+			return
+		}
 		for {
 			i := Index(s, sep)
 			if i < 0 {

src/bytes/iter_test.go

@@ -0,0 +1,56 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package bytes_test
+
+import (
+	. "bytes"
+	"testing"
+)
+
+func BenchmarkSplitSeqEmptySeparator(b *testing.B) {
+	for range b.N {
+		for range SplitSeq(benchInputHard, nil) {
+		}
+	}
+}
+
+func BenchmarkSplitSeqSingleByteSeparator(b *testing.B) {
+	sep := []byte("/")
+	for range b.N {
+		for range SplitSeq(benchInputHard, sep) {
+		}
+	}
+}
+
+func BenchmarkSplitSeqMultiByteSeparator(b *testing.B) {
+	sep := []byte("hello")
+	for range b.N {
+		for range SplitSeq(benchInputHard, sep) {
+		}
+	}
+}
+
+func BenchmarkSplitAfterSeqEmptySeparator(b *testing.B) {
+	for range b.N {
+		for range SplitAfterSeq(benchInputHard, nil) {
+		}
+	}
+}
+
+func BenchmarkSplitAfterSeqSingleByteSeparator(b *testing.B) {
+	sep := []byte("/")
+	for range b.N {
+		for range SplitAfterSeq(benchInputHard, sep) {
+		}
+	}
+}
+
+func BenchmarkSplitAfterSeqMultiByteSeparator(b *testing.B) {
+	sep := []byte("hello")
+	for range b.N {
+		for range SplitAfterSeq(benchInputHard, sep) {
+		}
+	}
+}

src/strings/iter.go

@@ -32,8 +32,7 @@ func Lines(s string) iter.Seq[string] {
 }
 
 // explodeSeq returns an iterator over the runes in s.
-func explodeSeq(s string) iter.Seq[string] {
-	return func(yield func(string) bool) {
+func explodeSeq(s string, yield func(string) bool) {
 	for len(s) > 0 {
 		_, size := utf8.DecodeRuneInString(s)
 		if !yield(s[:size]) {
@@ -42,15 +41,15 @@ func explodeSeq(s string) iter.Seq[string] {
 		s = s[size:]
 	}
 }
-}
 
 // splitSeq is SplitSeq or SplitAfterSeq, configured by how many
 // bytes of sep to include in the results (none or all).
 func splitSeq(s, sep string, sepSave int) iter.Seq[string] {
-	if len(sep) == 0 {
-		return explodeSeq(s)
-	}
 	return func(yield func(string) bool) {
+		if len(sep) == 0 {
+			explodeSeq(s, yield)
+			return
+		}
 		for {
 			i := Index(s, sep)
 			if i < 0 {

src/strings/iter_test.go

@@ -0,0 +1,52 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package strings_test
+
+import (
+	. "strings"
+	"testing"
+)
+
+func BenchmarkSplitSeqEmptySeparator(b *testing.B) {
+	for range b.N {
+		for range SplitSeq(benchInputHard, "") {
+		}
+	}
+}
+
+func BenchmarkSplitSeqSingleByteSeparator(b *testing.B) {
+	for range b.N {
+		for range SplitSeq(benchInputHard, "/") {
+		}
+	}
+}
+
+func BenchmarkSplitSeqMultiByteSeparator(b *testing.B) {
+	for range b.N {
+		for range SplitSeq(benchInputHard, "hello") {
+		}
+	}
+}
+
+func BenchmarkSplitAfterSeqEmptySeparator(b *testing.B) {
+	for range b.N {
+		for range SplitAfterSeq(benchInputHard, "") {
+		}
+	}
+}
+
+func BenchmarkSplitAfterSeqSingleByteSeparator(b *testing.B) {
+	for range b.N {
+		for range SplitAfterSeq(benchInputHard, "/") {
+		}
+	}
+}
+
+func BenchmarkSplitAfterSeqMultiByteSeparator(b *testing.B) {
+	for range b.N {
+		for range SplitAfterSeq(benchInputHard, "hello") {
+		}
+	}
+}