iter.mx raw

   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