generative flat map, groovy

generativeFlatMap.groovy

List.metaClass.foldn = { init, closure ->
	def l = delegate.clone()
	def acc = init
	def it = l.listIterator()
	while ( it.hasNext() ) {
		def e = it.next()
		it.remove()
		closure( acc, e, it )
	}
	return acc
}

def it1 = { i, n ->
	def c,
		a = [i],
		q = a as Queue
	while ( (c = q.poll()) != null ) {
		def m = n( c )
		a += m
		q += m
	}
	return a
}

def it2 = { i, n ->
	return [i].foldn( [i] ) { a, v, iter ->
		n( v ).each {
			iter.add it
			a << it
		}
		return a
	}
}

class C {
	public static <T> Set<T> collectWhileNested( T init, Closure<Collection<T>> fn ) {
		T curr
		Set<T> acc = [init] as Set<T>
		Queue<T> queue = acc as ArrayDeque<T>

		while ( (curr = queue.poll()) != null ) {
			def add = fn( curr )
			if ( add ) {
				acc += add
				queue += add
			}
		}

		return acc
	}
}

def l = [1,2,3,4]
def i = 1
def c = { it == 1 ? [2,3,4] : [] }

def lwo = { it == 1 ? [2, 3] : [] }
def lwt = { it == 2 ? [4, 5] : lwo( it ) }
def lop = { it == 2 ? [3, 4] : lwo( it ) }

/*
println it1( i, c )
println it2( i, c )

def r = l.foldn([]) { acc, v, it ->
	acc << v
}
*/
println C.collectWhileNested( 1, lwo )
println C.collectWhileNested( 1, lwt )
println C.collectWhileNested( 1, lop )