| 43 | } // namespace |
| 44 | |
| 45 | void DfsTraversal(const GraphTopologyView& graph_view, |
| 46 | const absl::Span<const NodeDef* const> from, |
| 47 | const TraversalDirection direction, |
| 48 | const DfsPredicates& predicates, |
| 49 | const DfsCallbacks& callbacks) { |
| 50 | std::vector<DfsStackElem> stack; |
| 51 | stack.reserve(from.size()); |
| 52 | |
| 53 | for (const NodeDef* node : from) { |
| 54 | const absl::optional<int> node_idx = graph_view.GetNodeIndex(*node); |
| 55 | DCHECK(node_idx.has_value()) << "Illegal start node: " << node->name(); |
| 56 | if (node_idx.has_value()) { |
| 57 | stack.emplace_back(node_idx.value()); |
| 58 | } |
| 59 | } |
| 60 | |
| 61 | absl::flat_hash_map<int, NodeState> node_state; |
| 62 | while (!stack.empty()) { |
| 63 | DfsStackElem w = stack.back(); |
| 64 | stack.pop_back(); |
| 65 | |
| 66 | NodeState& state = node_state[w.node]; |
| 67 | if (state == NodeState::kDone) continue; |
| 68 | |
| 69 | // Skip nodes that we should not enter. |
| 70 | if (predicates.enter && !predicates.enter(graph_view.GetNode(w.node))) { |
| 71 | state = NodeState::kDone; |
| 72 | continue; |
| 73 | } |
| 74 | |
| 75 | // We've processed all the children of this node. |
| 76 | if (w.children_visited) { |
| 77 | state = NodeState::kDone; |
| 78 | if (callbacks.post_order) { |
| 79 | callbacks.post_order(graph_view.GetNode(w.node)); |
| 80 | } |
| 81 | continue; |
| 82 | } |
| 83 | |
| 84 | // Loop detected. |
| 85 | if (state == NodeState::kVisiting) { |
| 86 | if (callbacks.on_back_edge) { |
| 87 | callbacks.on_back_edge(graph_view.GetNode(w.src), |
| 88 | graph_view.GetNode(w.node)); |
| 89 | } |
| 90 | continue; |
| 91 | } |
| 92 | |
| 93 | state = NodeState::kVisiting; |
| 94 | if (callbacks.pre_order) { |
| 95 | callbacks.pre_order(graph_view.GetNode(w.node)); |
| 96 | } |
| 97 | |
| 98 | // Enqueue the node again with the children_visited flag set to true. |
| 99 | stack.emplace_back(w.node, true, w.src); |
| 100 | |
| 101 | // Check if we can continue traversal from the current node. |
| 102 | if (predicates.advance && !predicates.advance(graph_view.GetNode(w.node))) { |