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  func 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.
    57  func 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.
    63  func 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.
    71  func 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.
    81  func 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.
    92  func 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
View as plain text