001    /*
002     * Copyright 2010-2013 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.jet.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)) {
071                return;
072            }
073            handler.beforeChildren(current);
074            for (N neighbor : neighbors.getNeighbors(current)) {
075                doDfs(neighbor, neighbors, visited, handler);
076            }
077            handler.afterChildren(current);
078        }
079    
080        public interface NodeHandler<N, R> {
081    
082            void beforeChildren(N current);
083    
084            void afterChildren(N current);
085    
086            R result();
087        }
088    
089        public interface Neighbors<N> {
090            @NotNull
091            Iterable<N> getNeighbors(N current);
092        }
093    
094        public interface Visited<N> {
095            boolean checkAndMarkVisited(N current);
096        }
097    
098        public static abstract class AbstractNodeHandler<N, R> implements NodeHandler<N, R> {
099            @Override
100            public void beforeChildren(N current) {
101            }
102    
103            @Override
104            public void afterChildren(N current) {
105            }
106        }
107    
108        public static class VisitedWithSet<N> implements Visited<N> {
109            private final Set<N> visited;
110    
111            public VisitedWithSet() {
112                this(new HashSet<N>());
113            }
114    
115            public VisitedWithSet(@NotNull Set<N> visited) {
116                this.visited = visited;
117            }
118    
119            @Override
120            public boolean checkAndMarkVisited(N current) {
121                return visited.add(current);
122            }
123        }
124    
125        public static abstract class NodeHandlerWithListResult<N, R> extends AbstractNodeHandler<N, List<R>> {
126            protected final LinkedList<R> result = new LinkedList<R>();
127    
128            @Override
129            public List<R> result() {
130                return result;
131            }
132        }
133    
134        public static class TopologicalOrder<N> extends NodeHandlerWithListResult<N, N> {
135            @Override
136            public void afterChildren(N current) {
137                result.addFirst(current);
138            }
139        }
140    }