Make Graph::UpdateEdge() be O(e) instead of O(E)

where:
- E = number of edges in the graph
- e = number of edges on the node of interest

e is necessarily <= E and is typically really small
(# of inputs to an operation + control edges)

PiperOrigin-RevId: 210624296

2 files changed, 9 insertions, 15 deletions

diff --git a/tensorflow/core/graph/graph.cc b/tensorflow/core/graph/graph.cc
index ade9266231..1630ab7a15 100644
--- a/tensorflow/core/graph/graph.cc
+++ b/tensorflow/core/graph/graph.cc
@@ -495,6 +495,15 @@ void Graph::RemoveControlEdge(const Edge* e) {
   RemoveEdge(e);
 }
 
+namespace {
+const Edge* FindEdge(const Node* dst, int index) {
+  for (const Edge* e : dst->in_edges()) {
+    if (e->dst_input() == index) return e;
+  }
+  return nullptr;
+}
+}  // namespace
+
 Status Graph::UpdateEdge(Node* new_src, int new_src_index, Node* dst,
                          int dst_index) {
   TF_RETURN_IF_ERROR(IsValidOutputTensor(new_src, new_src_index));
@@ -512,17 +521,6 @@ Status Graph::UpdateEdge(Node* new_src, int new_src_index, Node* dst,
   return Status::OK();
 }
 
-const Edge* Graph::FindEdge(const Node* dst, int index) {
-  for (const Edge* e : edges_) {
-    // edges_ will contain null edges if RemoveEdge() was called.
-    if (e == nullptr) continue;
-    if (e->dst() == dst && e->dst_input() == index) {
-      return e;
-    }
-  }
-  return nullptr;
-}
-
 Status Graph::AddFunctionLibrary(const FunctionDefLibrary& fdef_lib) {
   // Need a new-enough consumer to support the functions we add to the graph.
   if (fdef_lib.function_size() > 0 && versions_->min_consumer() < 12) {
diff --git a/tensorflow/core/graph/graph.h b/tensorflow/core/graph/graph.h
index a147c94689..52e9f23a76 100644
--- a/tensorflow/core/graph/graph.h
+++ b/tensorflow/core/graph/graph.h
@@ -680,10 +680,6 @@ class Graph {
   // AddWhileContext() or Node::while_ctx(), but this manages the lifetime.
   std::map<string, WhileContext> while_ctxs_;
 
-  // Searches through edges_ for the Edge whose destination node and index
-  // matches dst. An edge with destination `dst` must exist in the graph.
-  const Edge* FindEdge(const Node* dst, int index);
-
   TF_DISALLOW_COPY_AND_ASSIGN(Graph);
 };