001    /*
002     * Copyright 2010-2015 JetBrains s.r.o.
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package org.jetbrains.kotlin.utils;
018    
019    import org.jetbrains.annotations.NotNull;
020    
021    import java.util.*;
022    
023    public class DFS {
024        public static <N, R> R dfs(@NotNull Collection<N> nodes, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited, @NotNull NodeHandler<N, R> handler) {
025            for (N node : nodes) {
026                doDfs(node, neighbors, visited, handler);
027            }
028            return handler.result();
029        }
030    
031        public static <N, R> R dfs(
032            @NotNull Collection<N> nodes,
033            @NotNull Neighbors<N> neighbors,
034            @NotNull NodeHandler<N, R> handler
035        ) {
036            return dfs(nodes, neighbors, new VisitedWithSet<N>(), handler);
037        }
038    
039        public static <N, R> R dfsFromNode(@NotNull N node, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited, @NotNull NodeHandler<N, R> handler) {
040            doDfs(node, neighbors, visited, handler);
041            return handler.result();
042        }
043    
044        public static <N> void dfsFromNode(
045                @NotNull N node,
046                @NotNull Neighbors<N> neighbors,
047                @NotNull Visited<N> visited
048        ) {
049            dfsFromNode(node, neighbors, visited, new AbstractNodeHandler<N, Void>() {
050                @Override
051                public Void result() {
052                    return null;
053                }
054            });
055        }
056    
057        public static <N> List<N> topologicalOrder(@NotNull Iterable<N> nodes, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited) {
058            TopologicalOrder<N> handler = new TopologicalOrder<N>();
059            for (N node : nodes) {
060                doDfs(node, neighbors, visited, handler);
061            }
062            return handler.result();
063        }
064    
065        public static <N> List<N> topologicalOrder(@NotNull Iterable<N> nodes, @NotNull Neighbors<N> neighbors) {
066            return topologicalOrder(nodes, neighbors, new VisitedWithSet<N>());
067        }
068    
069        private static <N> void doDfs(@NotNull N current, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited, @NotNull NodeHandler<N, ?> handler) {
070            if (!visited.checkAndMarkVisited(current)) return;
071            if (!handler.beforeChildren(current)) return;
072    
073            for (N neighbor : neighbors.getNeighbors(current)) {
074                doDfs(neighbor, neighbors, visited, handler);
075            }
076            handler.afterChildren(current);
077        }
078    
079        public interface NodeHandler<N, R> {
080            boolean beforeChildren(N current);
081    
082            void afterChildren(N current);
083    
084            R result();
085        }
086    
087        public interface Neighbors<N> {
088            @NotNull
089            Iterable<? extends N> getNeighbors(N current);
090        }
091    
092        public interface Visited<N> {
093            boolean checkAndMarkVisited(N current);
094        }
095    
096        public static abstract class AbstractNodeHandler<N, R> implements NodeHandler<N, R> {
097            @Override
098            public boolean beforeChildren(N current) {
099                return true;
100            }
101    
102            @Override
103            public void afterChildren(N current) {
104            }
105        }
106    
107        public static class VisitedWithSet<N> implements Visited<N> {
108            private final Set<N> visited;
109    
110            public VisitedWithSet() {
111                this(new HashSet<N>());
112            }
113    
114            public VisitedWithSet(@NotNull Set<N> visited) {
115                this.visited = visited;
116            }
117    
118            @Override
119            public boolean checkAndMarkVisited(N current) {
120                return visited.add(current);
121            }
122        }
123    
124        public static abstract class CollectingNodeHandler<N, R, C extends Iterable<R>> extends AbstractNodeHandler<N, C> {
125            @NotNull
126            protected final C result;
127    
128            protected CollectingNodeHandler(@NotNull C result) {
129                this.result = result;
130            }
131    
132            @Override
133            @NotNull
134            public C result() {
135                return result;
136            }
137        }
138    
139        public static abstract class NodeHandlerWithListResult<N, R> extends CollectingNodeHandler<N, R, LinkedList<R>> {
140            protected NodeHandlerWithListResult() {
141                super(new LinkedList<R>());
142            }
143        }
144    
145        public static class TopologicalOrder<N> extends NodeHandlerWithListResult<N, N> {
146            @Override
147            public void afterChildren(N current) {
148                result.addFirst(current);
149            }
150        }
151    }