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  // in the usual order.
    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 with descending indices.
    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 that yields the slice elements in order.
    37  func 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.
    49  // If seq is empty, the result preserves the nilness of s.
    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  // If seq is empty, the result is nil.
    59  func Collect[E any](seq iter.Seq[E]) []E {
    60  	return AppendSeq([]E(nil), seq)
    61  }
    62  
    63  // Sorted collects values from seq into a new slice, sorts the slice,
    64  // and returns it.
    65  // If seq is empty, the result is nil.
    66  func Sorted[E cmp.Ordered](seq iter.Seq[E]) []E {
    67  	s := Collect(seq)
    68  	Sort(s)
    69  	return s
    70  }
    71  
    72  // SortedFunc collects values from seq into a new slice, sorts the slice
    73  // using the comparison function, and returns it.
    74  // If seq is empty, the result is nil.
    75  func SortedFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
    76  	s := Collect(seq)
    77  	SortFunc(s, cmp)
    78  	return s
    79  }
    80  
    81  // SortedStableFunc collects values from seq into a new slice.
    82  // It then sorts the slice while keeping the original order of equal elements,
    83  // using the comparison function to compare elements.
    84  // It returns the new slice.
    85  // If seq is empty, the result is nil.
    86  func SortedStableFunc[E any](seq iter.Seq[E], cmp func(E, E) int) []E {
    87  	s := Collect(seq)
    88  	SortStableFunc(s, cmp)
    89  	return s
    90  }
    91  
    92  // Chunk returns an iterator over consecutive sub-slices of up to n elements of s.
    93  // All but the last sub-slice will have size n.
    94  // All sub-slices are clipped to have no capacity beyond the length.
    95  // If s is empty, the sequence is empty: there is no empty slice in the sequence.
    96  // Chunk panics if n is less than 1.
    97  func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice] {
    98  	if n < 1 {
    99  		panic("cannot be less than 1")
   100  	}
   101  
   102  	return func(yield func(Slice) bool) {
   103  		for i := 0; i < len(s); i += n {
   104  			// Clamp the last chunk to the slice bound as necessary.
   105  			end := min(n, len(s[i:]))
   106  
   107  			// Set the capacity of each chunk so that appending to a chunk does
   108  			// not modify the original slice.
   109  			if !yield(s[i : i+end : i+end]) {
   110  				return
   111  			}
   112  		}
   113  	}
   114  }
   115
View as plain text