split-iterator.go

package util

import (
	"container/list"
)

// Iterator A lazy, unbounded collection of elements.
type Iterator interface {
	// Consume one element from the iterator. This may block for an indefinite
	// amount of time, until an element becomes available.
	Next() SomeType
}

// Split creates two new Iterators, one that returns the elements that satisfy
// the given predicate, another one for those elements that don't.
//
// Boundary conditions:
// - It may happen that there is no demand on one of the target iterators for a
//   short period of time.
// - It is ok to drop elements from the parent iterator if their target iterator
//   did not collect them in a reasonable amount of time.
func Split(iterator Iterator, predicate func(SomeType) bool) (trueIterator, falseIterator Iterator) {
	sIt := &SplitIterator{
		source:     &iterator,
		predicate:  predicate,
		trueQueue:  list.New(),
		falseQueue: list.New(),
	}
	trueIterator = PredicatedIterator{
		desiredPredicateResult: true,
		sIt:                    sIt,
	}
	falseIterator = PredicatedIterator{
		desiredPredicateResult: false,
		sIt:                    sIt,
	}
	return
}

type SplitIterator struct {
	source     *Iterator
	trueQueue  *list.List
	falseQueue *list.List
	predicate  func(SomeType) bool
}

func (si SplitIterator) getNext(desiredPredicateResult bool) SomeType {
	queue := si.falseQueue
	if desiredPredicateResult {
		queue = si.trueQueue
	}

	if queue.Len() > 0 {
		e := queue.Front()
		val := e.Value.(SomeType)
		queue.Remove(e)
		return val
	}

	val := (*si.source).Next()
	state := si.predicate(val)
	if desiredPredicateResult == state {
		return val
	}

	if state {
		si.trueQueue.PushBack(val)
	} else {
		si.falseQueue.PushBack(val)
	}
	return si.getNext(desiredPredicateResult)
}

type PredicatedIterator struct {
	sIt                    *SplitIterator
	desiredPredicateResult bool
}

func (it PredicatedIterator) Next() SomeType {
	return it.sIt.getNext(it.desiredPredicateResult)
}

// Don't worry about this.
type SomeType struct{}