Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 107 additions & 41 deletions itertools.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,14 @@ import (
)

type Iter chan interface{}
type Predicate func (interface{}) bool
type Mapper func (interface{}) interface{}
type MultiMapper func (...interface{}) interface{}
type Reducer func (memo interface{}, element interface{}) interface{}
type Predicate func(interface{}) bool
type Mapper func(interface{}) interface{}
type MultiMapper func(...interface{}) interface{}
type Reducer func(memo interface{}, element interface{}) interface{}

func New(els ... interface{}) Iter {
func New(els ...interface{}) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -25,9 +25,9 @@ func New(els ... interface{}) Iter {
return c
}

func Int64(els ... int64) Iter {
func Int64(els ...int64) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -36,9 +36,9 @@ func Int64(els ... int64) Iter {
return c
}

func Int32(els ... int32) Iter {
func Int32(els ...int32) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -47,9 +47,9 @@ func Int32(els ... int32) Iter {
return c
}

func Float64(els ... float64) Iter {
func Float64(els ...float64) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -58,9 +58,9 @@ func Float64(els ... float64) Iter {
return c
}

func Float32(els ... float32) Iter {
func Float32(els ...float32) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -69,19 +69,19 @@ func Float32(els ... float32) Iter {
return c
}

func Uint(els ... uint) Iter {
func Uint(els ...uint) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
close(c)
}()
return c
}
func Uint64(els ... uint64) Iter {
func Uint64(els ...uint64) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -90,9 +90,9 @@ func Uint64(els ... uint64) Iter {
return c
}

func Uint32(els ... uint32) Iter {
func Uint32(els ...uint32) Iter {
c := make(Iter)
go func () {
go func() {
for _, el := range els {
c <- el
}
Expand All @@ -112,7 +112,7 @@ func List(it Iter) []interface{} {
// Count from i to infinity
func Count(i int) Iter {
c := make(Iter)
go func () {
go func() {
for ; true; i++ {
c <- i
}
Expand All @@ -123,7 +123,7 @@ func Count(i int) Iter {
// Cycle through an iterator infinitely (requires memory)
func Cycle(it Iter) Iter {
c, a := make(Iter), make([]interface{}, 0, 1)
go func () {
go func() {
for el := range it {
a = append(a, el)
c <- el
Expand All @@ -140,7 +140,7 @@ func Cycle(it Iter) Iter {
// Repeat an element n times or infinitely
func Repeat(el interface{}, n ...int) Iter {
c := make(Iter)
go func () {
go func() {
for i := 0; len(n) == 0 || i < n[0]; i++ {
c <- el
}
Expand All @@ -166,7 +166,7 @@ func Chain(its ...Iter) Iter {
// Elements after pred(el) == true
func DropWhile(pred Predicate, it Iter) Iter {
c := make(Iter)
go func () {
go func() {
for el := range it {
if drop := pred(el); !drop {
c <- el
Expand All @@ -181,11 +181,10 @@ func DropWhile(pred Predicate, it Iter) Iter {
return c
}


// Elements before pred(el) == false
func TakeWhile(pred Predicate, it Iter) Iter {
c := make(Iter)
go func () {
go func() {
for el := range it {
if take := pred(el); take {
c <- el
Expand All @@ -201,7 +200,7 @@ func TakeWhile(pred Predicate, it Iter) Iter {
// Filter out any elements where pred(el) == false
func Filter(pred Predicate, it Iter) Iter {
c := make(Iter)
go func () {
go func() {
for el := range it {
if keep := pred(el); keep {
c <- el
Expand All @@ -215,7 +214,7 @@ func Filter(pred Predicate, it Iter) Iter {
// Filter out any elements where pred(el) == true
func FilterFalse(pred Predicate, it Iter) Iter {
c := make(Iter)
go func () {
go func() {
for el := range it {
if drop := pred(el); !drop {
c <- el
Expand All @@ -227,18 +226,18 @@ func FilterFalse(pred Predicate, it Iter) Iter {
}

// Sub-iterator from start (inclusive) to [stop (exclusive) every [step (default 1)]]
func Slice(it Iter, startstopstep...int) Iter {
func Slice(it Iter, startstopstep ...int) Iter {
start, stop, step := 0, 0, 1
if len(startstopstep) == 1 {
start = startstopstep[0]
} else if len(startstopstep) == 2 {
start, stop = startstopstep[0], startstopstep[1]
} else if len(startstopstep) >= 3 {
} else if len(startstopstep) >= 3 {
start, stop, step = startstopstep[0], startstopstep[1], startstopstep[2]
}

c := make(Iter)
go func () {
go func() {
i := 0
// Start
for el := range it {
Expand All @@ -250,15 +249,15 @@ func Slice(it Iter, startstopstep...int) Iter {
}

// Stop
i, j := i + 1, 1
i, j := i+1, 1
for el := range it {
if stop > 0 && i >= stop {
break
} else if j % step == 0 {
} else if j%step == 0 {
c <- el
}

i, j = i + 1, j + 1
i, j = i+1, j+1
}

close(c)
Expand All @@ -269,7 +268,7 @@ func Slice(it Iter, startstopstep...int) Iter {
// Map an iterator to fn(el) for el in it
func Map(fn Mapper, it Iter) Iter {
c := make(Iter)
go func () {
go func() {
for el := range it {
c <- fn(el)
}
Expand All @@ -283,11 +282,11 @@ func Map(fn Mapper, it Iter) Iter {
func MultiMap(fn MultiMapper, its ...Iter) Iter {
c := make(Iter)
go func() {
Outer:
Outer:
for {
els := make([]interface{}, len(its))
for i, it := range its {
if el, ok := <- it; ok {
if el, ok := <-it; ok {
els[i] = el
} else {
break Outer
Expand All @@ -309,7 +308,7 @@ func MultiMapLongest(fn MultiMapper, its ...Iter) Iter {
els := make([]interface{}, len(its))
n := 0
for i, it := range its {
if el, ok := <- it; ok {
if el, ok := <-it; ok {
els[i] = el
} else {
n += 1
Expand Down Expand Up @@ -353,7 +352,7 @@ func Zip(its ...Iter) Iter {
for {
els := make([]interface{}, len(its))
for i, it := range its {
if el, ok := <- it; ok {
if el, ok := <-it; ok {
els[i] = el
} else {
return
Expand All @@ -374,7 +373,7 @@ func ZipLongest(its ...Iter) Iter {
els := make([]interface{}, len(its))
n := 0
for i, it := range its {
if el, ok := <- it; ok {
if el, ok := <-it; ok {
els[i] = el
} else {
n += 1
Expand Down Expand Up @@ -415,7 +414,7 @@ func Tee(it Iter, n int) []Iter {
if len(deques[i]) == 0 {
mutex.Lock()
if len(deques[i]) == 0 {
if newval, ok := <- it; ok {
if newval, ok := <-it; ok {
for i, d := range deques {
deques[i] = append(d, newval)
}
Expand Down Expand Up @@ -443,3 +442,70 @@ func Tee2(it Iter) (Iter, Iter) {
iters := Tee(it, 2)
return iters[0], iters[1]
}

// Permutations returns iterator of all permutations of a finite iterator
func Permutations(it Iter, r ...int) chan []interface{} {
c := make(chan []interface{})
go func() {
pool := List(it)
n := len(pool)
if len(r) == 0 {
r = append(r, n)
}
indices := []int{}
for i := 0; i < n; i++ {
indices = append(indices, i)
}
cycles := []int{}
for i := n; i >= n-r[0]+1; i-- {
cycles = append(cycles, i)
}

{
tempSlice := []interface{}{}
for _, i := range indices[:r[0]] {
tempSlice = append(tempSlice, pool[i])
}
c <- tempSlice
}

if n == 0 {
close(c)
return
}

for {
var broken bool
for i := r[0] - 1; i >= 0; i-- {
broken = false
cycles[i]--
if cycles[i] == 0 {
tempIndex := indices[i]
for t := i; t < len(indices)-1; t++ {
indices[t] = indices[t+1]
}
indices[len(indices)-1] = tempIndex
cycles[i] = n - i
} else {
j := cycles[i]
indices[i], indices[len(indices)-j] = indices[len(indices)-j], indices[i]
{
tempSlice := []interface{}{}
for _, t := range indices[:r[0]] {
tempSlice = append(tempSlice, pool[t])
}
c <- tempSlice
}
broken = true
break
}

}
if !broken {
close(c)
return
}
}
}()
return c
}
9 changes: 9 additions & 0 deletions itertools_bench_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,12 @@ func BenchmarkNoReduce(b *testing.B) {
}
}
}

func BenchmarkPermutations(b *testing.B) {
for i := 0; i < b.N; i++ {
input := Int32(1, 2, 3, 4, 5, 6)
for _ = range Permutations(input, 3) {

}
}
}
Loading