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
017package org.jetbrains.jet.utils;
018
019import com.google.common.collect.Lists;
020import com.google.common.collect.Sets;
021import org.jetbrains.annotations.NotNull;
022
023import java.util.Collection;
024import java.util.LinkedList;
025import java.util.List;
026import java.util.Set;
027
028public 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}