Java で yield return ぽいことを実現しよう！ #Java

Java には yield return のような構文は存在しないのですが、 Is there a Java equivalent to C#'s 'yield' keyword? の回答にあるとおり、頑張ればそれっぽい機能を実現することができます。例えば同 stackoverflow の質問の回答では

の二つのライブラリが具体例として紹介されています。

この二つのライブラリ、前者は Java の文法に yield のキーワードを追加していたりバイトコードいじいじするなどして頑張り過ぎていてちょっとカジュアル利用が難しい一方で、後者はマルチスレッドで yield ライクな機能をクラスとして提供していてカジュアルに利用できるかと思いきや、Maven が提供されていなかったり単一のクラスに収まっているわけでもなく、どちらにせよカジュアルな利用は難しい状況にあります。

そういうわけで、ないものは作ってしまえ精神を発揮して yield return ぽい機能を作ってみました。

Yielder.java

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

/**
 * Java で yield return っぽいことを実現する機能を提供します。
 *
 * 以下は使い方の一例です。
 * <pre>
 * class YielderDemo {
 *   public static void main(String[] args) {
 *     for (Integer integer : Yielder.newIterable(new Yielder<Integer>() {
 *       public void run() {
 *         for (int i = 0; i < 10; i++) {
 *           yield(i);
 *         }
 *         for (int i = 30; i < 50; i++) {
 *           yield(i);
 *         }
 *       }
 *     })) {
 *       System.out.println(integer);
 *     }
 *   }
 * }
 * </pre>
 *
 * @author KOMIYA Atsushi
 */
public abstract class Yielder<T> {
    private static class Item<T> {
        T object;
        boolean ended = false;

        private Item<T> set(T object) {
            this.object = object;
            return this;
        }

        private Item<T> end() {
            this.ended = true;
            return this;
        }
    }

    private static class Context<T> {
        private final BlockingQueue<Item<T>> queue;
        private final Item<T>[] items;
        private int index;
        private boolean endReceived;
        private boolean ended;
        private Exception thrownException;

        Context(int queueSize) {
            queue = new ArrayBlockingQueue<>(queueSize);
            items = new Item[queueSize + 2];
            for (int i = 0; i < items.length; i++) {
                items[i] = new Item<>();
            }
        }

        synchronized boolean yield(T returnValue) {
            if (endReceived) {
                throw new IllegalStateException("#yieldEnd() has been called.");
            }

            try {
                Item<T> item = items[index++].set(returnValue);
                while (!queue.offer(item, 1000, TimeUnit.SECONDS)) {
                    if (ended) {
                        throw new IllegalStateException("iteration has been ended.");
                    }
                }

                if (index == items.length) {
                    index = 0;
                }

                return true;

            } catch (Exception e) {
                thrownException = e;
                return false;
            }
        }

        synchronized boolean yieldEnd() {
            if (endReceived) {
                throw new IllegalStateException("#yieldEnd() has been called.");
            }

            endReceived = true;

            try {
                Item<T> item = items[index++].end();
                while (!queue.offer(item, 1000, TimeUnit.SECONDS)) {
                    if (ended) {
                        throw new IllegalStateException("iteration has been ended.");
                    }
                }

                if (index == items.length) {
                    index = 0;
                }

                return true;

            } catch (InterruptedException e) {
                thrownException = e;
                return false;
            }
        }

        synchronized Exception thrownException() {
            return thrownException;
        }

        Item<T> pop() throws InterruptedException {
            Item<T> result = queue.take();
            if (result.ended) {
                this.ended = true;
            }

            return result;
        }
    }

    public static class YielderIterable<T> implements Iterable<T>, AutoCloseable {
        private final Context<T> context;
        private final Thread thread;

        private boolean iteratorCalled;
        private boolean closed;

        private YielderIterable(int queueSize, Yielder<T> yielder) {
            this.context = new Context<>(queueSize);
            this.thread = new Thread() {
                @Override
                public void run() {
                    try {
                        yielder.context.set(YielderIterable.this.context);
                        yielder.run();
                        context.yieldEnd();

                    } catch (Exception e) {
                        e.printStackTrace();

                    } finally {
                        yielder.context.remove();
                    }
                }
            };

            this.thread.start();
        }

        @Override
        public Iterator<T> iterator() {
            if (iteratorCalled) {
                throw new IllegalStateException("#iterator() has been called");
            }

            iteratorCalled = true;

            return new Iterator<T>() {
                Item<T> nextItem;

                @Override
                public boolean hasNext() {
                    if (nextItem != null) {
                        return !nextItem.ended;
                    }

                    try {
                        nextItem = context.pop();

                    } catch (InterruptedException e) {
                        throw new IllegalStateException(e);
                    }

                    return !nextItem.ended;
                }

                @Override
                public T next() {
                    if (!hasNext()) {
                        throw new NoSuchElementException();
                    }

                    Item<T> result = nextItem;
                    nextItem = null;

                    return result.object;
                }
            };
        }

        @Override
        public void close() throws Exception {
            if (closed) {
                return;
            }

            closed = true;
            context.endReceived = true;
            context.ended = true;

            thread.interrupt();
        }
    }

    private static final int DEFAULT_QUEUE_SIZE = 1;
    private ThreadLocal<Context<T>> context = new ThreadLocal<>();

    public static <T> Iterable<T> newIterable(final Yielder<T> yielder) {
        return new YielderIterable<>(DEFAULT_QUEUE_SIZE, yielder);
    }

    protected boolean yield(T returnValue) {
        return context.get().yield(returnValue);
    }

    protected Exception thrownException() {
        return context.get().thrownException();
    }

    public abstract void run();
}

class YielderDemo {
    public static void main(String[] args) {
        for (Integer integer : Yielder.newIterable(new Yielder<Integer>() {
            @Override
            public void run() {
                for (int i = 0; i < 10; i++) {
                    yield(i);
                }

                for (int i = 30; i < 50; i++) {
                    yield(i);
                }
            }
        })) {
            System.out.println(integer);
        }
    }
}

あんまりテストできていないので正しく動かないこともあるかもしれないけど、そこは自己責任で！