Source file src/pkg/exp/iterable/iterable.go

// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The iterable package provides several traversal and searching methods.
// It can be used on anything that satisfies the Iterable interface,
// including vector, though certain functions, such as Map, can also be used on
// something that would produce an infinite amount of data.
package iterable

import (
    "container/list"
    "container/vector"
)

type Iterable interface {
    // Iter should return a fresh channel each time it is called.
    Iter() <-chan interface{}
}

func not(f func(interface{}) bool) func(interface{}) bool {
    return func(e interface{}) bool { return !f(e) }
}

// All tests whether f is true for every element of iter.
func All(iter Iterable, f func(interface{}) bool) bool {
    for e := range iter.Iter() {
        if !f(e) {
            return false
        }
    }
    return true
}

// Any tests whether f is true for at least one element of iter.
func Any(iter Iterable, f func(interface{}) bool) bool {
    return !All(iter, not(f))
}

// Data returns a slice containing the elements of iter.
func Data(iter Iterable) []interface{} {
    var v vector.Vector
    for e := range iter.Iter() {
        v.Push(e)
    }
    return v
}

// filteredIterable is a struct that implements Iterable with each element
// passed through a filter.
type filteredIterable struct {
    it Iterable
    f  func(interface{}) bool
}

func (f *filteredIterable) iterate(out chan<- interface{}) {
    for e := range f.it.Iter() {
        if f.f(e) {
            out <- e
        }
    }
    close(out)
}

func (f *filteredIterable) Iter() <-chan interface{} {
    ch := make(chan interface{})
    go f.iterate(ch)
    return ch
}

// Filter returns an Iterable that returns the elements of iter that satisfy f.
func Filter(iter Iterable, f func(interface{}) bool) Iterable {
    return &filteredIterable{iter, f}
}

// Find returns the first element of iter that satisfies f.
// Returns nil if no such element is found.
func Find(iter Iterable, f func(interface{}) bool) interface{} {
    for e := range Filter(iter, f).Iter() {
        return e
    }
    return nil
}

// Injector is a type representing a function that takes two arguments,
// an accumulated value and an element, and returns the next accumulated value.
// See the Inject function.
type Injector func(interface{}, interface{}) interface{}

// Inject combines the elements of iter by repeatedly calling f with an
// accumulated value and each element in order. The starting accumulated value
// is initial, and after each call the accumulated value is set to the return
// value of f. For instance, to compute a sum:
//   var arr IntArray = []int{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
//   sum := iterable.Inject(arr, 0,
//                          func(ax interface {}, x interface {}) interface {} {
//                            return ax.(int) + x.(int) }).(int)
func Inject(iter Iterable, initial interface{}, f Injector) interface{} {
    acc := initial
    for e := range iter.Iter() {
        acc = f(acc, e)
    }
    return acc
}

// mappedIterable is a helper struct that implements Iterable, returned by Map.
type mappedIterable struct {
    it Iterable
    f  func(interface{}) interface{}
}

func (m *mappedIterable) iterate(out chan<- interface{}) {
    for e := range m.it.Iter() {
        out <- m.f(e)
    }
    close(out)
}

func (m *mappedIterable) Iter() <-chan interface{} {
    ch := make(chan interface{})
    go m.iterate(ch)
    return ch
}

// Map returns an Iterable that returns the result of applying f to each
// element of iter.
func Map(iter Iterable, f func(interface{}) interface{}) Iterable {
    return &mappedIterable{iter, f}
}

// Partition(iter, f) returns Filter(iter, f) and Filter(iter, !f).
func Partition(iter Iterable, f func(interface{}) bool) (Iterable, Iterable) {
    return Filter(iter, f), Filter(iter, not(f))
}

// A Func is a function that, when called, sends the
// iterable values on a channel.
type Func func(chan<- interface{})

// Iter creates and returns a new channel; it starts a
// goroutine running f to send values to the channel.
func (f Func) Iter() <-chan interface{} {
    ch := make(chan interface{})
    go f(ch)
    return ch
}

// Take returns an Iterable that contains the first n elements of iter.
func Take(iter Iterable, n int) Iterable { return Slice(iter, 0, n) }

// TakeWhile returns an Iterable that contains elements from iter while f is true.
func TakeWhile(iter Iterable, f func(interface{}) bool) Iterable {
    return Func(func(ch chan<- interface{}) {
        for v := range iter.Iter() {
            if !f(v) {
                break
            }
            ch <- v
        }
        close(ch)
    })
}

// Drop returns an Iterable that returns each element of iter after the first n elements.
func Drop(iter Iterable, n int) Iterable {
    return Func(func(ch chan<- interface{}) {
        m := n
        for v := range iter.Iter() {
            if m > 0 {
                m--
                continue
            }
            ch <- v
        }
        close(ch)
    })
}

// DropWhile returns an Iterable that returns each element of iter after the initial sequence for which f returns true.
func DropWhile(iter Iterable, f func(interface{}) bool) Iterable {
    return Func(func(ch chan<- interface{}) {
        drop := true
        for v := range iter.Iter() {
            if drop {
                if f(v) {
                    continue
                }
                drop = false
            }
            ch <- v
        }
        close(ch)
    })
}

// Cycle repeats the values of iter in order infinitely.
func Cycle(iter Iterable) Iterable {
    return Func(func(ch chan<- interface{}) {
        for {
            for v := range iter.Iter() {
                ch <- v
            }
        }
    })
}

// Chain returns an Iterable that concatentates all values from the specified Iterables.
func Chain(args []Iterable) Iterable {
    return Func(func(ch chan<- interface{}) {
        for _, e := range args {
            for v := range e.Iter() {
                ch <- v
            }
        }
        close(ch)
    })
}

// Zip returns an Iterable of []interface{} consisting of the next element from
// each input Iterable.  The length of the returned Iterable is the minimum of
// the lengths of the input Iterables.
func Zip(args []Iterable) Iterable {
    return Func(func(ch chan<- interface{}) {
        defer close(ch)
        if len(args) == 0 {
            return
        }
        iters := make([]<-chan interface{}, len(args))
        for i := 0; i < len(iters); i++ {
            iters[i] = args[i].Iter()
        }
        for {
            out := make([]interface{}, len(args))
            for i, v := range iters {
                out[i] = <-v
                if closed(v) {
                    return
                }
            }
            ch <- out
        }
    })
}

// ZipWith returns an Iterable containing the result of executing f using arguments read from a and b.
func ZipWith2(f func(c, d interface{}) interface{}, a, b Iterable) Iterable {
    return Map(Zip([]Iterable{a, b}), func(a1 interface{}) interface{} {
        arr := a1.([]interface{})
        return f(arr[0], arr[1])
    })
}

// ZipWith returns an Iterable containing the result of executing f using arguments read from a, b and c.
func ZipWith3(f func(d, e, f interface{}) interface{}, a, b, c Iterable) Iterable {
    return Map(Zip([]Iterable{a, b, c}), func(a1 interface{}) interface{} {
        arr := a1.([]interface{})
        return f(arr[0], arr[1], arr[2])
    })
}

// Slice returns an Iterable that contains the elements from iter
// with indexes in [start, stop).
func Slice(iter Iterable, start, stop int) Iterable {
    return Func(func(ch chan<- interface{}) {
        defer close(ch)
        i := 0
        for v := range iter.Iter() {
            switch {
            case i >= stop:
                return
            case i >= start:
                ch <- v
            }
            i++
        }
    })
}

// Repeat generates an infinite stream of v.
func Repeat(v interface{}) Iterable {
    return Func(func(ch chan<- interface{}) {
        for {
            ch <- v
        }
    })
}

// RepeatTimes generates a stream of n copies of v.
func RepeatTimes(v interface{}, n int) Iterable {
    return Func(func(ch chan<- interface{}) {
        for i := 0; i < n; i++ {
            ch <- v
        }
        close(ch)
    })
}

// Group is the type for elements returned by the GroupBy function.
type Group struct {
    Key  interface{} // key value for matching items
    Vals Iterable    // Iterable for receiving values in the group
}

// Key defines the interface required by the GroupBy function.
type Grouper interface {
    // Return the key for the given value
    Key(interface{}) interface{}

    // Compute equality for the given keys
    Equal(a, b interface{}) bool
}

// GroupBy combines sequences of logically identical values from iter using k
// to generate a key to compare values.  Each value emitted by the returned
// Iterable is of type Group, which contains the key used for matching the
// values for the group, and an Iterable for retrieving all the values in the
// group.
func GroupBy(iter Iterable, k Grouper) Iterable {
    return Func(func(ch chan<- interface{}) {
        var curkey interface{}
        var lst *list.List
        // Basic strategy is to read one group at a time into a list prior to emitting the Group value
        for v := range iter.Iter() {
            kv := k.Key(v)
            if lst == nil || !k.Equal(curkey, kv) {
                if lst != nil {
                    ch <- Group{curkey, lst}
                }
                lst = list.New()
                curkey = kv
            }
            lst.PushBack(v)
        }
        if lst != nil {
            ch <- Group{curkey, lst}
        }
        close(ch)
    })
}

// Unique removes duplicate values which occur consecutively using id to compute keys.
func Unique(iter Iterable, id Grouper) Iterable {
    return Map(GroupBy(iter, id), func(v interface{}) interface{} { return v.(Group).Key })
}