001package ball.util;
002/*-
003 * ##########################################################################
004 * Utilities
005 * %%
006 * Copyright (C) 2008 - 2022 Allen D. Ball
007 * %%
008 * Licensed under the Apache License, Version 2.0 (the "License");
009 * you may not use this file except in compliance with the License.
010 * You may obtain a copy of the License at
011 *
012 *      http://www.apache.org/licenses/LICENSE-2.0
013 *
014 * Unless required by applicable law or agreed to in writing, software
015 * distributed under the License is distributed on an "AS IS" BASIS,
016 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
017 * See the License for the specific language governing permissions and
018 * limitations under the License.
019 * ##########################################################################
020 */
021import java.util.Arrays;
022import java.util.Collection;
023import java.util.Iterator;
024import java.util.Objects;
025import java.util.Spliterator;
026import java.util.Spliterators.AbstractSpliterator;
027import java.util.function.Consumer;
028import java.util.function.Function;
029import java.util.function.Supplier;
030import java.util.stream.Stream;
031import java.util.stream.StreamSupport;
032import lombok.ToString;
033
034/**
035 * {@link Spliterator} implemenation to build {@link Stream}s to walk a tree
036 * of {@code <T>} nodes.  See {@link #walk(Object,Function)}.
037 *
038 * @param       <T>             The type of node.
039 *
040 * @author {@link.uri mailto:ball@hcf.dev Allen D. Ball}
041 */
042@ToString
043public class Walker<T> extends AbstractSpliterator<T> {
044    private final Stream<Supplier<Spliterator<T>>> stream;
045    private Iterator<Supplier<Spliterator<T>>> iterator = null;
046    private Spliterator<? extends T> spliterator = null;
047
048    private Walker(Collection<? extends T> nodes, Function<? super T,Collection<? extends T>> childrenOf) {
049        super(Long.MAX_VALUE, IMMUTABLE | NONNULL);
050
051        stream =
052            nodes.stream()
053            .map(t -> (() -> new Walker<T>(t, childrenOf)));
054    }
055
056    private Walker(T node, Function<? super T,Collection<? extends T>> childrenOf) {
057        super(Long.MAX_VALUE, IMMUTABLE | NONNULL);
058
059        stream =
060            Stream.of(node)
061            .filter(Objects::nonNull)
062            .map(childrenOf)
063            .filter(Objects::nonNull)
064            .flatMap(Collection::stream)
065            .map(t -> (() -> new Walker<T>(t, childrenOf)));
066        spliterator = Stream.of(node).spliterator();
067    }
068
069    @Override
070    public Spliterator<T> trySplit() {
071        if (iterator == null) {
072            iterator = stream.iterator();
073        }
074
075        return iterator.hasNext() ? iterator.next().get() : null;
076    }
077
078    @Override
079    public boolean tryAdvance(Consumer<? super T> consumer) {
080        boolean accepted = false;
081
082        while (! accepted) {
083            if (spliterator == null) {
084                spliterator = trySplit();
085            }
086
087            if (spliterator != null) {
088                accepted = spliterator.tryAdvance(consumer);
089
090                if (! accepted) {
091                    spliterator = null;
092                }
093            } else {
094                break;
095            }
096        }
097
098        return accepted;
099    }
100
101    /**
102     * Entry-point to create a {@link Stream} for traversing a tree of type
103     * {@code <T>} nodes.  The caller need only supply the root node and a
104     * {@link Function} to calculate the children nodes.
105     *
106     * @param   root            The root node.
107     * @param   childrenOf      The {@link Function} to get the children of
108     *                          a node.
109     * @param   <T>             The type of node.
110     *
111     * @return  A {@link Stream}.
112     */
113    public static <T> Stream<T> walk(T root, Function<? super T,Collection<? extends T>> childrenOf) {
114        return StreamSupport.stream(new Walker<>(root, childrenOf), false);
115    }
116
117    /**
118     * Entry-point to create a {@link Stream} for traversing a tree of type
119     * {@code <T>} nodes.  The caller need only supply the root nodes and a
120     * {@link Function} to calculate the children nodes.
121     *
122     * @param   roots           The root nodes.
123     * @param   childrenOf      The {@link Function} to get the children of
124     *                          a node.
125     * @param   <T>             The type of node.
126     *
127     * @return  A {@link Stream}.
128     */
129    public static <T> Stream<T> walk(Collection<? extends T> roots, Function<? super T,Collection<? extends T>> childrenOf) {
130        return StreamSupport.stream(new Walker<>(roots, childrenOf), false);
131    }
132}