iter.go

     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  
     5  package slices
     6  
     7  import (
     8  	"cmp"
     9  	"iter"
    10  )
    11  
    12  // All returns an iterator over index-value pairs in the slice.
    13  // The indexes range in the usual order, from 0 through len(s)-1.
    14  func 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. The indexes range from len(s)-1 down to 0.
    26  func 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 over the slice elements,
    37  // starting with s[0].
    38  func Values[Slice ~[]E, E any](s Slice) iter.Seq[E] {
    39  	return func(yield func(E) bool) {
    40  		for _, v := range s {
    41  			if !yield(v) {
    42  				return
    43  			}
    44  		}
    45  	}
    46  }
    47  
    48  // AppendSeq appends the values from seq to the slice and
    49  // returns the extended slice.
    50  func AppendSeq[Slice ~[]E, E any](s Slice, seq iter.Seq[E]) Slice {
    51  	for v := range seq {
    52  		s = append(s, v)
    53  	}
    54  	return s
    55  }
    56  
    57  // Collect collects values from seq into a new slice and returns it.
    58  func Collect[E any](seq iter.Seq[E]) []E {
    59  	return AppendSeq([]E(nil), seq)
    60  }
    61  
    62  // Sorted collects values from seq into a new slice, sorts the slice,
    63  // and returns it.
    64  func Sorted[E cmp.Ordered](seq iter.Seq[E]) []E {
    65  	s := Collect(seq)
    66  	Sort(s)
    67  	return s
    68  }
    69  
    70  // SortedFunc collects values from seq into a new slice, sorts the slice
    71  // using the comparison function, and returns it.
    72  func SortedFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
    73  	s := Collect(seq)
    74  	SortFunc(s, cmp)
    75  	return s
    76  }
    77  
    78  // SortedStableFunc collects values from seq into a new slice.
    79  // It then sorts the slice while keeping the original order of equal elements,
    80  // using the comparison function to compare elements.
    81  // It returns the new slice.
    82  func SortedStableFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
    83  	s := Collect(seq)
    84  	SortStableFunc(s, cmp)
    85  	return s
    86  }
    87  
    88  // Chunk returns an iterator over consecutive sub-slices of up to n elements of s.
    89  // All but the last sub-slice will have size n.
    90  // All sub-slices are clipped to have no capacity beyond the length.
    91  // If s is empty, the sequence is empty: there is no empty slice in the sequence.
    92  // Chunk panics if n is less than 1.
    93  func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice] {
    94  	if n < 1 {
    95  		panic("cannot be less than 1")
    96  	}
    97  
    98  	return func(yield func(Slice) bool) {
    99  		for i := 0; i < len(s); i += n {
   100  			// Clamp the last chunk to the slice bound as necessary.
   101  			end := min(n, len(s[i:]))
   102  
   103  			// Set the capacity of each chunk so that appending to a chunk does
   104  			// not modify the original slice.
   105  			if !yield(s[i : i+end : i+end]) {
   106  				return
   107  			}
   108  		}
   109  	}
   110  }
   111
View as plain text