Files
lo/slice.go
2025-09-25 04:15:30 +02:00

888 lines
25 KiB
Go

package lo
import (
"sort"
"github.com/samber/lo/internal/constraints"
"github.com/samber/lo/mutable"
)
// Filter iterates over elements of collection, returning a slice of all elements predicate returns true for.
// Play: https://go.dev/play/p/Apjg3WeSi7K
func Filter[T any, Slice ~[]T](collection Slice, predicate func(item T, index int) bool) Slice {
result := make(Slice, 0, len(collection))
for i := range collection {
if predicate(collection[i], i) {
result = append(result, collection[i])
}
}
return result
}
// Map manipulates a slice and transforms it to a slice of another type.
// Play: https://go.dev/play/p/OkPcYAhBo0D
func Map[T any, R any](collection []T, iteratee func(item T, index int) R) []R {
result := make([]R, len(collection))
for i := range collection {
result[i] = iteratee(collection[i], i)
}
return result
}
// UniqMap manipulates a slice and transforms it to a slice of another type with unique values.
// Play: https://go.dev/play/p/fygzLBhvUdB
func UniqMap[T any, R comparable](collection []T, iteratee func(item T, index int) R) []R {
result := make([]R, 0, len(collection))
seen := make(map[R]struct{}, len(collection))
for i := range collection {
r := iteratee(collection[i], i)
if _, ok := seen[r]; !ok {
result = append(result, r)
seen[r] = struct{}{}
}
}
return result
}
// FilterMap returns a slice obtained after both filtering and mapping using the given callback function.
// The callback function should return two values:
// - the result of the mapping operation and
// - whether the result element should be included or not.
//
// Play: https://go.dev/play/p/CgHYNUpOd1I
func FilterMap[T any, R any](collection []T, callback func(item T, index int) (R, bool)) []R {
result := make([]R, 0, len(collection))
for i := range collection {
if r, ok := callback(collection[i], i); ok {
result = append(result, r)
}
}
return result
}
// FlatMap manipulates a slice and transforms and flattens it to a slice of another type.
// The transform function can either return a slice or a `nil`, and in the `nil` case
// no value is added to the final slice.
// Play: https://go.dev/play/p/pFCF5WVB225
func FlatMap[T any, R any](collection []T, iteratee func(item T, index int) []R) []R {
result := make([]R, 0, len(collection))
for i := range collection {
result = append(result, iteratee(collection[i], i)...)
}
return result
}
// Reduce reduces collection to a value which is the accumulated result of running each element in collection
// through accumulator, where each successive invocation is supplied the return value of the previous.
// Play: https://go.dev/play/p/CgHYNUpOd1I
func Reduce[T any, R any](collection []T, accumulator func(agg R, item T, index int) R, initial R) R {
for i := range collection {
initial = accumulator(initial, collection[i], i)
}
return initial
}
// ReduceRight is like Reduce except that it iterates over elements of collection from right to left.
// Play: https://go.dev/play/p/Fq3W70l7wXF
func ReduceRight[T any, R any](collection []T, accumulator func(agg R, item T, index int) R, initial R) R {
for i := len(collection) - 1; i >= 0; i-- {
initial = accumulator(initial, collection[i], i)
}
return initial
}
// ForEach iterates over elements of collection and invokes iteratee for each element.
// Play: https://go.dev/play/p/oofyiUPRf8t
func ForEach[T any](collection []T, iteratee func(item T, index int)) {
for i := range collection {
iteratee(collection[i], i)
}
}
// ForEachWhile iterates over elements of collection and invokes iteratee for each element
// collection return value decide to continue or break, like do while().
// Play: https://go.dev/play/p/QnLGt35tnow
func ForEachWhile[T any](collection []T, iteratee func(item T, index int) bool) {
for i := range collection {
if !iteratee(collection[i], i) {
break
}
}
}
// Times invokes the iteratee n times, returning a slice of the results of each invocation.
// The iteratee is invoked with index as argument.
// Play: https://go.dev/play/p/vgQj3Glr6lT
func Times[T any](count int, iteratee func(index int) T) []T {
result := make([]T, count)
for i := 0; i < count; i++ {
result[i] = iteratee(i)
}
return result
}
// Uniq returns a duplicate-free version of a slice, in which only the first occurrence of each element is kept.
// The order of result values is determined by the order they occur in the slice.
// Play: https://go.dev/play/p/DTzbeXZ6iEN
func Uniq[T comparable, Slice ~[]T](collection Slice) Slice {
result := make(Slice, 0, len(collection))
seen := make(map[T]struct{}, len(collection))
for i := range collection {
if _, ok := seen[collection[i]]; ok {
continue
}
seen[collection[i]] = struct{}{}
result = append(result, collection[i])
}
return result
}
// UniqBy returns a duplicate-free version of a slice, in which only the first occurrence of each element is kept.
// The order of result values is determined by the order they occur in the slice. It accepts `iteratee` which is
// invoked for each element in the slice to generate the criterion by which uniqueness is computed.
// Play: https://go.dev/play/p/g42Z3QSb53u
func UniqBy[T any, U comparable, Slice ~[]T](collection Slice, iteratee func(item T) U) Slice {
result := make(Slice, 0, len(collection))
seen := make(map[U]struct{}, len(collection))
for i := range collection {
key := iteratee(collection[i])
if _, ok := seen[key]; ok {
continue
}
seen[key] = struct{}{}
result = append(result, collection[i])
}
return result
}
// GroupBy returns an object composed of keys generated from the results of running each element of collection through iteratee.
// Play: https://go.dev/play/p/XnQBd_v6brd
func GroupBy[T any, U comparable, Slice ~[]T](collection Slice, iteratee func(item T) U) map[U]Slice {
result := map[U]Slice{}
for i := range collection {
key := iteratee(collection[i])
result[key] = append(result[key], collection[i])
}
return result
}
// GroupByMap returns an object composed of keys generated from the results of running each element of collection through iteratee.
// Play: https://go.dev/play/p/iMeruQ3_W80
func GroupByMap[T any, K comparable, V any](collection []T, iteratee func(item T) (K, V)) map[K][]V {
result := map[K][]V{}
for i := range collection {
k, v := iteratee(collection[i])
result[k] = append(result[k], v)
}
return result
}
// Chunk returns a slice of elements split into groups of length size. If the slice can't be split evenly,
// the final chunk will be the remaining elements.
// Play: https://go.dev/play/p/kEMkFbdu85g
func Chunk[T any, Slice ~[]T](collection Slice, size int) []Slice {
if size <= 0 {
panic("lo.Chunk: size must be greater than 0")
}
chunksNum := len(collection) / size
if len(collection)%size != 0 {
chunksNum++
}
result := make([]Slice, 0, chunksNum)
for i := 0; i < chunksNum; i++ {
last := (i + 1) * size
if last > len(collection) {
last = len(collection)
}
// Copy chunk in a new slice, to prevent memory leak and free memory from initial collection.
newSlice := make(Slice, last-i*size)
copy(newSlice, collection[i*size:last])
result = append(result, newSlice)
}
return result
}
// PartitionBy returns a slice of elements split into groups. The order of grouped values is
// determined by the order they occur in collection. The grouping is generated from the results
// of running each element of collection through iteratee.
// Play: https://go.dev/play/p/NfQ_nGjkgXW
func PartitionBy[T any, K comparable, Slice ~[]T](collection Slice, iteratee func(item T) K) []Slice {
result := []Slice{}
seen := map[K]int{}
for i := range collection {
key := iteratee(collection[i])
resultIndex, ok := seen[key]
if !ok {
resultIndex = len(result)
seen[key] = resultIndex
result = append(result, Slice{})
}
result[resultIndex] = append(result[resultIndex], collection[i])
}
return result
// unordered:
// groups := GroupBy[T, K](collection, iteratee)
// return Values[K, []T](groups)
}
// Flatten returns a slice a single level deep.
// Play: https://go.dev/play/p/rbp9ORaMpjw
func Flatten[T any, Slice ~[]T](collection []Slice) Slice {
totalLen := 0
for i := range collection {
totalLen += len(collection[i])
}
result := make(Slice, 0, totalLen)
for i := range collection {
result = append(result, collection[i]...)
}
return result
}
// Interleave round-robin alternating input slices and sequentially appending value at index into result.
// Play: https://go.dev/play/p/-RJkTLQEDVt
func Interleave[T any, Slice ~[]T](collections ...Slice) Slice {
if len(collections) == 0 {
return Slice{}
}
maxSize := 0
totalSize := 0
for i := range collections {
size := len(collections[i])
totalSize += size
if size > maxSize {
maxSize = size
}
}
if maxSize == 0 {
return Slice{}
}
result := make(Slice, totalSize)
resultIdx := 0
for i := 0; i < maxSize; i++ {
for j := range collections {
if len(collections[j])-1 < i {
continue
}
result[resultIdx] = collections[j][i]
resultIdx++
}
}
return result
}
// Shuffle returns a slice of shuffled values. Uses the Fisher-Yates shuffle algorithm.
// Play: https://go.dev/play/p/ZTGG7OUCdnp
//
// Deprecated: use mutable.Shuffle() instead.
func Shuffle[T any, Slice ~[]T](collection Slice) Slice {
mutable.Shuffle(collection)
return collection
}
// Reverse reverses a slice so that the first element becomes the last, the second element becomes the second to last, and so on.
// Play: https://go.dev/play/p/iv2e9jslfBM
//
// Deprecated: use mutable.Reverse() instead.
func Reverse[T any, Slice ~[]T](collection Slice) Slice {
mutable.Reverse(collection)
return collection
}
// Fill fills elements of a slice with `initial` value.
// Play: https://go.dev/play/p/VwR34GzqEub
func Fill[T Clonable[T], Slice ~[]T](collection Slice, initial T) Slice {
result := make(Slice, 0, len(collection))
for range collection {
result = append(result, initial.Clone())
}
return result
}
// Repeat builds a slice with N copies of initial value.
// Play: https://go.dev/play/p/g3uHXbmc3b6
func Repeat[T Clonable[T]](count int, initial T) []T {
result := make([]T, 0, count)
for i := 0; i < count; i++ {
result = append(result, initial.Clone())
}
return result
}
// RepeatBy builds a slice with values returned by N calls of callback.
// Play: https://go.dev/play/p/ozZLCtX_hNU
func RepeatBy[T any](count int, predicate func(index int) T) []T {
result := make([]T, 0, count)
for i := 0; i < count; i++ {
result = append(result, predicate(i))
}
return result
}
// KeyBy transforms a slice or a slice of structs to a map based on a pivot callback.
// Play: https://go.dev/play/p/ccUiUL_Lnel
func KeyBy[K comparable, V any](collection []V, iteratee func(item V) K) map[K]V {
result := make(map[K]V, len(collection))
for i := range collection {
k := iteratee(collection[i])
result[k] = collection[i]
}
return result
}
// Associate returns a map containing key-value pairs provided by transform function applied to elements of the given slice.
// If any of two pairs have the same key the last one gets added to the map.
// The order of keys in returned map is not specified and is not guaranteed to be the same from the original slice.
// Play: https://go.dev/play/p/WHa2CfMO3Lr
func Associate[T any, K comparable, V any](collection []T, transform func(item T) (K, V)) map[K]V {
result := make(map[K]V, len(collection))
for i := range collection {
k, v := transform(collection[i])
result[k] = v
}
return result
}
// SliceToMap returns a map containing key-value pairs provided by transform function applied to elements of the given slice.
// If any of two pairs have the same key the last one gets added to the map.
// The order of keys in returned map is not specified and is not guaranteed to be the same from the original slice.
// Alias of Associate().
// Play: https://go.dev/play/p/WHa2CfMO3Lr
func SliceToMap[T any, K comparable, V any](collection []T, transform func(item T) (K, V)) map[K]V {
return Associate(collection, transform)
}
// FilterSliceToMap returns a map containing key-value pairs provided by transform function applied to elements of the given slice.
// If any of two pairs have the same key the last one gets added to the map.
// The order of keys in returned map is not specified and is not guaranteed to be the same from the original slice.
// The third return value of the transform function is a boolean that indicates whether the key-value pair should be included in the map.
// Play: https://go.dev/play/p/2z0rDz2ZSGU
func FilterSliceToMap[T any, K comparable, V any](collection []T, transform func(item T) (K, V, bool)) map[K]V {
result := make(map[K]V, len(collection))
for i := range collection {
k, v, ok := transform(collection[i])
if ok {
result[k] = v
}
}
return result
}
// Keyify returns a map with each unique element of the slice as a key.
// Play: https://go.dev/play/p/RYhhM_csqIG
func Keyify[T comparable, Slice ~[]T](collection Slice) map[T]struct{} {
result := make(map[T]struct{}, len(collection))
for i := range collection {
result[collection[i]] = struct{}{}
}
return result
}
// Drop drops n elements from the beginning of a slice.
// Play: https://go.dev/play/p/JswS7vXRJP2
func Drop[T any, Slice ~[]T](collection Slice, n int) Slice {
if len(collection) <= n {
return make(Slice, 0)
}
result := make(Slice, 0, len(collection)-n)
return append(result, collection[n:]...)
}
// DropRight drops n elements from the end of a slice.
// Play: https://go.dev/play/p/GG0nXkSJJa3
func DropRight[T any, Slice ~[]T](collection Slice, n int) Slice {
if len(collection) <= n {
return Slice{}
}
result := make(Slice, 0, len(collection)-n)
return append(result, collection[:len(collection)-n]...)
}
// DropWhile drops elements from the beginning of a slice while the predicate returns true.
// Play: https://go.dev/play/p/7gBPYw2IK16
func DropWhile[T any, Slice ~[]T](collection Slice, predicate func(item T) bool) Slice {
i := 0
for ; i < len(collection); i++ {
if !predicate(collection[i]) {
break
}
}
result := make(Slice, 0, len(collection)-i)
return append(result, collection[i:]...)
}
// DropRightWhile drops elements from the end of a slice while the predicate returns true.
// Play: https://go.dev/play/p/3-n71oEC0Hz
func DropRightWhile[T any, Slice ~[]T](collection Slice, predicate func(item T) bool) Slice {
i := len(collection) - 1
for ; i >= 0; i-- {
if !predicate(collection[i]) {
break
}
}
result := make(Slice, 0, i+1)
return append(result, collection[:i+1]...)
}
// DropByIndex drops elements from a slice by the index.
// A negative index will drop elements from the end of the slice.
// Play: https://go.dev/play/p/bPIH4npZRxS
func DropByIndex[T any, Slice ~[]T](collection Slice, indexes ...int) Slice {
initialSize := len(collection)
if initialSize == 0 {
return make(Slice, 0)
}
for i := range indexes {
if indexes[i] < 0 {
indexes[i] = initialSize + indexes[i]
}
}
indexes = Uniq(indexes)
sort.Ints(indexes)
result := make(Slice, 0, initialSize)
result = append(result, collection...)
for i := range indexes {
if indexes[i]-i < 0 || indexes[i]-i >= initialSize-i {
continue
}
result = append(result[:indexes[i]-i], result[indexes[i]-i+1:]...)
}
return result
}
// Reject is the opposite of Filter, this method returns the elements of collection that predicate does not return true for.
// Play: https://go.dev/play/p/pFCF5WVB225
func Reject[T any, Slice ~[]T](collection Slice, predicate func(item T, index int) bool) Slice {
result := Slice{}
for i := range collection {
if !predicate(collection[i], i) {
result = append(result, collection[i])
}
}
return result
}
// RejectMap is the opposite of FilterMap, this method returns a slice obtained after both filtering and mapping using the given callback function.
// The callback function should return two values:
// - the result of the mapping operation and
// - whether the result element should be included or not.
//
// Play: https://go.dev/play/p/W9Ug9r0QFkL
func RejectMap[T any, R any](collection []T, callback func(item T, index int) (R, bool)) []R {
result := []R{}
for i := range collection {
if r, ok := callback(collection[i], i); !ok {
result = append(result, r)
}
}
return result
}
// FilterReject mixes Filter and Reject, this method returns two slices, one for the elements of collection that
// predicate returns true for and one for the elements that predicate does not return true for.
// Play: https://go.dev/play/p/lHSEGSznJjB
func FilterReject[T any, Slice ~[]T](collection Slice, predicate func(T, int) bool) (kept Slice, rejected Slice) {
kept = make(Slice, 0, len(collection))
rejected = make(Slice, 0, len(collection))
for i := range collection {
if predicate(collection[i], i) {
kept = append(kept, collection[i])
} else {
rejected = append(rejected, collection[i])
}
}
return kept, rejected
}
// Count counts the number of elements in the collection that equal value.
// Play: https://go.dev/play/p/Y3FlK54yveC
func Count[T comparable](collection []T, value T) int {
var count int
for i := range collection {
if collection[i] == value {
count++
}
}
return count
}
// CountBy counts the number of elements in the collection for which predicate is true.
// Play: https://go.dev/play/p/ByQbNYQQi4X
func CountBy[T any](collection []T, predicate func(item T) bool) int {
var count int
for i := range collection {
if predicate(collection[i]) {
count++
}
}
return count
}
// CountValues counts the number of each element in the collection.
// Play: https://go.dev/play/p/-p-PyLT4dfy
func CountValues[T comparable](collection []T) map[T]int {
result := make(map[T]int)
for i := range collection {
result[collection[i]]++
}
return result
}
// CountValuesBy counts the number of each element returned from mapper function.
// Is equivalent to chaining lo.Map and lo.CountValues.
// Play: https://go.dev/play/p/2U0dG1SnOmS
func CountValuesBy[T any, U comparable](collection []T, mapper func(item T) U) map[U]int {
result := make(map[U]int)
for i := range collection {
result[mapper(collection[i])]++
}
return result
}
// Subset returns a copy of a slice from `offset` up to `length` elements. Like `slice[start:start+length]`, but does not panic on overflow.
// Play: https://go.dev/play/p/tOQu1GhFcog
func Subset[T any, Slice ~[]T](collection Slice, offset int, length uint) Slice {
size := len(collection)
if offset < 0 {
offset = size + offset
if offset < 0 {
offset = 0
}
}
if offset > size {
return Slice{}
}
if length > uint(size)-uint(offset) {
length = uint(size - offset)
}
return collection[offset : offset+int(length)]
}
// Slice returns a copy of a slice from `start` up to, but not including `end`. Like `slice[start:end]`, but does not panic on overflow.
// Play: https://go.dev/play/p/8XWYhfMMA1h
func Slice[T any, Slice ~[]T](collection Slice, start int, end int) Slice {
size := len(collection)
if start >= end {
return Slice{}
}
if start > size {
start = size
}
if start < 0 {
start = 0
}
if end > size {
end = size
}
if end < 0 {
end = 0
}
return collection[start:end]
}
// Replace returns a copy of the slice with the first n non-overlapping instances of old replaced by new.
// Play: https://go.dev/play/p/XfPzmf9gql6
func Replace[T comparable, Slice ~[]T](collection Slice, old T, nEw T, n int) Slice {
result := make(Slice, len(collection))
copy(result, collection)
for i := range result {
if result[i] == old && n != 0 {
result[i] = nEw
n--
}
}
return result
}
// ReplaceAll returns a copy of the slice with all non-overlapping instances of old replaced by new.
// Play: https://go.dev/play/p/a9xZFUHfYcV
func ReplaceAll[T comparable, Slice ~[]T](collection Slice, old T, nEw T) Slice {
return Replace(collection, old, nEw, -1)
}
// Compact returns a slice of all non-zero elements.
// Play: https://go.dev/play/p/tXiy-iK6PAc
func Compact[T comparable, Slice ~[]T](collection Slice) Slice {
var zero T
result := make(Slice, 0, len(collection))
for i := range collection {
if collection[i] != zero {
result = append(result, collection[i])
}
}
return result
}
// IsSorted checks if a slice is sorted.
// Play: https://go.dev/play/p/mc3qR-t4mcx
func IsSorted[T constraints.Ordered](collection []T) bool {
for i := 1; i < len(collection); i++ {
if collection[i-1] > collection[i] {
return false
}
}
return true
}
// IsSortedByKey checks if a slice is sorted by iteratee.
func IsSortedByKey[T any, K constraints.Ordered](collection []T, iteratee func(item T) K) bool {
size := len(collection)
for i := 0; i < size-1; i++ {
if iteratee(collection[i]) > iteratee(collection[i+1]) {
return false
}
}
return true
}
// Splice inserts multiple elements at index i. A negative index counts back
// from the end of the slice. The helper is protected against overflow errors.
// Play: https://go.dev/play/p/G5_GhkeSUBA
func Splice[T any, Slice ~[]T](collection Slice, i int, elements ...T) Slice {
sizeCollection := len(collection)
sizeElements := len(elements)
output := make(Slice, 0, sizeCollection+sizeElements) // preallocate memory for the output slice
switch {
case sizeElements == 0:
return append(output, collection...) // simple copy
case i > sizeCollection:
// positive overflow
return append(append(output, collection...), elements...)
case i < -sizeCollection:
// negative overflow
return append(append(output, elements...), collection...)
case i < 0:
// backward
i = sizeCollection + i
}
return append(append(append(output, collection[:i]...), elements...), collection[i:]...)
}
// Cut slices collection around the first instance of separator, returning the part of collection
// before and after separator. The found result reports whether separator appears in collection.
// If separator does not appear in s, cut returns collection, empty slice of []T, false.
// Play: https://go.dev/play/p/GiL3qhpIP3f
func Cut[T comparable, Slice ~[]T](collection Slice, separator Slice) (before Slice, after Slice, found bool) {
if len(separator) == 0 {
return make(Slice, 0), collection, true
}
for i := 0; i+len(separator) <= len(collection); i++ {
match := true
for j := 0; j < len(separator); j++ {
if collection[i+j] != separator[j] {
match = false
break
}
}
if match {
return collection[:i], collection[i+len(separator):], true
}
}
return collection, make(Slice, 0), false
}
// CutPrefix returns collection without the provided leading prefix []T
// and reports whether it found the prefix.
// If s doesn't start with prefix, CutPrefix returns collection, false.
// If prefix is the empty []T, CutPrefix returns collection, true.
// Play: https://go.dev/play/p/7Plak4a1ICl
func CutPrefix[T comparable, Slice ~[]T](collection Slice, separator Slice) (after Slice, found bool) {
if len(separator) == 0 {
return collection, true
}
if len(separator) > len(collection) {
return collection, false
}
for i := range separator {
if collection[i] != separator[i] {
return collection, false
}
}
return collection[len(separator):], true
}
// CutSuffix returns collection without the provided ending suffix []T and reports
// whether it found the suffix. If s doesn't end with suffix, CutSuffix returns collection, false.
// If suffix is the empty []T, CutSuffix returns collection, true.
// Play: https://go.dev/play/p/7FKfBFvPTaT
func CutSuffix[T comparable, Slice ~[]T](collection Slice, separator Slice) (before Slice, found bool) {
if len(separator) == 0 {
return collection, true
}
if len(separator) > len(collection) {
return collection, false
}
start := len(collection) - len(separator)
for i := range separator {
if collection[start+i] != separator[i] {
return collection, false
}
}
return collection[:start], true
}
// Trim removes all the leading and trailing cutset from the collection.
// Play: https://go.dev/play/p/1an9mxLdRG5
func Trim[T comparable, Slice ~[]T](collection Slice, cutset Slice) Slice {
return TrimLeft(TrimRight(collection, cutset), cutset)
}
// TrimLeft removes all the leading cutset from the collection.
// Play: https://go.dev/play/p/74aqfAYLmyi
func TrimLeft[T comparable, Slice ~[]T](collection Slice, cutset Slice) Slice {
set := Keyify(cutset)
return DropWhile(collection, func(item T) bool {
_, ok := set[item]
return ok
})
}
// TrimPrefix removes all the leading prefix from the collection.
// Play: https://go.dev/play/p/SHO6X-YegPg
func TrimPrefix[T comparable, Slice ~[]T](collection Slice, prefix Slice) Slice {
if len(prefix) == 0 {
return collection
}
for {
if !HasPrefix(collection, prefix) {
return collection
}
collection = collection[len(prefix):]
}
}
// TrimRight removes all the trailing cutset from the collection.
// Play: https://go.dev/play/p/MRpAfR6sf0g
func TrimRight[T comparable, Slice ~[]T](collection Slice, cutset Slice) Slice {
set := Keyify(cutset)
return DropRightWhile(collection, func(item T) bool {
_, ok := set[item]
return ok
})
}
// TrimSuffix removes all the trailing suffix from the collection.
// Play: https://go.dev/play/p/IjEUrV0iofq
func TrimSuffix[T comparable, Slice ~[]T](collection Slice, suffix Slice) Slice {
if len(suffix) == 0 {
return collection
}
for {
if !HasSuffix(collection, suffix) {
return collection
}
collection = collection[:len(collection)-len(suffix)]
}
}