├── .github ├── dependabot.yaml └── workflows │ └── ci.yaml ├── LICENSE ├── README.md ├── chan.go ├── chan_test.go ├── count.go ├── count_test.go ├── func.go ├── func_test.go ├── go.mod ├── go.sum ├── group.go ├── group_test.go ├── io.go ├── io_test.go ├── iter.go ├── keyvalue.go ├── keyvalue_test.go ├── slice.go ├── slice_test.go ├── take.go └── take_test.go /.github/dependabot.yaml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | 4 | - package-ecosystem: "github-actions" 5 | directory: "/" 6 | schedule: 7 | # Check for updates to GitHub Actions every weekday. 8 | interval: "daily" 9 | 10 | - package-ecosystem: "gomod" 11 | directory: "/" 12 | schedule: 13 | # Check for updates to go modules every weekday. 14 | interval: "daily" 15 | -------------------------------------------------------------------------------- /.github/workflows/ci.yaml: -------------------------------------------------------------------------------- 1 | name: CI 2 | on: [push, pull_request] 3 | 4 | jobs: 5 | build_test: 6 | name: Build and Test 7 | strategy: 8 | matrix: 9 | go: ['1.18.x'] 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v3 13 | - uses: actions/setup-go@v3 14 | with: 15 | go-version: ${{ matrix.go }} 16 | stable: false 17 | - uses: actions/cache@v3 18 | with: 19 | path: ~/go/pkg/mod 20 | key: ubuntu-go-${{ hashFiles('**/go.sum') }} 21 | restore-keys: | 22 | ubuntu-go- 23 | - name: Test 24 | run: go test -mod readonly -race ./... 25 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Francesco Banconi 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![GoDoc](https://godoc.org/github.com/frankban/iterate?status.svg)](https://godoc.org/github.com/frankban/iterate) 2 | [![Build Status](https://github.com/frankban/iterate/actions/workflows/ci.yaml/badge.svg)](https://github.com/frankban/iterate/actions/workflows/ci.yaml) 3 | 4 | # iterate 5 | 6 | `go get github.com/frankban/iterate` 7 | 8 | Iterators and lazy evaluation in Go. 9 | This package uses generics, so requires Go 1.18 at least. 10 | 11 | For a complete API reference, see the [package documentation](https://pkg.go.dev/github.com/frankban/iterate). 12 | 13 | All functions work with implementations of the Iterator interface. 14 | 15 | ```go 16 | // Iterator is implemented by types producing values of type T. 17 | type Iterator[T any] interface { 18 | // Next advances the iterator. The next value can be then retrieved using 19 | // the Value method. False is returned when the iteration is done. Further 20 | // calls to Next should just return false with no other side effects. When 21 | // iterating produces an error, false is returned, and Err() returns the 22 | // error. 23 | Next() bool 24 | 25 | // Value returns the most recent value generated by a call to Next. It may 26 | // be called any number of times between calls to Next. If called after Next 27 | // has returned false, it returns the zero value. 28 | Value() T 29 | 30 | // Err returns the first error occurred while iterating. 31 | Err() error 32 | } 33 | ``` 34 | Implementations are typically used in for loops, for instance: 35 | ```go 36 | for iterator.Next() { 37 | v := iterator.Value() 38 | // Do something with v. 39 | } 40 | if err := iterator.Err(); err != nil { 41 | // Handle error. 42 | } 43 | ``` 44 | Depending on the implementation, producing values might lead to errors. For this 45 | reason it is important to always check Err() after iterating. 46 | -------------------------------------------------------------------------------- /chan.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // FromChannel returns an iterator producing values from the given channel. The 6 | // iteration stops when the channel is closed. The returned error is always nil. 7 | func FromChannel[T any](ch <-chan T) Iterator[T] { 8 | return &channelIterator[T]{ 9 | ch: ch, 10 | } 11 | } 12 | 13 | type channelIterator[T any] struct { 14 | ch <-chan T 15 | value T 16 | } 17 | 18 | // Next implements Iterator[T].Next. 19 | func (it *channelIterator[T]) Next() bool { 20 | val, ok := <-it.ch 21 | if ok { 22 | it.value = val 23 | return true 24 | } 25 | it.value = *new(T) 26 | return false 27 | } 28 | 29 | // Value implements Iterator[T].Value by returning values from the channel. 30 | func (it *channelIterator[T]) Value() T { 31 | return it.value 32 | } 33 | 34 | // Err implements Iterator[T].Err. The returned error is always nil. 35 | func (it *channelIterator[T]) Err() error { 36 | return nil 37 | } 38 | -------------------------------------------------------------------------------- /chan_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "testing" 7 | 8 | "github.com/go-quicktest/qt" 9 | 10 | it "github.com/frankban/iterate" 11 | ) 12 | 13 | func TestFromChan(t *testing.T) { 14 | ch := make(chan int) 15 | iter := it.FromChannel(ch) 16 | go func() { 17 | for i := 0; i < 5; i++ { 18 | ch <- i 19 | } 20 | close(ch) 21 | }() 22 | 23 | vs, err := it.ToSlice(iter) 24 | qt.Assert(t, qt.IsNil(err)) 25 | qt.Assert(t, qt.DeepEquals(vs, []int{0, 1, 2, 3, 4})) 26 | } 27 | -------------------------------------------------------------------------------- /count.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | import "golang.org/x/exp/constraints" 6 | 7 | // Count returns an iterator counting consecutive values from start to stop with 8 | // the given step. The returned error is always nil. 9 | // 10 | // For instance: 11 | // 12 | // counter := it.Count(0, 10, 1) 13 | // for counter.Next() { 14 | // v := counter.Value() 15 | // // v is 0, then 1, then 2 and so on till 9. 16 | // } 17 | // 18 | func Count(start, stop, step int) Iterator[int] { 19 | return &counter{ 20 | start: start - step, 21 | stop: stop, 22 | step: step, 23 | } 24 | } 25 | 26 | type counter struct { 27 | start, stop, step int 28 | } 29 | 30 | // Next implements Iterator[T].Next. 31 | func (it *counter) Next() bool { 32 | it.start += it.step 33 | if it.start == it.stop { 34 | it.start, it.stop = 0, it.step 35 | return false 36 | } 37 | return true 38 | } 39 | 40 | // Value implements Iterator[T].Value by returning the next number in the count. 41 | func (it *counter) Value() int { 42 | return it.start 43 | } 44 | 45 | // Err implements Iterator[T].Err. The returned error is always nil. 46 | func (it *counter) Err() error { 47 | return nil 48 | } 49 | 50 | // Enumerate returns an iterator that produces key/value pairs in which the keys 51 | // are iterator indexes (starting from 0) and the values are produced by the 52 | // given iterator. 53 | // 54 | // For instance: 55 | // 56 | // letters = it.Enumerate(it.FromSlice([]string{"a", "b", "c"})) 57 | // for letters.Next() { 58 | // kv := letters.Value() 59 | // // kv is (0, "a"), then (1, "b"), then (2, "c") 60 | // } 61 | // 62 | func Enumerate[T any](it Iterator[T]) Iterator[KeyValue[int, T]] { 63 | return &enumerator[T]{ 64 | source: it, 65 | idx: -1, 66 | } 67 | } 68 | 69 | type enumerator[T any] struct { 70 | source Iterator[T] 71 | idx int 72 | } 73 | 74 | // Next implements Iterator[T].Next. 75 | func (it *enumerator[T]) Next() bool { 76 | for it.source.Next() { 77 | it.idx++ 78 | return true 79 | } 80 | it.idx = 0 81 | return false 82 | } 83 | 84 | // Value implements Iterator[T].Next enumerating values from the source iterator. 85 | func (it *enumerator[T]) Value() KeyValue[int, T] { 86 | return KeyValue[int, T]{ 87 | Key: it.idx, 88 | Value: it.source.Value(), 89 | } 90 | } 91 | 92 | // Err implements Iterator[T].Err by propagating the error from the source 93 | // iterator. 94 | func (it *enumerator[T]) Err() error { 95 | return it.source.Err() 96 | } 97 | 98 | // Sum sums the values produced by the given iterator. 99 | // 100 | // For instance: 101 | // 102 | // sum, err := it.Sum(it.Count(1, 11, 1)) // sum is 55 103 | // sum, err := it.Sum(it.FromSlice([]string{"hello", " ", "world"})) // sum is "hello world" 104 | // 105 | func Sum[T constraints.Ordered](it Iterator[T]) (sum T, err error) { 106 | return Reduce(it, func(a, v T) T { 107 | return a + v 108 | }, *new(T)) 109 | } 110 | -------------------------------------------------------------------------------- /count_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "fmt" 7 | "testing" 8 | 9 | "github.com/frankban/iterate" 10 | "github.com/go-quicktest/qt" 11 | 12 | it "github.com/frankban/iterate" 13 | ) 14 | 15 | func TestCount(t *testing.T) { 16 | tests := []struct { 17 | start, stop, step int 18 | want []int 19 | }{{ 20 | start: 0, 21 | stop: 10, 22 | step: 1, 23 | want: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, 24 | }, { 25 | start: 0, 26 | stop: 0, 27 | step: 0, 28 | want: nil, 29 | }, { 30 | start: 10, 31 | stop: -10, 32 | step: -5, 33 | want: []int{10, 5, 0, -5}, 34 | }} 35 | for i := range tests { 36 | test := tests[i] 37 | t.Run(fmt.Sprintf("%d, %d, %d\n", test.start, test.stop, test.step), func(t *testing.T) { 38 | counter := it.Count(test.start, test.stop, test.step) 39 | got, err := it.ToSlice(counter) 40 | qt.Assert(t, qt.IsNil(err)) 41 | qt.Assert(t, qt.DeepEquals(got, test.want)) 42 | 43 | // Further calls to next return false and produce the zero value. 44 | qt.Assert(t, qt.IsFalse(counter.Next())) 45 | qt.Assert(t, qt.Equals(counter.Value(), 0)) 46 | }) 47 | } 48 | } 49 | 50 | func TestCountSkipValue(t *testing.T) { 51 | counter := it.Count(0, 100, 3) 52 | counter.Next() 53 | counter.Next() 54 | counter.Next() 55 | qt.Assert(t, qt.Equals(counter.Value(), 6)) 56 | qt.Assert(t, qt.Equals(counter.Value(), 6)) 57 | qt.Assert(t, qt.IsNil(counter.Err())) 58 | } 59 | 60 | func TestEnumerate(t *testing.T) { 61 | iter := it.Enumerate(it.FromSlice([]string{"a", "b", "c"})) 62 | got, err := it.ToSlice(iter) 63 | qt.Assert(t, qt.IsNil(err)) 64 | qt.Assert(t, qt.DeepEquals(got, []it.KeyValue[int, string]{{ 65 | Key: 0, 66 | Value: "a", 67 | }, { 68 | Key: 1, 69 | Value: "b", 70 | }, { 71 | Key: 2, 72 | Value: "c", 73 | }})) 74 | 75 | // Further calls to next return false and produce the zero value. 76 | qt.Assert(t, qt.IsFalse(iter.Next())) 77 | qt.Assert(t, qt.DeepEquals(iter.Value(), it.KeyValue[int, string]{})) 78 | } 79 | 80 | func TestEnumerateSkipValue(t *testing.T) { 81 | iter := it.Enumerate(it.FromSlice([]string{"a", "b", "c"})) 82 | iter.Next() 83 | iter.Next() 84 | want := it.KeyValue[int, string]{1, "b"} 85 | qt.Assert(t, qt.DeepEquals(iter.Value(), want)) 86 | qt.Assert(t, qt.DeepEquals(iter.Value(), want)) 87 | qt.Assert(t, qt.IsNil(iter.Err())) 88 | } 89 | 90 | func TestEnumerateError(t *testing.T) { 91 | iter := it.Enumerate[string](&errorIterator[string]{v: "v"}) 92 | got, err := it.ToSlice(iter) 93 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 94 | qt.Assert(t, qt.DeepEquals(got, []iterate.KeyValue[int, string]{{ 95 | Value: "v", 96 | }})) 97 | 98 | } 99 | 100 | func TestSum(t *testing.T) { 101 | sum1, err := it.Sum(it.FromSlice([]int{1, -1})) 102 | qt.Assert(t, qt.IsNil(err)) 103 | qt.Assert(t, qt.Equals(sum1, 0)) 104 | 105 | sum2, err := it.Sum(it.Count(1, 11, 1)) 106 | qt.Assert(t, qt.IsNil(err)) 107 | qt.Assert(t, qt.Equals(sum2, 55)) 108 | 109 | sum3, err := it.Sum(it.FromSlice([]float64{0.0, 0.0})) 110 | qt.Assert(t, qt.IsNil(err)) 111 | qt.Assert(t, qt.Equals(sum3, 0.0)) 112 | 113 | sum4, err := it.Sum(it.FromSlice([]string{"hello", " ", "world"})) 114 | qt.Assert(t, qt.IsNil(err)) 115 | qt.Assert(t, qt.Equals(sum4, "hello world")) 116 | } 117 | -------------------------------------------------------------------------------- /func.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // Filter returns an iterator producing values from the given interator, for 6 | // which predicate(v) is true. 7 | func Filter[T any](it Iterator[T], predicate func(v T) bool) Iterator[T] { 8 | return &filter[T]{ 9 | Iterator: it, 10 | predicate: predicate, 11 | } 12 | } 13 | 14 | type filter[T any] struct { 15 | Iterator[T] 16 | predicate func(v T) bool 17 | } 18 | 19 | // Next implements Iterator[T].Next by producing the next value satisfying the 20 | // predicate. 21 | func (it *filter[T]) Next() bool { 22 | for it.Iterator.Next() { 23 | if it.predicate(it.Iterator.Value()) { 24 | return true 25 | } 26 | } 27 | return false 28 | } 29 | 30 | // Map returns an iterator that computes the given function using values from 31 | // the given iterator. 32 | func Map[S, D any](source Iterator[S], f func(v S) D) Iterator[D] { 33 | return &mapper[S, D]{ 34 | Iterator: source, 35 | f: f, 36 | } 37 | } 38 | 39 | type mapper[S, D any] struct { 40 | Iterator[S] 41 | f func(v S) D 42 | } 43 | 44 | // Value implements Iterator[T].Value by returning the next corresponding value. 45 | func (it *mapper[S, D]) Value() D { 46 | return it.f(it.Iterator.Value()) 47 | } 48 | 49 | // Reduce applies f cumulatively to the values of the given iterator, from left 50 | // to right, so as to reduce the iterable to a single value. The first argument 51 | // is the accumulated value and the second argument is the value from the 52 | // iterator. 53 | // 54 | // For instance, for calculating the overall length of a slice of strings: 55 | // 56 | // iter := it.FromSlice([]string{"hello", "world"}) 57 | // length, err := it.Reduce(iter, func(a int, v string) int { 58 | // return a + len(v) 59 | // }, 0) 60 | // 61 | func Reduce[T, A any](it Iterator[T], f func(a A, v T) A, initial A) (A, error) { 62 | for it.Next() { 63 | initial = f(initial, it.Value()) 64 | } 65 | return initial, it.Err() 66 | } 67 | 68 | // Chain returns an iterator producing values from the concatenation of all the 69 | // given iterators. The iteration is stopped when all iterators are consumed or 70 | // when any of them has an error. 71 | func Chain[T any](base Iterator[T], others ...Iterator[T]) Iterator[T] { 72 | return &chain[T]{ 73 | Iterator: base, 74 | others: others, 75 | } 76 | } 77 | 78 | type chain[T any] struct { 79 | Iterator[T] 80 | others []Iterator[T] 81 | } 82 | 83 | // Next implements Iterator[T].Next by producing values until all iterators 84 | // are consumed. 85 | func (it *chain[T]) Next() bool { 86 | if it.Iterator.Next() { 87 | return true 88 | } 89 | if len(it.others) == 0 || it.Iterator.Err() != nil { 90 | return false 91 | } 92 | it.Iterator, it.others = it.others[0], it.others[1:] 93 | return it.Next() 94 | } 95 | 96 | // Repeat returns an iterator repeating values from the given iterator 97 | // endlessly. 98 | func Repeat[T any](it Iterator[T]) Iterator[T] { 99 | return &repeater[T]{ 100 | source: it, 101 | idx: -1, 102 | } 103 | } 104 | 105 | type repeater[T any] struct { 106 | source Iterator[T] 107 | idx int 108 | values []T 109 | } 110 | 111 | // Next implements Iterator[T].Next. 112 | func (it *repeater[T]) Next() bool { 113 | if it.source.Next() { 114 | it.values = append(it.values, it.source.Value()) 115 | it.idx++ 116 | return true 117 | } 118 | if len(it.values) == 0 || it.source.Err() != nil { 119 | return false 120 | } 121 | if it.idx+1 == len(it.values) { 122 | it.idx = 0 123 | } else { 124 | it.idx++ 125 | } 126 | return true 127 | } 128 | 129 | // Value implements Iterator[T].Value by returning values endlessly. 130 | func (it *repeater[T]) Value() T { 131 | return it.values[it.idx] 132 | } 133 | 134 | // Err implements Iterator[T].Err by propagating the error from the source 135 | // iterator. 136 | func (it *repeater[T]) Err() error { 137 | return it.source.Err() 138 | } 139 | 140 | // Tee returns an iterator that causes the given function to be called each time 141 | // a value is produced. 142 | func Tee[T any](it Iterator[T], f func(v T)) Iterator[T] { 143 | return Map(it, func(v T) T { 144 | f(v) 145 | return v 146 | }) 147 | } 148 | -------------------------------------------------------------------------------- /func_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "errors" 7 | "strings" 8 | "testing" 9 | "unicode" 10 | 11 | "github.com/go-quicktest/qt" 12 | 13 | it "github.com/frankban/iterate" 14 | ) 15 | 16 | func TestFilter(t *testing.T) { 17 | iter := it.FromSlice([]string{"these", "are", "the", "voyages"}) 18 | iter = it.Filter(iter, func(v string) bool { 19 | return len(v) == 3 20 | }) 21 | var vs []string 22 | for iter.Next() { 23 | vs = append(vs, iter.Value()) 24 | } 25 | qt.Assert(t, qt.IsNil(iter.Err())) 26 | qt.Assert(t, qt.DeepEquals(vs, []string{"are", "the"})) 27 | 28 | // Further calls to next return false and produce the zero value. 29 | qt.Assert(t, qt.IsFalse(iter.Next())) 30 | qt.Assert(t, qt.Equals(iter.Value(), "")) 31 | } 32 | 33 | func TestFilterSkipValue(t *testing.T) { 34 | iter := it.Filter(it.Count(10, 0, -1), func(v int) bool { 35 | return v%2 == 0 36 | }) 37 | iter.Next() 38 | iter.Next() 39 | 40 | qt.Assert(t, qt.Equals(iter.Value(), 8)) 41 | qt.Assert(t, qt.Equals(iter.Value(), 8)) 42 | qt.Assert(t, qt.IsNil(iter.Err())) 43 | } 44 | 45 | func TestFilterError(t *testing.T) { 46 | var iter it.Iterator[rune] = &errorIterator[rune]{ 47 | v: 'r', 48 | } 49 | iter = it.Filter(iter, unicode.IsLower) 50 | var vs []rune 51 | for iter.Next() { 52 | vs = append(vs, iter.Value()) 53 | } 54 | qt.Assert(t, qt.ErrorMatches(iter.Err(), "bad wolf")) 55 | qt.Assert(t, qt.DeepEquals(vs, []rune{'r'})) 56 | } 57 | 58 | func TestMap(t *testing.T) { 59 | iter := it.FromSlice([]string{"these", "are", "the", "voyages"}) 60 | iter = it.Map(iter, strings.ToUpper) 61 | got, err := it.ToSlice(iter) 62 | qt.Assert(t, qt.IsNil(err)) 63 | qt.Assert(t, qt.DeepEquals(got, []string{"THESE", "ARE", "THE", "VOYAGES"})) 64 | 65 | // Further calls to next return false and produce the zero value. 66 | qt.Assert(t, qt.IsFalse(iter.Next())) 67 | qt.Assert(t, qt.Equals(iter.Value(), "")) 68 | } 69 | 70 | func TestMapDifferentTypes(t *testing.T) { 71 | type rectangle struct { 72 | x, y int 73 | } 74 | 75 | iter := it.FromSlice([]rectangle{{ 76 | x: 1, y: 2, 77 | }, { 78 | x: 4, y: 5, 79 | }, { 80 | x: 10, y: 20, 81 | }}) 82 | areas := it.Map(iter, func(v rectangle) int { 83 | return v.x * v.y 84 | }) 85 | got, err := it.ToSlice(areas) 86 | qt.Assert(t, qt.IsNil(err)) 87 | qt.Assert(t, qt.DeepEquals(got, []int{2, 20, 200})) 88 | } 89 | 90 | func TestMapError(t *testing.T) { 91 | iter := it.Chain[int](it.Count(1, 5, 1), &errorIterator[int]{v: 5}) 92 | got, err := it.ToSlice(it.Map(iter, func(v int) int { 93 | return v * v 94 | })) 95 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 96 | qt.Assert(t, qt.DeepEquals(got, []int{1, 4, 9, 16, 25})) 97 | } 98 | 99 | func TestReduce(t *testing.T) { 100 | iter := it.FromSlice([]string{"hello", "world"}) 101 | length, err := it.Reduce(iter, func(a int, v string) int { 102 | return a + len(v) 103 | }, 0) 104 | qt.Assert(t, qt.IsNil(err)) 105 | qt.Assert(t, qt.Equals(length, 10)) 106 | } 107 | 108 | func TestReduceError(t *testing.T) { 109 | iter := it.Chain[int](it.Count(1, 5, 1), &errorIterator[int]{v: 5}) 110 | got, err := it.Reduce(iter, func(a, v int) int { 111 | return a + v 112 | }, 0) 113 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 114 | qt.Assert(t, qt.DeepEquals(got, 15)) 115 | } 116 | 117 | func TestChain(t *testing.T) { 118 | iter := it.Chain(it.FromSlice([]string{"1", "2"}), it.FromSlice([]string{"3", "4"})) 119 | got, err := it.ToSlice(iter) 120 | qt.Assert(t, qt.IsNil(err)) 121 | qt.Assert(t, qt.DeepEquals(got, []string{"1", "2", "3", "4"})) 122 | 123 | // Further calls to next return false and produce the zero value. 124 | qt.Assert(t, qt.IsFalse(iter.Next())) 125 | qt.Assert(t, qt.Equals(iter.Value(), "")) 126 | } 127 | 128 | func TestChainError(t *testing.T) { 129 | iter := it.Chain[int](it.Count(1, 5, 1), &errorIterator[int]{v: 5}) 130 | got, err := it.ToSlice(iter) 131 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 132 | qt.Assert(t, qt.DeepEquals(got, []int{1, 2, 3, 4, 5})) 133 | } 134 | 135 | func TestRepeat(t *testing.T) { 136 | iter := it.Limit(it.Repeat(it.Count(3, 0, -1)), 9) 137 | got, err := it.ToSlice(iter) 138 | qt.Assert(t, qt.IsNil(err)) 139 | qt.Assert(t, qt.DeepEquals(got, []int{3, 2, 1, 3, 2, 1, 3, 2, 1})) 140 | } 141 | 142 | func TestRepeatSkipValue(t *testing.T) { 143 | iter := it.Limit(it.Repeat(it.Count(3, 0, -1)), 9) 144 | iter.Next() 145 | iter.Next() 146 | qt.Assert(t, qt.Equals(iter.Value(), 2)) 147 | qt.Assert(t, qt.Equals(iter.Value(), 2)) 148 | qt.Assert(t, qt.IsNil(iter.Err())) 149 | } 150 | 151 | func TestRepeatError(t *testing.T) { 152 | iter := it.Repeat[string](&errorIterator[string]{v: "v"}) 153 | got, err := it.ToSlice(iter) 154 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 155 | qt.Assert(t, qt.DeepEquals(got, []string{"v"})) 156 | } 157 | 158 | // errorIterator is an iterator returning the given value and then an error. 159 | type errorIterator[T any] struct { 160 | v T 161 | calls int 162 | } 163 | 164 | func (it *errorIterator[T]) Next() bool { 165 | it.calls++ 166 | if it.calls == 1 { 167 | return true 168 | } 169 | it.v = *new(T) 170 | return false 171 | } 172 | 173 | func (it *errorIterator[T]) Value() T { 174 | return it.v 175 | } 176 | 177 | func (it *errorIterator[T]) Err() error { 178 | if it.calls > 1 { 179 | return errors.New("bad wolf") 180 | } 181 | return nil 182 | } 183 | 184 | func TestTee(t *testing.T) { 185 | var vs []int 186 | iter := it.Tee(it.Count(0, 10, 2), func(v int) { 187 | vs = append(vs, v) 188 | }) 189 | 190 | v, err := it.Next(iter) 191 | qt.Assert(t, qt.IsNil(err)) 192 | qt.Assert(t, qt.Equals(v, 0)) 193 | qt.Assert(t, qt.DeepEquals(vs, []int{0})) 194 | 195 | s, err := it.ToSlice(iter) 196 | qt.Assert(t, qt.IsNil(err)) 197 | qt.Assert(t, qt.DeepEquals(s, []int{2, 4, 6, 8})) 198 | qt.Assert(t, qt.DeepEquals(vs, []int{0, 2, 4, 6, 8})) 199 | } 200 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/frankban/iterate 2 | 3 | go 1.18 4 | 5 | require ( 6 | github.com/go-quicktest/qt v0.1.0 7 | golang.org/x/exp v0.0.0-20220321173239-a90fa8a75705 8 | ) 9 | 10 | require ( 11 | github.com/google/go-cmp v0.5.9 // indirect 12 | github.com/kr/pretty v0.3.0 // indirect 13 | github.com/kr/text v0.2.0 // indirect 14 | github.com/rogpeppe/go-internal v1.6.1 // indirect 15 | ) 16 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= 2 | github.com/go-quicktest/qt v0.1.0 h1:CUHOPpa5RfvcwLKo0yuz9ilHAnjySCoDZqY8n+oUcEQ= 3 | github.com/go-quicktest/qt v0.1.0/go.mod h1:rt7en+a0CcXVHA5cYg6YvOtRBsQflrHWbWUQKIrtl14= 4 | github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= 5 | github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 6 | github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= 7 | github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= 8 | github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= 9 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= 10 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= 11 | github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= 12 | github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= 13 | github.com/rogpeppe/go-internal v1.6.1 h1:/FiVV8dS/e+YqF2JvO3yXRFbBLTIuSDkuC7aBOAvL+k= 14 | github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= 15 | golang.org/x/exp v0.0.0-20220321173239-a90fa8a75705 h1:ba9YlqfDGTTQ5aZ2fwOoQ1hf32QySyQkR6ODGDzHlnE= 16 | golang.org/x/exp v0.0.0-20220321173239-a90fa8a75705/go.mod h1:lgLbSvA5ygNOMpwM/9anMpWVlVJ7Z+cHWq/eFuinpGE= 17 | gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 18 | gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= 19 | -------------------------------------------------------------------------------- /group.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // GroupBy returns an iterator returning key/value pairs in which the key is the 6 | // key used for grouping elements using the given function, and the value is an 7 | // iterator of values with that key. 8 | // 9 | // For instance: 10 | // 11 | // words := it.FromSlice([]string{"a", "be", "it", "no", "hello", "the", "are"}) 12 | // // Group words by length. 13 | // groups := it.GroupBy(words, func(v string) int { 14 | // return len(v) 15 | // }) 16 | // for groups.Next() { 17 | // kv := groups.Value() 18 | // // kv is (1, Iterator("a")), then (2, Iterator("be", "it", "no")), 19 | // // then (5, Iterator("hello")), then (3, Iterator("the", "are")) 20 | // } 21 | // 22 | func GroupBy[T any, K comparable](it Iterator[T], f func(v T) K) Iterator[KeyValue[K, Iterator[T]]] { 23 | return &grouper[T, K]{ 24 | source: it, 25 | f: f, 26 | pendingValues: make(map[int][]T), 27 | } 28 | } 29 | 30 | type grouper[T any, K comparable] struct { 31 | source Iterator[T] 32 | f func(v T) K 33 | id int 34 | key K 35 | pendingValues map[int][]T 36 | iter Iterator[T] 37 | } 38 | 39 | // Next implements Iterator[T].Next by iterating over groups. 40 | func (it *grouper[T, K]) Next() bool { 41 | if !it.next() { 42 | it.iter = nil 43 | it.key = *new(K) 44 | return false 45 | } 46 | if it.iter != nil { 47 | return true 48 | } 49 | return it.Next() 50 | } 51 | 52 | func (it *grouper[T, K]) next() bool { 53 | if !it.source.Next() { 54 | // The iteration is done. 55 | return false 56 | } 57 | 58 | val := it.source.Value() 59 | key := it.f(val) 60 | if it.id == 0 || key != it.key { 61 | // We either are at the beginning of the iteration, or the key just 62 | // changed. Store a new iterator to be returned on the next call to 63 | // Value. 64 | it.id++ 65 | it.iter = &groupKeyIterator[T, K]{ 66 | source: it, 67 | id: it.id, 68 | } 69 | it.key = key 70 | } 71 | 72 | // Store the produced value waiting for group iterators to retrieve it. 73 | it.pendingValues[it.id] = append(it.pendingValues[it.id], val) 74 | return true 75 | } 76 | 77 | // Value implements Iterator[T].Value by iterating over groups. 78 | func (it *grouper[T, K]) Value() KeyValue[K, Iterator[T]] { 79 | return KeyValue[K, Iterator[T]]{ 80 | Key: it.key, 81 | Value: it.iter, 82 | } 83 | } 84 | 85 | // Err implements Iterator[T].Err by propagating the error from the source 86 | // iterator. 87 | func (it *grouper[T, K]) Err() error { 88 | return it.source.Err() 89 | } 90 | 91 | // groupKeyIterator is the iterator returned for generating values for a 92 | // specific group key. 93 | type groupKeyIterator[T any, K comparable] struct { 94 | source *grouper[T, K] 95 | id int 96 | value T 97 | } 98 | 99 | // Next implements Iterator[T].Next. 100 | func (it *groupKeyIterator[T, K]) Next() bool { 101 | for { 102 | // Check whether there are pending values already. 103 | if len(it.source.pendingValues[it.id]) > 0 { 104 | it.value, it.source.pendingValues[it.id] = it.source.pendingValues[it.id][0], it.source.pendingValues[it.id][1:] 105 | return true 106 | } 107 | 108 | // Check whether the grouper is still iterating over this id, in which case we 109 | // can progress the iteration further and retry. 110 | if it.source.id == it.id && it.source.next() { 111 | continue 112 | } 113 | 114 | // The iteration for this group is over. 115 | it.value = *new(T) 116 | return false 117 | } 118 | } 119 | 120 | // Value implements Iterator[T].Value by returning values for a specific key. 121 | func (it *groupKeyIterator[T, K]) Value() T { 122 | return it.value 123 | } 124 | 125 | // Err implements Iterator[T].Err by propagating the error from the source 126 | // iterator. 127 | func (it *groupKeyIterator[T, K]) Err() error { 128 | return it.source.Err() 129 | } 130 | -------------------------------------------------------------------------------- /group_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "testing" 7 | 8 | "github.com/go-quicktest/qt" 9 | 10 | it "github.com/frankban/iterate" 11 | ) 12 | 13 | func TestGroupBy(t *testing.T) { 14 | // Group words by length. 15 | words := it.FromSlice([]string{ 16 | "a", 17 | "be", "it", "no", 18 | "hello", 19 | "the", "are", 20 | "be", "it", 21 | "again", 22 | }) 23 | groups := it.GroupBy(words, func(v string) int { 24 | return len(v) 25 | }) 26 | 27 | // Consume groups. 28 | 29 | // Forward the groups iterator a couple of times. 30 | qt.Assert(t, qt.IsTrue(groups.Next())) 31 | k, group1 := groups.Value().Split() 32 | qt.Assert(t, qt.Equals(k, 1)) 33 | 34 | qt.Assert(t, qt.IsTrue(groups.Next())) 35 | k, group2 := groups.Value().Split() 36 | qt.Assert(t, qt.Equals(k, 2)) 37 | 38 | // Consume the second group. 39 | qt.Assert(t, qt.IsTrue(group2.Next())) 40 | qt.Assert(t, qt.Equals(group2.Value(), "be")) 41 | 42 | qt.Assert(t, qt.IsTrue(group2.Next())) 43 | qt.Assert(t, qt.Equals(group2.Value(), "it")) 44 | 45 | qt.Assert(t, qt.IsTrue(group2.Next())) 46 | qt.Assert(t, qt.Equals(group2.Value(), "no")) 47 | 48 | // Produce another group. 49 | qt.Assert(t, qt.IsTrue(groups.Next())) 50 | k, group3 := groups.Value().Split() 51 | qt.Assert(t, qt.Equals(k, 5)) 52 | 53 | // Consume the first group. 54 | qt.Assert(t, qt.IsTrue(group1.Next())) 55 | qt.Assert(t, qt.Equals(group1.Value(), "a")) 56 | 57 | // Consume the third group. 58 | qt.Assert(t, qt.IsTrue(group3.Next())) 59 | qt.Assert(t, qt.Equals(group3.Value(), "hello")) 60 | 61 | qt.Assert(t, qt.IsFalse(group3.Next())) 62 | qt.Assert(t, qt.Equals(group3.Value(), "")) 63 | 64 | // Produce the fourth group. 65 | qt.Assert(t, qt.IsTrue(groups.Next())) 66 | k, group4 := groups.Value().Split() 67 | qt.Assert(t, qt.Equals(k, 3)) 68 | 69 | // Consume the fourth group. 70 | qt.Assert(t, qt.IsTrue(group4.Next())) 71 | qt.Assert(t, qt.Equals(group4.Value(), "the")) 72 | 73 | qt.Assert(t, qt.IsTrue(group4.Next())) 74 | qt.Assert(t, qt.Equals(group4.Value(), "are")) 75 | 76 | qt.Assert(t, qt.IsFalse(group4.Next())) 77 | qt.Assert(t, qt.Equals(group4.Value(), "")) 78 | 79 | // Produce the fifth group. 80 | qt.Assert(t, qt.IsTrue(groups.Next())) 81 | k, group5 := groups.Value().Split() 82 | qt.Assert(t, qt.Equals(k, 2)) 83 | 84 | // Produce the sixth group. 85 | qt.Assert(t, qt.IsTrue(groups.Next())) 86 | k, group6 := groups.Value().Split() 87 | qt.Assert(t, qt.Equals(k, 5)) 88 | 89 | // There are no other groups. 90 | qt.Assert(t, qt.IsFalse(groups.Next())) 91 | k, group7 := groups.Value().Split() 92 | qt.Assert(t, qt.Equals(k, 0)) 93 | qt.Assert(t, qt.IsNil(group7)) 94 | 95 | // Consume the sixth group. 96 | qt.Assert(t, qt.IsTrue(group6.Next())) 97 | qt.Assert(t, qt.Equals(group6.Value(), "again")) 98 | 99 | // Consume the fifth group. 100 | qt.Assert(t, qt.IsTrue(group5.Next())) 101 | qt.Assert(t, qt.Equals(group5.Value(), "be")) 102 | 103 | qt.Assert(t, qt.IsTrue(group5.Next())) 104 | qt.Assert(t, qt.Equals(group5.Value(), "it")) 105 | 106 | qt.Assert(t, qt.IsFalse(group5.Next())) 107 | qt.Assert(t, qt.Equals(group5.Value(), "")) 108 | 109 | // Check errors. 110 | qt.Assert(t, qt.IsNil(groups.Err())) 111 | qt.Assert(t, qt.IsNil(group1.Err())) 112 | qt.Assert(t, qt.IsNil(group2.Err())) 113 | qt.Assert(t, qt.IsNil(group3.Err())) 114 | qt.Assert(t, qt.IsNil(group4.Err())) 115 | qt.Assert(t, qt.IsNil(group5.Err())) 116 | qt.Assert(t, qt.IsNil(group4.Err())) 117 | } 118 | -------------------------------------------------------------------------------- /io.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | import ( 6 | "bufio" 7 | "io" 8 | ) 9 | 10 | // Lines returns an iterator producing lines from the given reader. 11 | func Lines(r io.Reader) Iterator[string] { 12 | return &lineReader{ 13 | scanner: *bufio.NewScanner(r), 14 | } 15 | } 16 | 17 | type lineReader struct { 18 | scanner bufio.Scanner 19 | } 20 | 21 | // Next implements Iterator[T].Next. 22 | func (it *lineReader) Next() bool { 23 | return it.scanner.Scan() 24 | } 25 | 26 | // Value implements Iterator[T].Value by returning the next line from the 27 | // reader. 28 | func (it *lineReader) Value() string { 29 | return it.scanner.Text() 30 | } 31 | 32 | // Err implements Iterator[T].Err by propagating any errors occurred while 33 | // reading, except for io.EOF. 34 | func (it *lineReader) Err() error { 35 | return it.scanner.Err() 36 | } 37 | 38 | // Bytes returns an iterator producing bytes from the given reader. 39 | func Bytes(r io.Reader) Iterator[byte] { 40 | return &byteReader{ 41 | r: *bufio.NewReader(r), 42 | } 43 | } 44 | 45 | type byteReader struct { 46 | r bufio.Reader 47 | b byte 48 | err error 49 | } 50 | 51 | // Next implements Iterator[T].Next by producing the next byte in the reader. 52 | func (it *byteReader) Next() bool { 53 | if it.err != nil { 54 | return false 55 | } 56 | it.b, it.err = it.r.ReadByte() 57 | if it.err == nil { 58 | return true 59 | } 60 | if it.err == io.EOF { 61 | it.err = nil 62 | } 63 | return false 64 | } 65 | 66 | // Value implements Iterator[T].Value by returning the next byte from the 67 | // reader. 68 | func (it *byteReader) Value() byte { 69 | return it.b 70 | } 71 | 72 | // Err implements Iterator[T].Err by propagating any errors occurred while 73 | // reading, except for io.EOF. 74 | func (it *byteReader) Err() error { 75 | return it.err 76 | } 77 | -------------------------------------------------------------------------------- /io_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "bytes" 7 | "errors" 8 | "strings" 9 | "testing" 10 | 11 | "github.com/go-quicktest/qt" 12 | 13 | it "github.com/frankban/iterate" 14 | ) 15 | 16 | func TestLines(t *testing.T) { 17 | r := strings.NewReader("hello\nworld") 18 | lines := it.Lines(r) 19 | 20 | qt.Assert(t, qt.IsTrue(lines.Next())) 21 | qt.Assert(t, qt.Equals(lines.Value(), "hello")) 22 | qt.Assert(t, qt.IsTrue(lines.Next())) 23 | qt.Assert(t, qt.Equals(lines.Value(), "world")) 24 | 25 | // Further calls to next return false and produce the zero value. 26 | qt.Assert(t, qt.IsFalse(lines.Next())) 27 | qt.Assert(t, qt.Equals(lines.Value(), "")) 28 | 29 | qt.Assert(t, qt.IsNil(lines.Err())) 30 | } 31 | 32 | func TestLinesError(t *testing.T) { 33 | lines := it.Lines(errReader{}) 34 | qt.Assert(t, qt.IsFalse(lines.Next())) 35 | qt.Assert(t, qt.Equals(lines.Value(), "")) 36 | qt.Assert(t, qt.ErrorMatches(lines.Err(), "bad wolf")) 37 | } 38 | 39 | func TestBytes(t *testing.T) { 40 | r := bytes.NewReader([]byte("hello\nworld")) 41 | b, err := it.ToSlice(it.Bytes(r)) 42 | qt.Assert(t, qt.IsNil(err)) 43 | qt.Assert(t, qt.Equals(string(b), "hello\nworld")) 44 | } 45 | 46 | func TestBytesError(t *testing.T) { 47 | bytes := it.Bytes(errReader{}) 48 | qt.Assert(t, qt.IsFalse(bytes.Next())) 49 | qt.Assert(t, qt.Equals(bytes.Value(), 0)) 50 | qt.Assert(t, qt.ErrorMatches(bytes.Err(), "bad wolf")) 51 | } 52 | 53 | // errReader is a io.Reader implementation that always return an error. 54 | type errReader struct{} 55 | 56 | func (errReader) Read(p []byte) (n int, err error) { 57 | return 0, errors.New("bad wolf") 58 | } 59 | -------------------------------------------------------------------------------- /iter.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // Iterator is implemented by types producing values of type T. Implementations 6 | // are typically used in for loops, for instance: 7 | // 8 | // for iterator.Next() { 9 | // v := iterator.Value() 10 | // // Do something with v. 11 | // } 12 | // if err := iterator.Err(); err != nil { 13 | // // Handle error. 14 | // } 15 | // 16 | // Depending on the implementation, producing values might lead to errors. For 17 | // this reason it is important to always check Err() after iterating. 18 | type Iterator[T any] interface { 19 | // Next advances the iterator. The next value can be then retrieved using 20 | // the Value method. False is returned when the iteration is done. Further 21 | // calls to Next should just return false with no other side effects. When 22 | // iterating produces an error, false is returned, and Err() returns the 23 | // error. 24 | Next() bool 25 | 26 | // Value returns the most recent value generated by a call to Next. It may 27 | // be called any number of times between calls to Next. If called after Next 28 | // has returned false, it returns the zero value. 29 | Value() T 30 | 31 | // Err returns the first error occurred while iterating. 32 | Err() error 33 | } 34 | -------------------------------------------------------------------------------- /keyvalue.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // KeyValue represents a key value pair. 6 | type KeyValue[K comparable, V any] struct { 7 | Key K 8 | Value V 9 | } 10 | 11 | // Split return the key and the value. 12 | func (kv KeyValue[K, V]) Split() (K, V) { 13 | return kv.Key, kv.Value 14 | } 15 | 16 | // Zip returns an iterator of key/value pairs produced by the given key/value 17 | // iterators. The shorter of the two iterators is used. An error is returned if 18 | // any of the two iterators returns an error. 19 | func Zip[K comparable, V any](keys Iterator[K], values Iterator[V]) Iterator[KeyValue[K, V]] { 20 | return &zipper[K, V]{ 21 | keys: keys, 22 | values: values, 23 | } 24 | } 25 | 26 | type zipper[K comparable, V any] struct { 27 | keys Iterator[K] 28 | values Iterator[V] 29 | stopped bool 30 | } 31 | 32 | // Next implements Iterator[T].Next. 33 | func (it *zipper[K, V]) Next() bool { 34 | if it.keys.Next() && it.values.Next() { 35 | return true 36 | } 37 | it.stopped = true 38 | return false 39 | } 40 | 41 | // Value implements Iterator[T].Value by returning key/value pairs. 42 | func (it *zipper[K, V]) Value() KeyValue[K, V] { 43 | if it.stopped { 44 | return KeyValue[K, V]{} 45 | } 46 | return KeyValue[K, V]{ 47 | Key: it.keys.Value(), 48 | Value: it.values.Value(), 49 | } 50 | } 51 | 52 | // Err implements Iterator[T].Err by propagating the error from the keys and 53 | // values iterators. 54 | func (it *zipper[K, V]) Err() error { 55 | if err := it.keys.Err(); err != nil { 56 | return err 57 | } 58 | return it.values.Err() 59 | } 60 | 61 | // Unzip returns a key iterator and a value iterator with pairs produced by the 62 | // given key/value iterator. 63 | func Unzip[K comparable, V any](kvs Iterator[KeyValue[K, V]]) (Iterator[K], Iterator[V]) { 64 | u := unzipper[K, V]{ 65 | kvs: kvs, 66 | } 67 | return u.iterators() 68 | } 69 | 70 | type unzipper[K comparable, V any] struct { 71 | kvs Iterator[KeyValue[K, V]] 72 | keys []K 73 | values []V 74 | } 75 | 76 | func (u *unzipper[K, V]) iterators() (Iterator[K], Iterator[V]) { 77 | return &keyIterator[K, V]{ 78 | u: u, 79 | }, &valueIterator[K, V]{ 80 | u: u, 81 | } 82 | } 83 | 84 | type keyIterator[K comparable, V any] struct { 85 | u *unzipper[K, V] 86 | key K 87 | } 88 | 89 | // Next implements Iterator[T].Next. 90 | func (it *keyIterator[K, V]) Next() bool { 91 | // TODO(frankban): make this thread safe. 92 | if len(it.u.keys) != 0 { 93 | it.key, it.u.keys = it.u.keys[0], it.u.keys[1:] 94 | return true 95 | } 96 | if it.u.kvs.Next() { 97 | kv := it.u.kvs.Value() 98 | it.key = kv.Key 99 | it.u.values = append(it.u.values, kv.Value) 100 | return true 101 | } 102 | it.key = *new(K) 103 | return false 104 | } 105 | 106 | // Value implements Iterator[T].Value by returning keys from the key/value 107 | // iterator stored in the unzipper. 108 | func (it *keyIterator[K, V]) Value() K { 109 | return it.key 110 | } 111 | 112 | // Err implements Iterator[T].Err by propagating the error from the key/value 113 | // source iterator. 114 | func (it *keyIterator[K, V]) Err() error { 115 | return it.u.kvs.Err() 116 | } 117 | 118 | type valueIterator[K comparable, V any] struct { 119 | u *unzipper[K, V] 120 | value V 121 | } 122 | 123 | // Next implements Iterator[T].Next. 124 | func (it *valueIterator[K, V]) Next() bool { 125 | // TODO(frankban): make this thread safe. 126 | if len(it.u.values) != 0 { 127 | it.value, it.u.values = it.u.values[0], it.u.values[1:] 128 | return true 129 | } 130 | if it.u.kvs.Next() { 131 | kv := it.u.kvs.Value() 132 | it.value = kv.Value 133 | it.u.keys = append(it.u.keys, kv.Key) 134 | return true 135 | } 136 | it.value = *new(V) 137 | return false 138 | } 139 | 140 | // Value implements Iterator[T].Value by returning values from the key/value 141 | // iterator stored in the unzipper. 142 | func (it *valueIterator[K, V]) Value() V { 143 | return it.value 144 | } 145 | 146 | // Err implements Iterator[T].Err by propagating the error from the key/value 147 | // source iterator. 148 | func (it *valueIterator[K, V]) Err() error { 149 | return it.u.kvs.Err() 150 | } 151 | 152 | // ToMap returns a map with the values produced by the given key/value iterator. 153 | // An error is returned if the iterator returns an error, in which case the 154 | // returned map includes the key/value pairs already consumed. 155 | func ToMap[K comparable, V any](it Iterator[KeyValue[K, V]]) (map[K]V, error) { 156 | m := make(map[K]V) 157 | for it.Next() { 158 | kv := it.Value() 159 | m[kv.Key] = kv.Value 160 | } 161 | return m, it.Err() 162 | } 163 | -------------------------------------------------------------------------------- /keyvalue_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "testing" 7 | 8 | "github.com/go-quicktest/qt" 9 | 10 | it "github.com/frankban/iterate" 11 | ) 12 | 13 | func TestKeyValueSplit(t *testing.T) { 14 | kv := it.KeyValue[int, string]{ 15 | Key: 42, 16 | Value: "engage", 17 | } 18 | k, v := kv.Split() 19 | qt.Assert(t, qt.Equals(k, 42)) 20 | qt.Assert(t, qt.Equals(v, "engage")) 21 | } 22 | 23 | func TestZip(t *testing.T) { 24 | keys := it.FromSlice([]string{"a", "b", "c"}) 25 | values := it.FromSlice([]int{4, 3, 2, 1}) 26 | iter := it.Zip(keys, values) 27 | 28 | got, err := it.ToSlice(iter) 29 | qt.Assert(t, qt.IsNil(err)) 30 | qt.Assert(t, qt.DeepEquals(got, []it.KeyValue[string, int]{{ 31 | Key: "a", 32 | Value: 4, 33 | }, { 34 | Key: "b", 35 | Value: 3, 36 | }, { 37 | Key: "c", 38 | Value: 2, 39 | }})) 40 | 41 | // Further calls to next return false and produce the zero value. 42 | qt.Assert(t, qt.IsFalse(iter.Next())) 43 | qt.Assert(t, qt.Equals(iter.Value(), it.KeyValue[string, int]{})) 44 | 45 | // The remaining value can still be retrieved. 46 | vs, err := it.ToSlice(values) 47 | qt.Assert(t, qt.IsNil(err)) 48 | qt.Assert(t, qt.DeepEquals(vs, []int{1})) 49 | } 50 | 51 | func TestZipError(t *testing.T) { 52 | keys := it.FromSlice([]string{"a", "b", "c"}) 53 | iter := it.Zip[string, int](keys, &errorIterator[int]{ 54 | v: 42, 55 | }) 56 | 57 | got, err := it.ToSlice(iter) 58 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 59 | qt.Assert(t, qt.DeepEquals(got, []it.KeyValue[string, int]{{ 60 | Key: "a", 61 | Value: 42, 62 | }})) 63 | 64 | } 65 | 66 | func TestUnzip(t *testing.T) { 67 | iter := it.FromSlice(makeKeyValues()) 68 | keys, values := it.Unzip(iter) 69 | 70 | qt.Assert(t, qt.IsTrue(keys.Next())) 71 | qt.Assert(t, qt.Equals(keys.Value(), 1)) 72 | 73 | qt.Assert(t, qt.IsTrue(values.Next())) 74 | qt.Assert(t, qt.Equals(values.Value(), "these")) 75 | 76 | qt.Assert(t, qt.IsTrue(values.Next())) 77 | qt.Assert(t, qt.Equals(values.Value(), "are")) 78 | 79 | qt.Assert(t, qt.IsTrue(values.Next())) 80 | qt.Assert(t, qt.Equals(values.Value(), "the")) 81 | 82 | qt.Assert(t, qt.IsTrue(keys.Next())) 83 | qt.Assert(t, qt.Equals(keys.Value(), 2)) 84 | 85 | qt.Assert(t, qt.IsTrue(values.Next())) 86 | qt.Assert(t, qt.Equals(values.Value(), "voyages")) 87 | 88 | qt.Assert(t, qt.IsFalse(values.Next())) 89 | qt.Assert(t, qt.Equals(values.Value(), "")) 90 | 91 | qt.Assert(t, qt.IsTrue(keys.Next())) 92 | qt.Assert(t, qt.Equals(keys.Value(), 42)) 93 | 94 | qt.Assert(t, qt.IsTrue(keys.Next())) 95 | qt.Assert(t, qt.Equals(keys.Value(), 47)) 96 | 97 | qt.Assert(t, qt.IsFalse(keys.Next())) 98 | qt.Assert(t, qt.Equals(keys.Value(), 0)) 99 | 100 | qt.Assert(t, qt.IsNil(keys.Err())) 101 | qt.Assert(t, qt.IsNil(values.Err())) 102 | } 103 | 104 | func TestUnzipSkipValues(t *testing.T) { 105 | iter := it.FromSlice(makeKeyValues()) 106 | keys, values := it.Unzip(iter) 107 | 108 | keys.Next() 109 | keys.Next() 110 | keys.Next() 111 | values.Next() 112 | values.Next() 113 | 114 | qt.Assert(t, qt.Equals(keys.Value(), 42)) 115 | qt.Assert(t, qt.Equals(keys.Value(), 42)) 116 | qt.Assert(t, qt.Equals(values.Value(), "are")) 117 | 118 | qt.Assert(t, qt.IsNil(keys.Err())) 119 | qt.Assert(t, qt.IsNil(values.Err())) 120 | } 121 | 122 | func TestUnzipError(t *testing.T) { 123 | keys, values := it.Unzip[int, string](&errorIterator[it.KeyValue[int, string]]{ 124 | v: it.KeyValue[int, string]{ 125 | Key: 1, 126 | Value: "engage", 127 | }, 128 | }) 129 | 130 | qt.Assert(t, qt.IsTrue(keys.Next())) 131 | qt.Assert(t, qt.Equals(keys.Value(), 1)) 132 | 133 | qt.Assert(t, qt.IsFalse(keys.Next())) 134 | qt.Assert(t, qt.Equals(keys.Value(), 0)) 135 | 136 | qt.Assert(t, qt.ErrorMatches(keys.Err(), "bad wolf")) 137 | // The values iterator is still able to produce values but it's already in 138 | // error. 139 | qt.Assert(t, qt.ErrorMatches(values.Err(), "bad wolf")) 140 | 141 | qt.Assert(t, qt.IsTrue(values.Next())) 142 | qt.Assert(t, qt.Equals(values.Value(), "engage")) 143 | 144 | qt.Assert(t, qt.IsFalse(values.Next())) 145 | qt.Assert(t, qt.Equals(values.Value(), "")) 146 | } 147 | 148 | func TestZipUnzipRoundtrip(t *testing.T) { 149 | s := makeKeyValues() 150 | got, err := it.ToSlice(it.Zip(it.Unzip(it.FromSlice(s)))) 151 | qt.Assert(t, qt.IsNil(err)) 152 | qt.Assert(t, qt.DeepEquals(got, s)) 153 | } 154 | 155 | func TestToMap(t *testing.T) { 156 | iter := it.FromSlice(makeKeyValues()) 157 | m, err := it.ToMap(iter) 158 | qt.Assert(t, qt.IsNil(err)) 159 | qt.Assert(t, qt.DeepEquals(m, map[int]string{ 160 | 1: "these", 161 | 2: "are", 162 | 42: "the", 163 | 47: "voyages", 164 | })) 165 | 166 | } 167 | 168 | func TestToMapError(t *testing.T) { 169 | m, err := it.ToMap[int, string](&errorIterator[it.KeyValue[int, string]]{ 170 | v: it.KeyValue[int, string]{ 171 | Key: 1, 172 | Value: "engage", 173 | }, 174 | }) 175 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 176 | qt.Assert(t, qt.DeepEquals(m, map[int]string{ 177 | 1: "engage", 178 | })) 179 | 180 | } 181 | 182 | func makeKeyValues() []it.KeyValue[int, string] { 183 | return []it.KeyValue[int, string]{{ 184 | Key: 1, 185 | Value: "these", 186 | }, { 187 | Key: 2, 188 | Value: "are", 189 | }, { 190 | Key: 42, 191 | Value: "the", 192 | }, { 193 | Key: 47, 194 | Value: "voyages", 195 | }} 196 | } 197 | -------------------------------------------------------------------------------- /slice.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // FromSlice returns an iterator producing values from the given slice. 6 | func FromSlice[T any](s []T) Iterator[T] { 7 | return &sliceIterator[T]{ 8 | s: s, 9 | first: true, 10 | } 11 | } 12 | 13 | type sliceIterator[T any] struct { 14 | s []T 15 | first bool 16 | } 17 | 18 | // Next implements Iterator[T].Next. 19 | func (it *sliceIterator[T]) Next() bool { 20 | if len(it.s) == 0 { 21 | return false 22 | } 23 | if it.first { 24 | it.first = false 25 | } else { 26 | it.s = it.s[1:] 27 | } 28 | return len(it.s) > 0 29 | } 30 | 31 | // Value implements Iterator[T].Value by returning values from a slice. 32 | func (it *sliceIterator[T]) Value() T { 33 | if len(it.s) == 0 { 34 | return *new(T) 35 | } 36 | return it.s[0] 37 | } 38 | 39 | // Err implements Iterator[T].Err. The returned error is always nil. 40 | func (it *sliceIterator[T]) Err() error { 41 | return nil 42 | } 43 | 44 | // ToSlice consumes the given iterator and returns a slice of produced values or 45 | // an error occurred while iterating. This function should not be used with 46 | // infinite iterators (see TakeWhile or Limit). 47 | func ToSlice[T any](it Iterator[T]) (s []T, err error) { 48 | for it.Next() { 49 | s = append(s, it.Value()) 50 | } 51 | return s, it.Err() 52 | } 53 | -------------------------------------------------------------------------------- /slice_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "strings" 7 | "testing" 8 | 9 | "github.com/go-quicktest/qt" 10 | 11 | it "github.com/frankban/iterate" 12 | ) 13 | 14 | func TestFromSlice(t *testing.T) { 15 | iter := it.FromSlice([]byte("these are the voyages")) 16 | var b strings.Builder 17 | for iter.Next() { 18 | b.WriteByte(iter.Value()) 19 | } 20 | qt.Assert(t, qt.IsNil(iter.Err())) 21 | qt.Assert(t, qt.Equals(b.String(), "these are the voyages")) 22 | 23 | // Further calls to next return false and produce the zero value. 24 | qt.Assert(t, qt.IsFalse(iter.Next())) 25 | qt.Assert(t, qt.Equals(iter.Value(), 0)) 26 | } 27 | 28 | func TestFromSliceSkipValue(t *testing.T) { 29 | iter := it.FromSlice([]int{1, 2, 3}) 30 | iter.Next() 31 | iter.Next() 32 | qt.Assert(t, qt.Equals(iter.Value(), 2)) 33 | qt.Assert(t, qt.Equals(iter.Value(), 2)) 34 | qt.Assert(t, qt.IsNil(iter.Err())) 35 | } 36 | 37 | func TestToSlice(t *testing.T) { 38 | // Let's take advantage of this for testing some composition as well. 39 | multipleOf3 := func(v int) bool { 40 | return v%3 == 0 41 | } 42 | lessThan500 := func(_, v int) bool { 43 | return v < 500 44 | } 45 | s, err := it.ToSlice(it.TakeWhile(it.Filter(it.Count(0, 20, 1), multipleOf3), lessThan500)) 46 | qt.Assert(t, qt.IsNil(err)) 47 | qt.Assert(t, qt.DeepEquals(s, []int{0, 3, 6, 9, 12, 15, 18})) 48 | } 49 | 50 | func TestToSliceError(t *testing.T) { 51 | s, err := it.ToSlice[string](&errorIterator[string]{ 52 | v: "engage", 53 | }) 54 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 55 | qt.Assert(t, qt.DeepEquals(s, []string{"engage"})) 56 | } 57 | 58 | func TestSliceRoundtrip(t *testing.T) { 59 | want := []string{"these", "are", "the", "voyages"} 60 | got, err := it.ToSlice(it.FromSlice(want)) 61 | qt.Assert(t, qt.IsNil(err)) 62 | qt.Assert(t, qt.DeepEquals(got, want)) 63 | } 64 | -------------------------------------------------------------------------------- /take.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate 4 | 5 | // DropWhile returns an iterator discarding values from the given iterator while 6 | // predicate(v) is true. 7 | func DropWhile[T any](it Iterator[T], predicate func(idx int, v T) bool) Iterator[T] { 8 | return &dropper[T]{ 9 | Iterator: it, 10 | predicate: predicate, 11 | } 12 | } 13 | 14 | type dropper[T any] struct { 15 | Iterator[T] 16 | predicate func(idx int, v T) bool 17 | idx int 18 | started bool 19 | } 20 | 21 | // Next implements Iterator[T].Next by discarding values while they satisfy the 22 | // predicate. 23 | func (it *dropper[T]) Next() bool { 24 | for it.Iterator.Next() { 25 | if !it.started && it.predicate(it.idx, it.Iterator.Value()) { 26 | it.idx++ 27 | continue 28 | } 29 | it.started = true 30 | return true 31 | } 32 | return false 33 | } 34 | 35 | // TakeWhile returns an iterator producing values from the given iterator while 36 | // predicate(v) is true. 37 | func TakeWhile[T any](it Iterator[T], predicate func(idx int, v T) bool) Iterator[T] { 38 | return &taker[T]{ 39 | source: it, 40 | predicate: predicate, 41 | } 42 | } 43 | 44 | type taker[T any] struct { 45 | source Iterator[T] 46 | predicate func(idx int, v T) bool 47 | idx int 48 | stopped bool 49 | } 50 | 51 | // Next implements Iterator[T].Next. 52 | func (it *taker[T]) Next() bool { 53 | if !it.stopped && it.source.Next() && it.predicate(it.idx, it.source.Value()) { 54 | it.idx++ 55 | return true 56 | } 57 | it.stopped = true 58 | return false 59 | } 60 | 61 | // Value implements Iterator[T].Value by returning values until they satisfy the 62 | // predicate. 63 | func (it *taker[T]) Value() T { 64 | if it.stopped { 65 | return *new(T) 66 | } 67 | return it.source.Value() 68 | } 69 | 70 | // Err implements Iterator[T].Err by propagating the error from the source 71 | // iterator. 72 | func (it *taker[T]) Err() error { 73 | return it.source.Err() 74 | } 75 | 76 | // Limit returns an iterator limiting the number of values returned by the given 77 | // iterator. 78 | func Limit[T any](it Iterator[T], limit int) Iterator[T] { 79 | return TakeWhile(it, func(idx int, v T) bool { 80 | return idx < limit 81 | }) 82 | } 83 | 84 | // Next returns the next value produced by the iterator. 85 | func Next[T any](it Iterator[T]) (T, error) { 86 | it.Next() 87 | return it.Value(), it.Err() 88 | } 89 | -------------------------------------------------------------------------------- /take_test.go: -------------------------------------------------------------------------------- 1 | // Licensed under the MIT license, see LICENSE file for details. 2 | 3 | package iterate_test 4 | 5 | import ( 6 | "testing" 7 | 8 | "github.com/go-quicktest/qt" 9 | 10 | it "github.com/frankban/iterate" 11 | ) 12 | 13 | func TestDropWhile(t *testing.T) { 14 | iter := it.DropWhile(it.Count(0, 10, 1), func(idx, v int) bool { 15 | return idx < 2 || idx > 4 || v == 2 || v == 7 16 | }) 17 | 18 | vs, err := it.ToSlice(iter) 19 | qt.Assert(t, qt.IsNil(err)) 20 | qt.Assert(t, qt.DeepEquals(vs, []int{3, 4, 5, 6, 7, 8, 9})) 21 | } 22 | 23 | func TestDropWhileError(t *testing.T) { 24 | iter := it.DropWhile[string](&errorIterator[string]{ 25 | v: "ok", 26 | }, func(idx int, v string) bool { 27 | return v == "ok" 28 | }) 29 | 30 | vs, err := it.ToSlice(iter) 31 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 32 | qt.Assert(t, qt.IsNil(vs)) 33 | } 34 | 35 | func TestTakeWhile(t *testing.T) { 36 | iter := it.TakeWhile(it.Count(0, 20, 1), func(idx, v int) bool { 37 | return idx < 3 || idx > 4 || v == 3 || v == 7 38 | }) 39 | 40 | vs, err := it.ToSlice(iter) 41 | qt.Assert(t, qt.IsNil(err)) 42 | qt.Assert(t, qt.DeepEquals(vs, []int{0, 1, 2, 3})) 43 | 44 | // Consuming more items does not produce values. 45 | for i := 0; i < 5; i++ { 46 | v, err := it.Next(iter) 47 | qt.Assert(t, qt.IsNil(err)) 48 | qt.Assert(t, qt.Equals(v, 0)) 49 | } 50 | } 51 | 52 | func TestTakeWhileSkipValue(t *testing.T) { 53 | iter := it.TakeWhile(it.Count(0, 100, 10), func(idx, v int) bool { 54 | return idx < 5 55 | }) 56 | iter.Next() 57 | iter.Next() 58 | qt.Assert(t, qt.Equals(iter.Value(), 10)) 59 | qt.Assert(t, qt.Equals(iter.Value(), 10)) 60 | qt.Assert(t, qt.IsNil(iter.Err())) 61 | } 62 | 63 | func TestTakeWhileError(t *testing.T) { 64 | iter := it.TakeWhile[string](&errorIterator[string]{ 65 | v: "ok", 66 | }, func(idx int, v string) bool { 67 | return v == "ok" 68 | }) 69 | 70 | vs, err := it.ToSlice(iter) 71 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 72 | qt.Assert(t, qt.DeepEquals(vs, []string{"ok"})) 73 | } 74 | 75 | func TestLimit(t *testing.T) { 76 | iter := it.Limit(it.Count(0, 10, 1), 3) 77 | 78 | vs, err := it.ToSlice(iter) 79 | qt.Assert(t, qt.IsNil(err)) 80 | qt.Assert(t, qt.DeepEquals(vs, []int{0, 1, 2})) 81 | } 82 | 83 | func TestNext(t *testing.T) { 84 | iter := &errorIterator[string]{ 85 | v: "ok", 86 | } 87 | 88 | v, err := it.Next[string](iter) 89 | qt.Assert(t, qt.IsNil(err)) 90 | qt.Assert(t, qt.Equals(v, "ok")) 91 | 92 | v, err = it.Next[string](iter) 93 | qt.Assert(t, qt.ErrorMatches(err, "bad wolf")) 94 | qt.Assert(t, qt.Equals(v, "")) 95 | } 96 | --------------------------------------------------------------------------------