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 com.google.common.collect.Lists;
020    import com.google.common.collect.Sets;
021    import org.jetbrains.annotations.NotNull;
022    
023    import java.util.Collection;
024    import java.util.LinkedList;
025    import java.util.List;
026    import java.util.Set;
027    
028    public class DFS {
029        public static <N, R> R dfs(@NotNull Collection<N> nodes, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited, @NotNull NodeHandler<N, R> handler) {
030            for (N node : nodes) {
031                doDfs(node, neighbors, visited, handler);
032            }
033            return handler.result();
034        }
035    
036        public static <N, R> R dfs(
037            @NotNull Collection<N> nodes,
038            @NotNull Neighbors<N> neighbors,
039            @NotNull NodeHandler<N, R> handler
040        ) {
041            return dfs(nodes, neighbors, new VisitedWithSet<N>(), handler);
042        }
043    
044        public static <N, R> R dfsFromNode(@NotNull N node, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited, @NotNull NodeHandler<N, R> handler) {
045            doDfs(node, neighbors, visited, handler);
046            return handler.result();
047        }
048    
049        public static <N> void dfsFromNode(
050                @NotNull N node,
051                @NotNull Neighbors<N> neighbors,
052                @NotNull Visited<N> visited
053        ) {
054            dfsFromNode(node, neighbors, visited, new AbstractNodeHandler<N, Void>() {
055                @Override
056                public Void result() {
057                    return null;
058                }
059            });
060        }
061    
062        public static <N> List<N> topologicalOrder(@NotNull Iterable<N> nodes, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited) {
063            TopologicalOrder<N> handler = new TopologicalOrder<N>();
064            for (N node : nodes) {
065                doDfs(node, neighbors, visited, handler);
066            }
067            return handler.result();
068        }
069    
070        public static <N> List<N> topologicalOrder(@NotNull Iterable<N> nodes, @NotNull Neighbors<N> neighbors) {
071            return topologicalOrder(nodes, neighbors, new VisitedWithSet<N>());
072        }
073    
074        private static <N> void doDfs(@NotNull N current, @NotNull Neighbors<N> neighbors, @NotNull Visited<N> visited, @NotNull NodeHandler<N, ?> handler) {
075            if (!visited.checkAndMarkVisited(current)) {
076                return;
077            }
078            handler.beforeChildren(current);
079            for (N neighbor : neighbors.getNeighbors(current)) {
080                doDfs(neighbor, neighbors, visited, handler);
081            }
082            handler.afterChildren(current);
083        }
084    
085        public interface NodeHandler<N, R> {
086    
087            void beforeChildren(N current);
088    
089            void afterChildren(N current);
090    
091            R result();
092        }
093    
094        public interface Neighbors<N> {
095            @NotNull
096            Iterable<N> getNeighbors(N current);
097        }
098    
099        public interface Visited<N> {
100            boolean checkAndMarkVisited(N current);
101        }
102    
103        public static abstract class AbstractNodeHandler<N, R> implements NodeHandler<N, R> {
104            @Override
105            public void beforeChildren(N current) {
106            }
107    
108            @Override
109            public void afterChildren(N current) {
110            }
111        }
112    
113        public static class VisitedWithSet<N> implements Visited<N> {
114            private final Set<N> visited;
115    
116            public VisitedWithSet() {
117                this(Sets.<N>newHashSet());
118            }
119    
120            public VisitedWithSet(@NotNull Set<N> visited) {
121                this.visited = visited;
122            }
123    
124            @Override
125            public boolean checkAndMarkVisited(N current) {
126                return visited.add(current);
127            }
128        }
129    
130        public static abstract class NodeHandlerWithListResult<N, R> extends AbstractNodeHandler<N, List<R>> {
131            protected final LinkedList<R> result = Lists.newLinkedList();
132    
133            @Override
134            public List<R> result() {
135                return result;
136            }
137        }
138    
139        public static class TopologicalOrder<N> extends NodeHandlerWithListResult<N, N> {
140            @Override
141            public void afterChildren(N current) {
142                result.addFirst(current);
143            }
144        }
145    }