| 1407 | } // namespace |
| 1408 | |
| 1409 | Status MutableGraphView::SortTopologically( |
| 1410 | bool ignore_cycles, |
| 1411 | absl::Span<const TopologicalDependency> extra_dependencies) { |
| 1412 | if (!mutation_.updated_nodes_.empty() || !mutation_.new_nodes_.empty()) { |
| 1413 | // Cannot sort when there is an active mutation due to indices possibly |
| 1414 | // being changed or invalidated. |
| 1415 | return errors::InvalidArgument(kMutableGraphViewSortTopologicallyError, |
| 1416 | "active mutation exists."); |
| 1417 | } |
| 1418 | |
| 1419 | const int num_nodes = nodes_.size(); |
| 1420 | |
| 1421 | // Group extra dependencies by `from` node. |
| 1422 | absl::flat_hash_map<int, std::vector<int>> extra_dependencies_by_parent; |
| 1423 | for (const auto& extra_dependency : extra_dependencies) { |
| 1424 | if (extra_dependency.graph_view_ != this || |
| 1425 | extra_dependency.from_ == extra_dependency.to_ || |
| 1426 | extra_dependency.from_ < 0 || extra_dependency.from_ >= num_nodes || |
| 1427 | extra_dependency.to_ < 0 || extra_dependency.to_ >= num_nodes) { |
| 1428 | return errors::InvalidArgument(kMutableGraphViewSortTopologicallyError, |
| 1429 | "invalid extra dependencies."); |
| 1430 | } |
| 1431 | extra_dependencies_by_parent[extra_dependency.from_].push_back( |
| 1432 | extra_dependency.to_); |
| 1433 | } |
| 1434 | |
| 1435 | // Reversed colored post-order DFS traversal. This does not fail on cycles, |
| 1436 | // but there are no guarantees on ordering within a cycle. |
| 1437 | std::vector<TraversalState> traversal_state(num_nodes, PENDING); |
| 1438 | int curr_pos = num_nodes - 1; |
| 1439 | std::vector<int> order(num_nodes); |
| 1440 | std::vector<Edge> edges_in_cycle; |
| 1441 | |
| 1442 | auto push_onto_stack = [this]( |
| 1443 | const int curr_index, const int fanout_index, |
| 1444 | std::vector<RecursionStackEntry>* recursion_stack, |
| 1445 | std::vector<TraversalState>* traversal_state, |
| 1446 | std::vector<Edge>* edges_in_cycle) { |
| 1447 | // Ignore NextIteration -> Merge connections to break control flow cycles. |
| 1448 | if (IsNextIteration(graph_->node(curr_index)) && |
| 1449 | IsMerge(graph_->node(fanout_index))) { |
| 1450 | return; |
| 1451 | } |
| 1452 | auto& fanout_traversal_state = (*traversal_state)[fanout_index]; |
| 1453 | if (fanout_traversal_state == PROCESSING) { |
| 1454 | // Cycle detected. |
| 1455 | edges_in_cycle->push_back({curr_index, fanout_index}); |
| 1456 | } else if (fanout_traversal_state == PENDING) { |
| 1457 | // Unvisited node, simply add to stack for future traversal. |
| 1458 | recursion_stack->push_back({fanout_index, ENTER}); |
| 1459 | } |
| 1460 | }; |
| 1461 | |
| 1462 | auto process_fanouts = [this, &extra_dependencies_by_parent, |
| 1463 | &push_onto_stack]( |
| 1464 | const int curr_index, |
| 1465 | std::vector<RecursionStackEntry>* recursion_stack, |
| 1466 | std::vector<TraversalState>* traversal_state, |