Question

Writing iterators for custom collections in Java is quite complicated, because instead of writing straight-forward code that provides one element after the other, you essentially have to write a state machine:

public class CustomCollection<T> implements Iterable<T>
{
    private T[] data;
    private int size;

    @Override
    public Iterator<T> iterator()
    {
        return new Iterator<T>()
        {
            private int cursor = 0;

            @Override
            public boolean hasNext()
            {
                return cursor < size;
            }

            @Override
            public T next()
            {
                return data[cursor++];
            }

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

For collections more complicated than an array list or a linked list, getting these state machines correctly is a daunting task. In fact, the C# design team deemed writing custom iterators complicated enough to introduce special language support (yield return) for letting the compiler build the state machines.

Is something like yield return coming in the next version of Java? Or are there any library solutions that make my life easier when it comes to writing my own iterators in Java?

Answer 1

Java has always provided a mechanism for maintaining state and continuing execution at a later point in time: threads. The basic idea for my library solution is to let a ConcurrentIterable produce the elements in one thread, and let a ConcurrentIterator consume them in another, communicating via a bounded queue. (This is generally known as the producer/consumer pattern.)

First, here is a demonstration of the simplified usage:

public class CustomCollection<T> extends ConcurrentIterable<T>
{
    private T[] data;
    private int size;

    @Override
    protected void provideElements()
    {
        for (int i = 0; i < size; ++i)
        {
            provideElement(data[i]);
        }
    }
    // ...
}

Note the complete absence of state machines. All you have to do is derive from ConcurrentIterable and implement the method provideElements. Inside this method, you write straight-forward code which calls provideElement for each element in the collection.

Sometimes a client does not iterate through the entire collection, for example in a linear search. You can stop providing elements as soon as an abortion is detected by checking iterationAborted():

    @Override
    protected void provideElements()
    {
        for (int i = 0; i < size && !iterationAborted(); ++i)
        {
            provideElement(data[i]);
        }
    }

It is perfectly fine not to check iterationAborted(), as long as you do not care about the additional elements being generated. With infinite sequences, checking iterationAborted() is mandatory.

How can the producer detect that the consumer has stopped iterating? This is implemented by having a strong reference to a token in the consumer and a weak reference to that same token in the producer. When the consumer stops iterating, the token becomes eligible for garbage collection, and it will eventually become invisible to the producer. From then on, all new elements will simply be discarded.

(Without this precaution, under certain circumstances the bounded queue could eventually fill up, the producer would enter an infinite loop, and the contained elements would never be garbage collected.)

And now for the implementation details:

ConcurrentIterable.java

import java.util.Iterator;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;

public abstract class ConcurrentIterable<T> implements Iterable<T>
{
    private static final int CAP = 1000;
    private final ThreadLocal<CommunicationChannel<T>> channels
    = new ThreadLocal<CommunicationChannel<T>>();

    @Override
    public Iterator<T> iterator()
    {
        BlockingQueue<Option<T>> queue = new ArrayBlockingQueue<Option<T>>(CAP);
        Object token = new Object();
        final CommunicationChannel<T> channel
        = new CommunicationChannel<T>(queue, token);
        new Thread(new Runnable()
        {
            @Override
            public void run()
            {
                channels.set(channel);
                provideElements();
                enqueueSentinel();
            }
        }).start();
        return new ConcurrentIterator<T>(queue, token);
    }

    protected abstract void provideElements();

    protected final boolean iterationAborted()
    {
        return channels.get().iterationAborted();
    }

    protected final void provideElement(T element)
    {
        enqueue(Option.some(element));
    }

    private void enqueueSentinel()
    {
        enqueue(Option.<T> none());
    }

    private void enqueue(Option<T> element)
    {
        try
        {
            while (!offer(element))
            {
                System.gc();
            }
        }
        catch (InterruptedException ignore)
        {
            ignore.printStackTrace();
        }
    }

    private boolean offer(Option<T> element) throws InterruptedException
    {
        CommunicationChannel<T> channel = channels.get();
        return channel.iterationAborted()
            || channel.queue.offer(element, 1, TimeUnit.SECONDS);
    }
}

CommunicationChannel.java

import java.lang.ref.WeakReference;
import java.util.concurrent.BlockingQueue;

public class CommunicationChannel<T>
{
    public final BlockingQueue<Option<T>> queue;
    private final WeakReference<Object> token;

    public CommunicationChannel(BlockingQueue<Option<T>> queue, Object token)
    {
        this.queue = queue;
        this.token = new WeakReference<Object>(token);
    }

    public boolean iterationAborted()
    {
        return token.get() == null;
    }
}

ConcurrentIterator.java

import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.concurrent.BlockingQueue;

public class ConcurrentIterator<T> implements Iterator<T>
{
    private final BlockingQueue<Option<T>> queue;
    @SuppressWarnings("unused")
    private final Object token;
    private Option<T> next;

    public ConcurrentIterator(BlockingQueue<Option<T>> queue, Object token)
    {
        this.queue = queue;
        this.token = token;
    }

    @Override
    public boolean hasNext()
    {
        if (next == null)
        {
            try
            {
                next = queue.take();
            }
            catch (InterruptedException ignore)
            {
                ignore.printStackTrace();
            }
        }
        return next.present;
    }

    @Override
    public T next()
    {
        if (!hasNext()) throw new NoSuchElementException();
        T result = next.value;
        next = null;
        return result;
    }

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

Option.java

public class Option<T>
{
    public final T value;
    public final boolean present;

    private Option(T value, boolean present)
    {
        this.value = value;
        this.present = present;
    }

    public static <T> Option<T> some(T value)
    {
        return new Option<T>(value, true);
    }

    @SuppressWarnings("unchecked")
    public static <T> Option<T> none()
    {
        return none;
    }

    @SuppressWarnings({ "rawtypes", "unchecked" })
    private static final Option none = new Option(null, false);
}

Let me know what you think!