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