無限Streamに終了条件を設定する

例えばDOMのようなツリー構造があったとして、あるノードを渡された時に、ルートからそのノードまでの名前を連結してパス文字列を作成する、というプログラムを考えます。

Java7でも動くように、副作用ループを使った書き方だとこんな感じ。

public static String getTreePath(Node node, String delimiter) {
    LinkedList<String> nodePath = new LinkedList<>();
    for (Node n = node; n != null; n = n.getParent()) { //根でgetParentするとnullが帰るものとする
        nodePath.add(0, n.getName());
    }

    return delimiter + String.join(delimiter, nodePath);
}

んで、これをJava8のStream APIで書くとどう書けるのか考えてみて、最初に書いたのがこれ。

public static String getTreePath2(Node node, String delimiter) {
    List<String> names = Stream.iterate(node, Node::getParent)
            .filter(n -> n != null)
            .map(o -> o.getName())
            .collect(Collectors.toList());
    Collections.reverse(names);

    return delimiter + String.join(delimiter, names);
}

半分くらい書いたところで、あれ？って思ったんですが、案の定これは NullPointerException で落ちました。Stream.iterate に渡した関数(Node::getParemt)を連鎖的に評価させてその結果をリスト化するわけですが、getParent が null を返しても、そこで評価を止めることができないわけですね。

API Document とかいろいろ漁って調べてみましたが、単純に Stream.iterate や Stream.generate で生成した無限 Stream は、limitで決まった数で区切るか、zip で有限 Stream と結合するか、短絡評価をする終端操作で区切るしかなさそうです。「指定した条件を満たした最初の要素より前のものだけ抽出」のような中間操作ないし終端操作はない模様。

こういう場合は、 Stream の作成のしかたからカスタマイズする必要があって、API Documentによると、「最も単純だが低パフォーマンスな」作成方法は、Iterator の実装クラスを作って、そのインスタンスから Splitterator を作ることとありました。なるほど、Iterator なら生成関数と終了条件関数をセットで実装できますね。

で書いてみたのがこちら。

public static String getTreePath3(Node node, String delimiter) {

    Iterator<Node> iter = new Iterator<Node>() {
        private Node n = node;

        @Override
        public boolean hasNext() {
            return (n != null);
        }

        @Override
        public Node next() {
            Node pre = n;
            n = pre != null ? pre.getParent() : null;
            return pre;
        }
    };


    List<String> names = StreamSupport.stream(
            Spliterators.spliteratorUnknownSize(iter, Spliterator.ORDERED), false)
            .map(o -> o.getName())
            .collect(Collectors.toList());
    Collections.reverse(names);

    return delimiter + String.join(delimiter, names);
}

長ぇ。

最初のループを使った書き方とくらべても断然めんどくさいですね。

Stream を作るところを、別のメソッドに切り出して、軽く汎用的にしてみます。引数としては、 Stream.iterate の引数に終了条件 term を加えた形ですね。

private static <T> Stream<T> iterateLimited(T seed, UnaryOperator<T> generator, Predicate<T> term) {
    Iterator<T> iter = new Iterator<T>() {
        T current = seed;

        @Override
        public boolean hasNext() {
            return !term.test(current);
        }

        @Override
        public T next() {
            T pre = current;
            current = generator.apply(pre);
            return pre;
        }
    };
    return StreamSupport.stream(
            Spliterators.spliteratorUnknownSize(iter, Spliterator.ORDERED), false);
}

これを使うと、メインの処理はこんなふうに書けます。(ついでに Optional 使って null を隠蔽してみました)

public static String getTreePath4(Node node, String delimiter) {
    List<String> names = iterateLimited(
            Optional.ofNullable(node), o -> o.map(n -> n.getParent()), o -> !o.isPresent())
            .filter(o -> o.isPresent())
            .map(o -> o.get().getName())
            .collect(Collectors.toList());
    Collections.reverse(names);

    return delimiter + String.join(delimiter, names);
}

だいぶ「らしく」なったような気がします。
… iterateLimited みたいな関数は、標準で用意してもらいたいものですが…

(2017/12/06追記)
Java9では簡単に書けるようになりました。
noisyspot.hatenablog.com