JungのグラフをトポロジカルソートするStream

処理の実行順序を制御するための依存関係を書きたかったが、手軽なものがぱっと見つからなかったので継続の練習がてら書いてみた。
パフォーマンスは知らない。

トポロジカルソートの概要はこちら(Wikipedia)。
グラフはJungのもの。
有向edgeのheadにあるvertexよりもtailにあるvertexが先に出現することが保証される。
深さ優先探索の閉路検出版をCPSっぽくIteratorとして実装。
閉路がある場合、それを検出した時点でIllegalStateException。
未探索ノードから自分に向いているエッジが無いノードを見つけたら、そのノードと残りの探索手続きを保存している。

// usage
final Graph<V, E> graph = getDirectedGraph(); // 何らかの有向グラフ
TopologicalSort.stream(graph) // Stream<V>
    .forEach(...); // あとはお好きに

コード

import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import org.apache.commons.lang3.tuple.Pair;

import edu.uci.ics.jung.graph.Graph;

/**
 * 有向グラフのトポロジカルソート
 */
public class TopologicalSort {
	private TopologicalSort(){}
	/**
     * Iterator実装
	 * @param <T> Streamのノード型
	 */
	public static class TopologicalSortedIterator<T> implements Iterator<T> {
		private final Graph<T, ?> graph;

		private final Map<T, State> state;
		private enum State {
			NONE, TEMP, PERP
		}

		private Supplier<Pair<Optional<T>, Supplier<?>>> cont;

		/**
		 * @param graph 対象の有向グラフ: 閉路があると走査中にIllegalStateException
		 */
		public TopologicalSortedIterator(final Graph<T, ?> graph) {
			this.graph = graph;
			this.state = graph.getVertices().stream()
				.collect(Collectors.toMap(v -> v, v -> State.NONE));

			// 全nodeを操作する継続を初期設定
			final Collection<T> vertices = graph.getVertices();
			final Iterator<T> it = vertices.iterator();

			this.cont = () -> iterateVertices(it);

			// 最初の要素を準備、もしくは空なら終端へ
			prepareNext();
		}

		final Pair<Optional<T>, Supplier<?>> iterateVertices(final Iterator<T> it) {
			if(!it.hasNext()) {
				return Pair.of(Optional.empty(), () -> { throw new NoSuchElementException(); });
			} else {
				final T next = it.next();
				return visit(next, () -> iterateVertices(it));
			}
		};

		private Pair<Optional<T>, Supplier<?>> visit(final T node, Supplier<Pair<Optional<T>, Supplier<?>>> cont) {
			switch(state.get(node)) {
			case TEMP:
				// 一度訪れた箇所が確定する前に再訪
				throw new IllegalStateException("閉路検出");
			case PERP:
				// 既訪なのでこのノードは無視して継続する
				return cont.get();
			case NONE:
				// 未訪ノード
				state.put(node, State.TEMP);

				final Collection<T> inVertices = getInVertices(graph, node);
				final Iterator<T> it = inVertices.iterator();
				return visitRest(node, it, cont);
			default: throw new IllegalStateException();
			}
		}

		private static <V, E> Collection<V> getInVertices(Graph<V, E> graph, V node) {
			return graph.getInEdges(node).stream()
				.filter(graph::containsEdge)
				.map(graph::getSource)
				.collect(Collectors.toList());
		}

		private Pair<Optional<T>, Supplier<?>> visitRest(final T node, final Iterator<? extends T> it, Supplier<Pair<Optional<T>, Supplier<?>>> cont) {
			if(!it.hasNext()) { // 子が無い
				// 訪れる先が無いのでreturnが自分
				state.put(node, State.PERP);
				return Pair.of(Optional.of(node), cont);
			} else { // inEdgeが存在
				// headを訪れる
				final T head = it.next();
				// cont: tailに対してheadを訪れる->cont
				return visit(head, () -> visitRest(node, it, cont));
			}
		}

		/** 次の値が存在する場合present */
		private Optional<T> next = Optional.empty();

		/**
		 * 次の値と継続を取得しそれぞれ保持する
		 */
		@SuppressWarnings("unchecked")
		private void prepareNext() {
			final Pair<Optional<T>, Supplier<?>> p = cont.get();

			this.next = p.getLeft();
			this.cont = (Supplier<Pair<Optional<T>, Supplier<?>>>) p.getRight();
		}

		/*
		 * (非 Javadoc)
		 * @see java.util.Iterator#hasNext()
		 */
		@Override
		public boolean hasNext() {
			return next.isPresent();
		}

		/*
		 * (非 Javadoc)
		 * @see java.util.Iterator#next()
		 */
		@Override
		public T next() {
			try {
				return next.orElseThrow(() -> new NoSuchElementException());
			} finally {
				next.ifPresent(x -> prepareNext());
			}
		}
	}

	/**
	 * 対象のグラフをトポロジカルソートした結果をノードのStreamとして返す
	 * @param graph
	 * @return Stream
	 */
	public static <V, E> Stream<V> stream(final Graph<V, E> graph) {
		return StreamSupport.stream(
			Spliterators.spliterator(
				new TopologicalSortedIterator<V>(graph),
				graph.getVertexCount(),
				Spliterator.NONNULL),
			false);
	}
}