SelectManyIterator.java

package Blah;

import java.util.Iterator;
import java.util.NoSuchElementException;


public class SelectManyIterable<TSource, TResult> implements Iterable<TResult>
{
    public interface TransformProvider<TSource, TResult>
    {
        Iterator<TResult> selectInnerIterator(TSource sourceObject);
    }

    private final Iterable<TSource> source;
    private final TransformProvider<TSource, TResult> transformProvider;

    public SelectManyIterable(Iterable<TSource> source, TransformProvider<TSource, TResult> selector)
    {
        this.source = source;
        this.transformProvider = selector;
    }

    @Override
    public Iterator<TResult> iterator()
    {
        return new SelectManyIterator<TSource, TResult>(source.iterator(), transformProvider);
    }

    private static class SelectManyIterator<TSource, TResult> implements Iterator<TResult>
    {
        private final Iterator<TSource> iteratorSource;
        private final TransformProvider<TSource, TResult> selector;
        private Iterator<TResult> currentIterator = null;

        public SelectManyIterator(Iterator<TSource> source, TransformProvider<TSource, TResult> selector)
        {
            this.iteratorSource = source;
            this.selector = selector;
            loadNextIterator();
        }

        @Override
        public boolean hasNext()
        {
            if(currentIteratorHasMore())
            {
                return true;
            }
            else if(currentIteratorIsEmptyButSourceHasMore())
            {
                loadNextIterator();
                return hasNext();
            }
            else
            {
                return false;
            }
        }

        @Override
        public TResult next()
        {
            if(currentIteratorHasMore())
            {
                return currentIterator.next();
            }
            else if(currentIteratorIsEmptyButSourceHasMore())
            {
                loadNextIterator();
                return next();
            }
            else
            {
                throw new NoSuchElementException();
            }
        }

        @Override
        public void remove()
        {
            throw new UnsupportedOperationException();
        }

        private boolean currentIteratorHasMore()
        {
            return this.currentIterator != null &&
                   this.currentIterator.hasNext();
        }

        private boolean currentIteratorIsEmptyButSourceHasMore()
        {
            return currentIterator != null &&
                    !currentIterator.hasNext() &&
                    iteratorSource.hasNext();
        }

        private void loadNextIterator()
        {
            if(this.iteratorSource.hasNext())
            {
                this.currentIterator = selector.selectInnerIterator(this.iteratorSource.next());
            }
            else
            {
                this.currentIterator = null;
            }
        }
    }
}