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