Unverified Commit 2427b205 authored by Scott Cyphers's avatar Scott Cyphers Committed by GitHub

Implement traverse nodes with pointers (#3205)

parent 9b3d5732
...@@ -69,39 +69,34 @@ void ngraph::traverse_nodes(const NodeVector& subgraph_results, ...@@ -69,39 +69,34 @@ void ngraph::traverse_nodes(const NodeVector& subgraph_results,
bool include_control_deps, bool include_control_deps,
const NodeVector& subgraph_params) const NodeVector& subgraph_params)
{ {
std::unordered_set<std::shared_ptr<Node>> instances_seen{subgraph_params.begin(), std::unordered_set<Node*> instances_seen;
subgraph_params.end()}; std::deque<Node*> stack;
std::deque<std::shared_ptr<Node>> stack; for (auto& node_ptr : subgraph_params)
for (auto r : subgraph_results)
{ {
stack.push_front(r); instances_seen.insert(node_ptr.get());
}
for (auto& node_ptr : subgraph_results)
{
stack.push_front(node_ptr.get());
} }
while (stack.size() > 0) while (stack.size() > 0)
{ {
std::shared_ptr<Node> n = stack.front(); Node* n = stack.front();
stack.pop_front(); stack.pop_front();
if (instances_seen.count(n) == 0) if (instances_seen.insert(n).second)
{ {
instances_seen.insert(n); f(n->shared_from_this());
f(n); for (auto& arg : n->get_arguments())
for (auto arg : n->get_arguments())
{ {
if (instances_seen.count(arg) == 0) stack.push_front(arg.get());
{
stack.push_front(arg);
}
} }
if (include_control_deps) if (include_control_deps)
{ {
for (auto cdep : n->get_control_dependencies()) for (auto& cdep : n->get_control_dependencies())
{ {
if (instances_seen.count(cdep) == 0) stack.push_front(cdep.get());
{
stack.push_front(cdep);
}
} }
} }
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment