| 511 | } |
| 512 | |
| 513 | Status Conditional::ExtractBodies(Graph* graph) { |
| 514 | VLOG(2) << "Extracting bodies for " << name(); |
| 515 | for (auto b : {BranchType::kElseBranch, BranchType::kThenBranch}) { |
| 516 | bodies_[static_cast<int>(b)] = |
| 517 | absl::make_unique<Graph>(graph->op_registry()); |
| 518 | } |
| 519 | |
| 520 | auto find_branch = [&](const Edge* e) { |
| 521 | const auto& id = state_map_->LookupCondId(e->src()); |
| 522 | return IsSwitch(e->src()) ? BranchType(e->src_output()) |
| 523 | : state_map_->FindBranchOf(id, predicate_); |
| 524 | }; |
| 525 | |
| 526 | std::array<std::vector<Node*>, 2> stacks; |
| 527 | VLOG(5) << "Merges: " << NodesToString(merges_); |
| 528 | for (Node* m : merges_) { |
| 529 | VLOG(5) << "For merge: " << m->DebugString() << " " |
| 530 | << state_map_->CondStateToString(m); |
| 531 | for (auto e : m->in_edges()) { |
| 532 | if (e->IsControlEdge()) continue; |
| 533 | BranchType branch = find_branch(e); |
| 534 | TF_RET_CHECK(branch == BranchType::kThenBranch || |
| 535 | branch == BranchType::kElseBranch) |
| 536 | << "Error: " << e->src()->name() |
| 537 | << " is not on either then or else branch (" << Branch_Name(branch) |
| 538 | << ") for predicate " << DebugString(predicate_) << " [" |
| 539 | << DebugString(state_map_->LookupCondId(e->src())) << "]."; |
| 540 | Node* src = e->src(); |
| 541 | if (IsSwitch(src)) { |
| 542 | // Switch node outputs and dependencies are handled separately. |
| 543 | TF_RETURN_IF_ERROR(AddSwitch(src)); |
| 544 | } else { |
| 545 | stacks[static_cast<int>(branch)].push_back(src); |
| 546 | } |
| 547 | } |
| 548 | } |
| 549 | |
| 550 | for (auto branch : {BranchType::kElseBranch, BranchType::kThenBranch}) { |
| 551 | int branch_index = static_cast<int>(branch); |
| 552 | auto output = bodies_[branch_index].get(); |
| 553 | auto& stack = stacks[branch_index]; |
| 554 | VLOG(5) << "In branch: " << Branch_Name(branch) << " " |
| 555 | << NodesToString(stack); |
| 556 | std::vector<bool> visited(graph->num_node_ids(), false); |
| 557 | node_maps_[branch_index].resize(graph->num_node_ids(), nullptr); |
| 558 | auto& node_map = node_maps_[branch_index]; |
| 559 | |
| 560 | while (!stack.empty()) { |
| 561 | Node* n = stack.back(); |
| 562 | stack.pop_back(); |
| 563 | |
| 564 | if (visited.at(n->id())) continue; |
| 565 | visited[n->id()] = true; |
| 566 | |
| 567 | // Verify output edges and record control edges exitting scope. |
| 568 | for (const Edge* e : n->out_edges()) { |
| 569 | Node* dst = e->dst(); |
| 570 | if (IsMerge(dst)) continue; |
nothing calls this directly
no test coverage detected