slices: add Chunk

Chunk returns an iterator over consecutive sub-slices of up to n elements of s.

Fixes #53987.

Change-Id: I508274eca388db39550eb9e4d8abd5ce68d29d8d
Reviewed-on: https://go-review.googlesource.com/c/go/+/562935
Reviewed-by: Cherry Mui <cherryyz@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Auto-Submit: Ian Lance Taylor <iant@google.com>

diff --git a/api/next/53987.txt b/api/next/53987.txt
new file mode 100644 (file)
index 0000000..1861d0b
--- /dev/null
+++ b/api/next/53987.txt
@@ -0,0 +1 @@
+pkg slices, func Chunk[$0 interface{ ~[]$1 }, $1 interface{}]($0, int) iter.Seq[$0] #53987

diff --git a/doc/next/6-stdlib/3-iter.md b/doc/next/6-stdlib/3-iter.md
index bc74f4556cd0e03f81a4e46180be318f3f886a6d..6b52b7c7e5624bc3afc0e74bfed4385ddf77bec8 100644 (file)
--- a/doc/next/6-stdlib/3-iter.md
+++ b/doc/next/6-stdlib/3-iter.md
@@ -19,3 +19,5 @@ with iterators:
   comparison function.
 - [SortedStableFunc](/pkg/slices#SortedStableFunc) is like `SortFunc`
   but uses a stable sort algorithm.
+- [Chunk](/pkg/slices#Chunk) returns an iterator over consecutive
+  sub-slices of up to n elements of a slice.

diff --git a/doc/next/6-stdlib/99-minor/slices/53987.md b/doc/next/6-stdlib/99-minor/slices/53987.md
new file mode 100644 (file)
index 0000000..02d77cd
--- /dev/null
+++ b/doc/next/6-stdlib/99-minor/slices/53987.md
@@ -0,0 +1 @@
+<!-- see ../../3-iter.md -->

diff --git a/src/slices/example_test.go b/src/slices/example_test.go
index e1bda36e28b53b249eedb829f0aedac01f03dd24..cb601ada0ade65f3b0aeb903068630b0cb4bee5f 100644 (file)
--- a/src/slices/example_test.go
+++ b/src/slices/example_test.go
@@ -384,3 +384,30 @@ func ExampleRepeat() {
        // Output:
        // [0 1 2 3 0 1 2 3]
 }
+
+func ExampleChunk() {
+       type Person struct {
+               Name string
+               Age  int
+       }
+
+       type People []Person
+
+       people := People{
+               {"Gopher", 13},
+               {"Alice", 20},
+               {"Bob", 5},
+               {"Vera", 24},
+               {"Zac", 15},
+       }
+
+       // Chunk people into []Person 2 elements at a time.
+       for c := range slices.Chunk(people, 2) {
+               fmt.Println(c)
+       }
+
+       // Output:
+       // [{Gopher 13} {Alice 20}]
+       // [{Bob 5} {Vera 24}]
+       // [{Zac 15}]
+}

diff --git a/src/slices/iter.go b/src/slices/iter.go
index 985bd27a1079d56bc92fb2263838a2564e0fc95c..a0f642e423250f480f1067995e3c597ded53927d 100644 (file)
--- a/src/slices/iter.go
+++ b/src/slices/iter.go
@@ -84,3 +84,27 @@ func SortedStableFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
        SortStableFunc(s, cmp)
        return s
 }
+
+// Chunk returns an iterator over consecutive sub-slices of up to n elements of s.
+// All but the last sub-slice will have size n.
+// All sub-slices are clipped to have no capacity beyond the length.
+// If s is empty, the sequence is empty: there is no empty slice in the sequence.
+// Chunk panics if n is less than 1.
+func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice] {
+       if n < 1 {
+               panic("cannot be less than 1")
+       }
+
+       return func(yield func(Slice) bool) {
+               for i := 0; i < len(s); i += n {
+                       // Clamp the last chunk to the slice bound as necessary.
+                       end := min(n, len(s[i:]))
+
+                       // Set the capacity of each chunk so that appending to a chunk does
+                       // not modify the original slice.
+                       if !yield(s[i : i+end : i+end]) {
+                               return
+                       }
+               }
+       }
+}

diff --git a/src/slices/iter_test.go b/src/slices/iter_test.go
index 67520f60c9d5f2f6bae1097e65c21072a1d0e266..07d73e90e2bc561287f5181a88d4d71618a1d31c 100644 (file)
--- a/src/slices/iter_test.go
+++ b/src/slices/iter_test.go
@@ -182,3 +182,113 @@ func TestSortedStableFunc(t *testing.T) {
                t.Errorf("SortedStableFunc wasn't stable on %d reverse ints", n)
        }
 }
+
+func TestChunk(t *testing.T) {
+       cases := []struct {
+               name   string
+               s      []int
+               n      int
+               chunks [][]int
+       }{
+               {
+                       name:   "nil",
+                       s:      nil,
+                       n:      1,
+                       chunks: nil,
+               },
+               {
+                       name:   "empty",
+                       s:      []int{},
+                       n:      1,
+                       chunks: nil,
+               },
+               {
+                       name:   "short",
+                       s:      []int{1, 2},
+                       n:      3,
+                       chunks: [][]int{{1, 2}},
+               },
+               {
+                       name:   "one",
+                       s:      []int{1, 2},
+                       n:      2,
+                       chunks: [][]int{{1, 2}},
+               },
+               {
+                       name:   "even",
+                       s:      []int{1, 2, 3, 4},
+                       n:      2,
+                       chunks: [][]int{{1, 2}, {3, 4}},
+               },
+               {
+                       name:   "odd",
+                       s:      []int{1, 2, 3, 4, 5},
+                       n:      2,
+                       chunks: [][]int{{1, 2}, {3, 4}, {5}},
+               },
+       }
+
+       for _, tc := range cases {
+               t.Run(tc.name, func(t *testing.T) {
+                       var chunks [][]int
+                       for c := range Chunk(tc.s, tc.n) {
+                               chunks = append(chunks, c)
+                       }
+
+                       if !chunkEqual(chunks, tc.chunks) {
+                               t.Errorf("Chunk(%v, %d) = %v, want %v", tc.s, tc.n, chunks, tc.chunks)
+                       }
+
+                       if len(chunks) == 0 {
+                               return
+                       }
+
+                       // Verify that appending to the end of the first chunk does not
+                       // clobber the beginning of the next chunk.
+                       s := Clone(tc.s)
+                       chunks[0] = append(chunks[0], -1)
+                       if !Equal(s, tc.s) {
+                               t.Errorf("slice was clobbered: %v, want %v", s, tc.s)
+                       }
+               })
+       }
+}
+
+func TestChunkPanics(t *testing.T) {
+       for _, test := range []struct {
+               name string
+               x    []struct{}
+               n    int
+       }{
+               {
+                       name: "cannot be less than 1",
+                       x:    make([]struct{}, 0),
+                       n:    0,
+               },
+       } {
+               if !panics(func() { _ = Chunk(test.x, test.n) }) {
+                       t.Errorf("Chunk %s: got no panic, want panic", test.name)
+               }
+       }
+}
+
+func TestChunkRange(t *testing.T) {
+       // Verify Chunk iteration can be stopped.
+       var got [][]int
+       for c := range Chunk([]int{1, 2, 3, 4, -100}, 2) {
+               if len(got) == 2 {
+                       // Found enough values, break early.
+                       break
+               }
+
+               got = append(got, c)
+       }
+
+       if want := [][]int{{1, 2}, {3, 4}}; !chunkEqual(got, want) {
+               t.Errorf("Chunk iteration did not stop, got %v, want %v", got, want)
+       }
+}
+
+func chunkEqual[Slice ~[]E, E comparable](s1, s2 []Slice) bool {
+       return EqualFunc(s1, s2, Equal[Slice])
+}