1// Copyright 2024 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package slices
6
7import (
8	"cmp"
9	"iter"
10)
11
12// All returns an iterator over index-value pairs in the slice
13// in the usual order.
14func All[Slice ~[]E, E any](s Slice) iter.Seq2[int, E] {
15	return func(yield func(int, E) bool) {
16		for i, v := range s {
17			if !yield(i, v) {
18				return
19			}
20		}
21	}
22}
23
24// Backward returns an iterator over index-value pairs in the slice,
25// traversing it backward with descending indices.
26func Backward[Slice ~[]E, E any](s Slice) iter.Seq2[int, E] {
27	return func(yield func(int, E) bool) {
28		for i := len(s) - 1; i >= 0; i-- {
29			if !yield(i, s[i]) {
30				return
31			}
32		}
33	}
34}
35
36// Values returns an iterator that yields the slice elements in order.
37func Values[Slice ~[]E, E any](s Slice) iter.Seq[E] {
38	return func(yield func(E) bool) {
39		for _, v := range s {
40			if !yield(v) {
41				return
42			}
43		}
44	}
45}
46
47// AppendSeq appends the values from seq to the slice and
48// returns the extended slice.
49func AppendSeq[Slice ~[]E, E any](s Slice, seq iter.Seq[E]) Slice {
50	for v := range seq {
51		s = append(s, v)
52	}
53	return s
54}
55
56// Collect collects values from seq into a new slice and returns it.
57func Collect[E any](seq iter.Seq[E]) []E {
58	return AppendSeq([]E(nil), seq)
59}
60
61// Sorted collects values from seq into a new slice, sorts the slice,
62// and returns it.
63func Sorted[E cmp.Ordered](seq iter.Seq[E]) []E {
64	s := Collect(seq)
65	Sort(s)
66	return s
67}
68
69// SortedFunc collects values from seq into a new slice, sorts the slice
70// using the comparison function, and returns it.
71func SortedFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
72	s := Collect(seq)
73	SortFunc(s, cmp)
74	return s
75}
76
77// SortedStableFunc collects values from seq into a new slice.
78// It then sorts the slice while keeping the original order of equal elements,
79// using the comparison function to compare elements.
80// It returns the new slice.
81func SortedStableFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
82	s := Collect(seq)
83	SortStableFunc(s, cmp)
84	return s
85}
86
87// Chunk returns an iterator over consecutive sub-slices of up to n elements of s.
88// All but the last sub-slice will have size n.
89// All sub-slices are clipped to have no capacity beyond the length.
90// If s is empty, the sequence is empty: there is no empty slice in the sequence.
91// Chunk panics if n is less than 1.
92func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice] {
93	if n < 1 {
94		panic("cannot be less than 1")
95	}
96
97	return func(yield func(Slice) bool) {
98		for i := 0; i < len(s); i += n {
99			// Clamp the last chunk to the slice bound as necessary.
100			end := min(n, len(s[i:]))
101
102			// Set the capacity of each chunk so that appending to a chunk does
103			// not modify the original slice.
104			if !yield(s[i : i+end : i+end]) {
105				return
106			}
107		}
108	}
109}
110