From 12be7ab4876f77fecfab903df70791623219b3da Mon Sep 17 00:00:00 2001
From: Georgios Pinitas <georgios.pinitas@arm.com>
Date: Tue, 3 Jul 2018 12:06:23 +0100
Subject: COMPMID-1310: Create graph validation executables.

Change-Id: I9e0b57b1b83fe5a95777cdaeddba6ecef650bafc
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/138697
Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
Tested-by: Jenkins <bsgcomp@arm.com>
---
 src/graph/Graph.cpp | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

(limited to 'src/graph/Graph.cpp')

diff --git a/src/graph/Graph.cpp b/src/graph/Graph.cpp
index e1ffeed668..ed24f18943 100644
--- a/src/graph/Graph.cpp
+++ b/src/graph/Graph.cpp
@@ -41,9 +41,9 @@ bool Graph::remove_node(NodeID nid)
 
     std::unique_ptr<INode> &node = _nodes[nid];
 
-    // Remove node connections
     if(node)
     {
+        // Remove node connections
         for(auto &input_eid : node->_input_edges)
         {
             remove_connection(input_eid);
@@ -52,6 +52,10 @@ bool Graph::remove_node(NodeID nid)
         {
             remove_connection(outpud_eid);
         }
+
+        // Remove nid from tagged nodes
+        std::vector<NodeID> &tnodes = _tagged_nodes.at(node->type());
+        tnodes.erase(std::remove(tnodes.begin(), tnodes.end(), nid), tnodes.end());
     }
 
     node = nullptr;
@@ -164,9 +168,9 @@ GraphID Graph::id() const
     return _id;
 }
 
-const std::vector<NodeID> &Graph::inputs()
+const std::vector<NodeID> &Graph::nodes(NodeType type)
 {
-    return _tagged_nodes[NodeType::Input];
+    return _tagged_nodes[type];
 }
 
 std::vector<std::unique_ptr<INode>> &Graph::nodes()
-- 
cgit v1.2.1