├── .travis.yml ├── itertools_bench_test.go ├── itertools_test.go ├── README.md └── itertools.go /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | go: 1.1 3 | -------------------------------------------------------------------------------- /itertools_bench_test.go: -------------------------------------------------------------------------------- 1 | package itertools 2 | 3 | import "testing" 4 | 5 | func BenchmarkFilter(b *testing.B) { 6 | pred := func(i interface{}) bool { 7 | return i.(uint64)%2 == 1 8 | } 9 | 10 | for i := 0; i < b.N; i++ { 11 | Filter(pred, Uint64(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)) 12 | } 13 | } 14 | 15 | func BenchmarkNoFilter(b *testing.B) { 16 | for i := 0; i < b.N; i++ { 17 | input := []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} 18 | results := make([]uint64, 0) 19 | 20 | for _, n := range input { 21 | if n%2 == 1 { 22 | results = append(results, n) 23 | } 24 | } 25 | } 26 | } 27 | 28 | func BenchmarkMap(b *testing.B) { 29 | mapper := func(i interface{}) interface{} { 30 | return len(i.(string)) 31 | } 32 | 33 | for i := 0; i < b.N; i++ { 34 | Map(mapper, New("a", "ab", "abc", "abcd")) 35 | } 36 | } 37 | 38 | func BenchmarkNoMap(b *testing.B) { 39 | for i := 0; i < b.N; i++ { 40 | input := []string{"a", "ab", "abc", "abcd"} 41 | results := make([]int, 0) 42 | 43 | for _, w := range input { 44 | results = append(results, len(w)) 45 | } 46 | } 47 | } 48 | 49 | func BenchmarkReduce(b *testing.B) { 50 | summer := func(memo interface{}, el interface{}) interface{} { 51 | return memo.(float64) + el.(float64) 52 | } 53 | 54 | for i := 0; i < b.N; i++ { 55 | Reduce(Float64(.1, .2, .3, .22), summer, float64(0)) 56 | } 57 | } 58 | 59 | func BenchmarkNoReduce(b *testing.B) { 60 | for i := 0; i < b.N; i++ { 61 | input := []float64{.1, .2, .3, .22} 62 | result := 0.0 63 | 64 | for _, n := range input { 65 | result += n 66 | } 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /itertools_test.go: -------------------------------------------------------------------------------- 1 | package itertools 2 | 3 | import ( 4 | "testing" 5 | "reflect" 6 | ) 7 | 8 | // Test iterators for element equality. Allow it1 to be longer than it2 9 | func testIter(t *testing.T, it1, it2 Iter) { 10 | t.Log("Start") 11 | for el1 := range it1 { 12 | if el2, ok := <- it2; !ok { 13 | t.Error("it2 shorter than it1!", el1) 14 | return 15 | } else if !reflect.DeepEqual(el1, el2) { 16 | t.Error("Elements are not equal", el1, el2) 17 | } else { 18 | t.Log(el1, el2) 19 | } 20 | } 21 | t.Log("Stop") 22 | } 23 | 24 | // Test iterators for element equality. Don't allow it1 to be longer than it2 25 | func testIterEq(t *testing.T, it1, it2 Iter) { 26 | t.Log("Start") 27 | for el1 := range it1 { 28 | if el2, ok := <- it2; !ok { 29 | t.Error("it2 shorter than it1!", el1) 30 | return 31 | } else if !reflect.DeepEqual(el1, el2) { 32 | t.Error("Elements are not equal", el1, el2) 33 | } else { 34 | t.Log(el1, el2) 35 | } 36 | } 37 | if el2, ok := <- it2; ok { 38 | t.Error("it1 shorter than it2!", el2) 39 | } 40 | t.Log("Stop") 41 | } 42 | 43 | 44 | func TestList(t *testing.T) { 45 | list := List(New(1,2,3)) 46 | if !reflect.DeepEqual(list, []interface{}{1,2,3}) { 47 | t.Error("List didn't make a list", list) 48 | } 49 | } 50 | 51 | func TestCount(t *testing.T) { 52 | testIter(t, New(1,2,3,4,5,6,7,8,9), Count(1)) 53 | } 54 | 55 | func TestCycle(t *testing.T) { 56 | testIter(t, New("a", "b", "ccc", "a", "b", "ccc", "a"), Cycle(New("a", "b", "ccc"))) 57 | } 58 | 59 | func TestRepeat(t *testing.T) { 60 | testIterEq(t, Uint64(100, 100, 100, 100), Repeat(uint64(100), 4)) 61 | testIter(t, Uint64(100, 100, 100, 100), Repeat(uint64(100))) 62 | } 63 | 64 | func TestChain(t *testing.T) { 65 | testIterEq(t, Int32(1,2,3,4,5,5,4,3,2,1,100), Chain(Int32(1,2,3,4,5), Int32(5,4,3,2,1), Int32(100))) 66 | } 67 | 68 | 69 | func TestDropWhile(t *testing.T) { 70 | pred := func (i interface{}) bool { 71 | return i.(int) < 10 72 | } 73 | testIter(t, New(10,11,12,13,14,15), DropWhile(pred, Count(0))) 74 | } 75 | 76 | func TestTakeWhile(t *testing.T) { 77 | pred := func (i interface{}) bool { 78 | return i.(string)[:3] == "abc" 79 | } 80 | testIterEq(t, New("abcdef", "abcdaj"), TakeWhile(pred, Cycle(New("abcdef", "abcdaj", "ajcde")))) 81 | } 82 | 83 | func TestFilter(t *testing.T) { 84 | pred := func (i interface{}) bool { 85 | return i.(uint64) % 2 == 1 86 | } 87 | testIterEq(t, Uint64(1,3,5,7,9), Filter(pred, Uint64(1,2,3,4,5,6,7,8,9,10))) 88 | testIterEq(t, Uint64(2,4,6,8,10), FilterFalse(pred, Uint64(1,2,3,4,5,6,7,8,9,10))) 89 | } 90 | 91 | func TestSlice(t *testing.T) { 92 | testIter(t, New(5,6,7,8,9,10), Slice(Count(0), 5)) 93 | testIterEq(t, New(2,3,4,5,6,7,8), Slice(Count(0), 2, 9)) 94 | testIterEq(t, New(3,6,9), Slice(Count(0), 3, 11, 3)) 95 | } 96 | 97 | func TestMap(t *testing.T) { 98 | mapper := func (i interface{}) interface{} { 99 | return len(i.(string)) 100 | } 101 | testIterEq(t, New(1,2,3,4), Map(mapper, New("a", "ab", "abc", "abcd"))) 102 | } 103 | 104 | func TestMultiMap(t *testing.T) { 105 | multiMapper := func (is ...interface{}) interface{} { 106 | var s float64 107 | for _, i := range is { 108 | s += i.(float64) 109 | } 110 | return s 111 | } 112 | testIterEq(t, Float64(10.4, 3.2), MultiMap(multiMapper, Float64(5.2, 1.6, 2.2), Float64(5.2, 1.0), Float64(0, 0.6, 0))) 113 | } 114 | 115 | func TestZip(t *testing.T) { 116 | a, b, c := []interface{}{1,"a"}, []interface{}{2,nil}, []interface{}{3,nil} 117 | test1, test2 := New(a), New(a,b,c) 118 | 119 | testIterEq(t, test1, Zip(Count(1), New("a"))) 120 | testIterEq(t, test2, ZipLongest(Slice(Count(1), 0, 3), New("a"))) 121 | } 122 | 123 | func TestStarmap(t *testing.T) { 124 | multiMapper := func (is ...interface{}) interface{} { 125 | var s int = 1 126 | for _, i := range is { 127 | s *= i.(int) 128 | } 129 | return s 130 | } 131 | testIterEq(t, New(10, 20, 30), Starmap(multiMapper, Zip(New(1,2,3), Repeat(10, 3)))) 132 | } 133 | 134 | func TestReduce(t *testing.T) { 135 | summer := func (memo interface{}, el interface{}) interface{} { 136 | return memo.(float64) + el.(float64) 137 | } 138 | if float64(.82) - Reduce(Float64(.1,.2,.3,.22), summer, float64(0)).(float64) > .000001 { 139 | t.Error("Sum Reduce failed") 140 | } 141 | } 142 | 143 | func TestTee2(t *testing.T) { 144 | it1, it2 := Tee2(New(5,4,3,2,1)) 145 | for i := range it1 { 146 | j := <- it2 147 | if i != j { 148 | t.Error("Tees are not coming off equal") 149 | } 150 | } 151 | 152 | it1, it2 = Tee2(New(1,2,3,4,5,6)) 153 | testIterEq(t, New(1,2,3,4,5,6), it1) 154 | testIterEq(t, New(1,2,3,4,5,6), it2) 155 | } 156 | 157 | func TestTee(t *testing.T) { 158 | its := Tee(New(3,4,5), 3) 159 | if len(its) != 3 { 160 | t.Error("its length wrong") 161 | } 162 | for _, it := range its { 163 | testIter(t, New(3,4,5), it) 164 | } 165 | } 166 | 167 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 6 | 7 | # Itertools for golang 8 | 9 | [![Build Status](https://travis-ci.org/yanatan16/itertools.png?branch=master)](http://travis-ci.org/yanatan16/itertools) 10 | 11 | This package is a translation of the python `itertools` module. It includes all the usual suspects except for the cartesian product and permutation operators. All iterators are `chan interface{}` which allows some type ambiguity for these generic functions. It would be completely ok, however, to reproduce these functions in your package for your type-specific iterators such as `chan MyStruct`. I did this mostly as a thought exercise on converting python generators to Go. 12 | 13 | Full documentation is available on [godoc](http://godoc.org/github.com/yanatan16/itertools). 14 | 15 | # Implemented Functions 16 | 17 | ## Infinite Iterator Creators 18 | 19 | - `Count(i)` - Infinite count from i 20 | - `Cycle(iter)` - Infinite cycling of `iter` (requires memory) 21 | - `Repeat(element [, n])` - Repeat element n times (or infinitely) 22 | 23 | ## Finite Iterator Creators 24 | 25 | - `New(elements ...)` - Create from `interface{}` elements 26 | - `Int32(elements ...)` - Create from `int32` elements 27 | - `Int64(elements ...)` - Create from `int64` elements 28 | - `Uint(elements ...)` - Create from `uint` elements 29 | - `Uint32(elements ...)` - Create from `uint32` elements 30 | - `Uint64(elements ...)` - Create from `uint64` elements 31 | - `Float32(elements ...)` - Create from `float32` elements 32 | - `Float64(elements ...)` - Create from `float64` elements 33 | 34 | ## Iterator Destroyers 35 | 36 | - `Reduce(iter, reducer, memo)` - Reduce (or Foldl) across the iterator. 37 | - `List(iter)` - Create a list from the iterator 38 | 39 | ## Iterator Modifiers 40 | 41 | - `Chain(iters...)` - Chain together multiple iterators 42 | - `DropWhile(predicate, iter)` - Drop elements until `predicate(el) == false` 43 | - `TakeWhile(predicate, iter)` - Take elements until `predicate(el) == false` 44 | - `Filter(predicate, iter)` - Filter out elements when `predicate(el) == false` 45 | - `FilterFalse(predicate, iter)` - Filter out elements when `predicate(el) == true` 46 | - `Slice(iter, start[, stop[, step]])` - Drop elements until the start (0-based index). Stop upon stop (exclusive) unless not given. Step is 1 unless given. 47 | 48 | ## More Iterator Modifiers 49 | 50 | - `Map(mapper func(interface{}) interface{}, iter)` - Map each element to `mapper(el)` 51 | - `MultiMap(multiMapper func(interface{}...)interface{}, iters...)` - Map all the iterators as variadic arguments to `multiMaper(elements...)`. Stop on shortest iterator. 52 | - `MultiMapLongest(multiMapper func(interface{}...)interface{}, iters...)` - Same as MultiMap except stop on longest iterator. Shorter iterators are filled with `nil` after they are exhausted. 53 | - `Starmap(multiMapper func(interface{}...)interface{}, iter)` - If iter is an iterator of `[]interface{}`, then expand it into the `multiMapper`. 54 | - `Zip(iters...)` - Zip multiple iterators together 55 | - `ZipLongest(iters...)` - Zip multiple iterators together. Take the longest. Shorter ones are appended with `nil`. 56 | - `Tee(iter, n)` - Split an iterator into n equal versions. 57 | - `Tee2(iter)` - Split an iterator into two equal versions 58 | 59 | 60 | ## Benchmarks 61 | 62 | ``` 63 | BenchmarkFilter 200000 41188 ns/op 64 | BenchmarkNoFilter 5000000 640 ns/op 65 | BenchmarkMap 10000 108577 ns/op 66 | BenchmarkNoMap 10000000 321 ns/op 67 | BenchmarkReduce 1000000 1962 ns/op 68 | BenchmarkNoReduce 100000000 22.0 ns/op 69 | ``` 70 | 71 | Clearly this package is much slower than doing just the original operations. However, all large systems inevitably involve abstractions and layers of indirection which are convinient while adding overhead. A system designer should think carefully before choosing one over the other. To be clear, this package was written as a thought experiment, and while usable, might represent a pattern to be rewritten using your types in your project to save the cost of typecasting overhead. 72 | 73 | 74 | Thanks to [damienklinnert](/damienklinnert) for the benchmarks. 75 | 76 | # License 77 | 78 | Copyright (c) 2013 Jon Eisen 79 | 80 | Permission is hereby granted, free of charge, to any person obtaining a copy 81 | of this software and associated documentation files (the "Software"), to deal 82 | in the Software without restriction, including without limitation the rights 83 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 84 | copies of the Software, and to permit persons to whom the Software is 85 | furnished to do so, subject to the following conditions: 86 | 87 | The above copyright notice and this permission notice shall be included in 88 | all copies or substantial portions of the Software. 89 | 90 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 91 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 92 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 93 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 94 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 95 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 96 | THE SOFTWARE. 97 | 98 | 99 | -------------------------------------------------------------------------------- /itertools.go: -------------------------------------------------------------------------------- 1 | // Package itertools provides a translation of the python standard library module itertools. 2 | // Many of the functions have been brought over, althought not all. 3 | // In this implementation, chan interface{} has been used as all iterators; if more specific types are necessary, 4 | // feel free to copy the code to your project to be implemented with more specific types. 5 | package itertools 6 | 7 | import ( 8 | "sync" 9 | ) 10 | 11 | type Iter chan interface{} 12 | type Predicate func (interface{}) bool 13 | type Mapper func (interface{}) interface{} 14 | type MultiMapper func (...interface{}) interface{} 15 | type Reducer func (memo interface{}, element interface{}) interface{} 16 | 17 | func New(els ... interface{}) Iter { 18 | c := make(Iter) 19 | go func () { 20 | for _, el := range els { 21 | c <- el 22 | } 23 | close(c) 24 | }() 25 | return c 26 | } 27 | 28 | func Int64(els ... int64) Iter { 29 | c := make(Iter) 30 | go func () { 31 | for _, el := range els { 32 | c <- el 33 | } 34 | close(c) 35 | }() 36 | return c 37 | } 38 | 39 | func Int32(els ... int32) Iter { 40 | c := make(Iter) 41 | go func () { 42 | for _, el := range els { 43 | c <- el 44 | } 45 | close(c) 46 | }() 47 | return c 48 | } 49 | 50 | func Float64(els ... float64) Iter { 51 | c := make(Iter) 52 | go func () { 53 | for _, el := range els { 54 | c <- el 55 | } 56 | close(c) 57 | }() 58 | return c 59 | } 60 | 61 | func Float32(els ... float32) Iter { 62 | c := make(Iter) 63 | go func () { 64 | for _, el := range els { 65 | c <- el 66 | } 67 | close(c) 68 | }() 69 | return c 70 | } 71 | 72 | func Uint(els ... uint) Iter { 73 | c := make(Iter) 74 | go func () { 75 | for _, el := range els { 76 | c <- el 77 | } 78 | close(c) 79 | }() 80 | return c 81 | } 82 | func Uint64(els ... uint64) Iter { 83 | c := make(Iter) 84 | go func () { 85 | for _, el := range els { 86 | c <- el 87 | } 88 | close(c) 89 | }() 90 | return c 91 | } 92 | 93 | func Uint32(els ... uint32) Iter { 94 | c := make(Iter) 95 | go func () { 96 | for _, el := range els { 97 | c <- el 98 | } 99 | close(c) 100 | }() 101 | return c 102 | } 103 | 104 | func List(it Iter) []interface{} { 105 | arr := make([]interface{}, 0, 1) 106 | for el := range it { 107 | arr = append(arr, el) 108 | } 109 | return arr 110 | } 111 | 112 | // Count from i to infinity 113 | func Count(i int) Iter { 114 | c := make(Iter) 115 | go func () { 116 | for ; true; i++ { 117 | c <- i 118 | } 119 | }() 120 | return c 121 | } 122 | 123 | // Cycle through an iterator infinitely (requires memory) 124 | func Cycle(it Iter) Iter { 125 | c, a := make(Iter), make([]interface{}, 0, 1) 126 | go func () { 127 | for el := range it { 128 | a = append(a, el) 129 | c <- el 130 | } 131 | for { 132 | for _, el := range a { 133 | c <- el 134 | } 135 | } 136 | }() 137 | return c 138 | } 139 | 140 | // Repeat an element n times or infinitely 141 | func Repeat(el interface{}, n ...int) Iter { 142 | c := make(Iter) 143 | go func () { 144 | for i := 0; len(n) == 0 || i < n[0]; i++ { 145 | c <- el 146 | } 147 | close(c) 148 | }() 149 | return c 150 | } 151 | 152 | // Chain together multiple iterators 153 | func Chain(its ...Iter) Iter { 154 | c := make(Iter) 155 | go func() { 156 | for _, it := range its { 157 | for el := range it { 158 | c <- el 159 | } 160 | } 161 | close(c) 162 | }() 163 | return c 164 | } 165 | 166 | // Elements after pred(el) == true 167 | func DropWhile(pred Predicate, it Iter) Iter { 168 | c := make(Iter) 169 | go func () { 170 | for el := range it { 171 | if drop := pred(el); !drop { 172 | c <- el 173 | break 174 | } 175 | } 176 | for el := range it { 177 | c <- el 178 | } 179 | close(c) 180 | }() 181 | return c 182 | } 183 | 184 | 185 | // Elements before pred(el) == false 186 | func TakeWhile(pred Predicate, it Iter) Iter { 187 | c := make(Iter) 188 | go func () { 189 | for el := range it { 190 | if take := pred(el); take { 191 | c <- el 192 | } else { 193 | break 194 | } 195 | } 196 | close(c) 197 | }() 198 | return c 199 | } 200 | 201 | // Filter out any elements where pred(el) == false 202 | func Filter(pred Predicate, it Iter) Iter { 203 | c := make(Iter) 204 | go func () { 205 | for el := range it { 206 | if keep := pred(el); keep { 207 | c <- el 208 | } 209 | } 210 | close(c) 211 | }() 212 | return c 213 | } 214 | 215 | // Filter out any elements where pred(el) == true 216 | func FilterFalse(pred Predicate, it Iter) Iter { 217 | c := make(Iter) 218 | go func () { 219 | for el := range it { 220 | if drop := pred(el); !drop { 221 | c <- el 222 | } 223 | } 224 | close(c) 225 | }() 226 | return c 227 | } 228 | 229 | // Sub-iterator from start (inclusive) to [stop (exclusive) every [step (default 1)]] 230 | func Slice(it Iter, startstopstep...int) Iter { 231 | start, stop, step := 0, 0, 1 232 | if len(startstopstep) == 1 { 233 | start = startstopstep[0] 234 | } else if len(startstopstep) == 2 { 235 | start, stop = startstopstep[0], startstopstep[1] 236 | } else if len(startstopstep) >= 3 { 237 | start, stop, step = startstopstep[0], startstopstep[1], startstopstep[2] 238 | } 239 | 240 | c := make(Iter) 241 | go func () { 242 | i := 0 243 | // Start 244 | for el := range it { 245 | if i >= start { 246 | c <- el // inclusive 247 | break 248 | } 249 | i += 1 250 | } 251 | 252 | // Stop 253 | i, j := i + 1, 1 254 | for el := range it { 255 | if stop > 0 && i >= stop { 256 | break 257 | } else if j % step == 0 { 258 | c <- el 259 | } 260 | 261 | i, j = i + 1, j + 1 262 | } 263 | 264 | close(c) 265 | }() 266 | return c 267 | } 268 | 269 | // Map an iterator to fn(el) for el in it 270 | func Map(fn Mapper, it Iter) Iter { 271 | c := make(Iter) 272 | go func () { 273 | for el := range it { 274 | c <- fn(el) 275 | } 276 | close(c) 277 | }() 278 | return c 279 | } 280 | 281 | // Map p, q, ... to fn(pEl, qEl, ...) 282 | // Breaks on first closed channel 283 | func MultiMap(fn MultiMapper, its ...Iter) Iter { 284 | c := make(Iter) 285 | go func() { 286 | Outer: 287 | for { 288 | els := make([]interface{}, len(its)) 289 | for i, it := range its { 290 | if el, ok := <- it; ok { 291 | els[i] = el 292 | } else { 293 | break Outer 294 | } 295 | } 296 | c <- fn(els...) 297 | } 298 | close(c) 299 | }() 300 | return c 301 | } 302 | 303 | // Map p, q, ... to fn(pEl, qEl, ...) 304 | // Breaks on last closed channel 305 | func MultiMapLongest(fn MultiMapper, its ...Iter) Iter { 306 | c := make(Iter) 307 | go func() { 308 | for { 309 | els := make([]interface{}, len(its)) 310 | n := 0 311 | for i, it := range its { 312 | if el, ok := <- it; ok { 313 | els[i] = el 314 | } else { 315 | n += 1 316 | } 317 | } 318 | if n < len(its) { 319 | c <- fn(els...) 320 | } else { 321 | break 322 | } 323 | } 324 | close(c) 325 | }() 326 | return c 327 | } 328 | 329 | // Map an iterator if arrays to a fn(els...) 330 | // Iter must be an iterator of []interface{} (possibly created by Zip) 331 | // If not, Starmap will act like MultiMap with a single iterator 332 | func Starmap(fn MultiMapper, it Iter) Iter { 333 | c := make(Iter) 334 | go func() { 335 | for els := range it { 336 | if elements, ok := els.([]interface{}); ok { 337 | c <- fn(elements...) 338 | } else { 339 | c <- fn(els) 340 | } 341 | } 342 | close(c) 343 | }() 344 | return c 345 | } 346 | 347 | // Zip up multiple interators into one 348 | // Close on shortest iterator 349 | func Zip(its ...Iter) Iter { 350 | c := make(Iter) 351 | go func() { 352 | defer close(c) 353 | for { 354 | els := make([]interface{}, len(its)) 355 | for i, it := range its { 356 | if el, ok := <- it; ok { 357 | els[i] = el 358 | } else { 359 | return 360 | } 361 | } 362 | c <- els 363 | } 364 | }() 365 | return c 366 | } 367 | 368 | // Zip up multiple iterators into one 369 | // Close on longest iterator 370 | func ZipLongest(its ...Iter) Iter { 371 | c := make(Iter) 372 | go func() { 373 | for { 374 | els := make([]interface{}, len(its)) 375 | n := 0 376 | for i, it := range its { 377 | if el, ok := <- it; ok { 378 | els[i] = el 379 | } else { 380 | n += 1 381 | } 382 | } 383 | if n < len(its) { 384 | c <- els 385 | } else { 386 | break 387 | } 388 | } 389 | close(c) 390 | }() 391 | return c 392 | } 393 | 394 | // Reduce the iterator (aka fold) from the left 395 | func Reduce(it Iter, red Reducer, memo interface{}) interface{} { 396 | for el := range it { 397 | memo = red(memo, el) 398 | } 399 | return memo 400 | } 401 | 402 | // Split an iterator into n multiple iterators 403 | // Requires memory to keep values for n iterators 404 | func Tee(it Iter, n int) []Iter { 405 | deques := make([][]interface{}, n) 406 | iters := make([]Iter, n) 407 | for i := 0; i < n; i++ { 408 | iters[i] = make(Iter) 409 | } 410 | 411 | mutex := new(sync.Mutex) 412 | 413 | gen := func(myiter Iter, i int) { 414 | for { 415 | if len(deques[i]) == 0 { 416 | mutex.Lock() 417 | if len(deques[i]) == 0 { 418 | if newval, ok := <- it; ok { 419 | for i, d := range deques { 420 | deques[i] = append(d, newval) 421 | } 422 | } else { 423 | mutex.Unlock() 424 | close(myiter) 425 | break 426 | } 427 | } 428 | mutex.Unlock() 429 | } 430 | var popped interface{} 431 | popped, deques[i] = deques[i][0], deques[i][1:] 432 | myiter <- popped 433 | } 434 | } 435 | for i, iter := range iters { 436 | go gen(iter, i) 437 | } 438 | return iters 439 | } 440 | 441 | // Helper to tee just into two iterators 442 | func Tee2(it Iter) (Iter, Iter) { 443 | iters := Tee(it, 2) 444 | return iters[0], iters[1] 445 | } 446 | --------------------------------------------------------------------------------