1 files changed, 40 insertions, 10 deletions

diff --git a/host/include/uhd/rfnoc/node_ctrl_base.ipp b/host/include/uhd/rfnoc/node_ctrl_base.ipp
index 79a8fe2a9..266df445a 100644
--- a/host/include/uhd/rfnoc/node_ctrl_base.ipp
+++ b/host/include/uhd/rfnoc/node_ctrl_base.ipp
@@ -27,27 +27,42 @@ namespace uhd {
         // List of return values:
         std::set< T_sptr > results_s;
         // To avoid cycles:
-        std::set< sptr > explored;
+        std::set < sptr > explored;
         // Initialize our search queue with ourself:
-        std::set< sptr > search_q;
-        search_q.insert(shared_from_this());
-        std::set< sptr > next_q;
+        std::set < node_map_pair_t > search_q;
+        // FIXME:  Port is initialized to ANY_PORT, but it should really be
+        // passed in by the caller.
+        search_q.insert(node_map_pair_t(ANY_PORT, shared_from_this()));
+        std::set < node_map_pair_t > next_q;
 
         while (iters++ < MAX_ITER) {
             next_q.clear();
-            BOOST_FOREACH(const sptr &this_node, search_q) {
+            BOOST_FOREACH(const node_map_pair_t node_pair, search_q) {
+                sptr node = node_pair.second.lock();
+                if (not node)
+                {
+                    continue;
+                }
+                size_t our_port = node_pair.first;
                 // Add this node to the list of explored nodes
-                explored.insert(this_node);
+                explored.insert(node);
                 // Create set of all child nodes of this_node that are not in explored:
-                std::set< sptr > next_nodes;
+                std::set< node_map_pair_t > next_nodes;
                 {
-                    node_map_t all_next_nodes = downstream ? this_node->list_downstream_nodes() : this_node->list_upstream_nodes();
+                    node_map_t all_next_nodes = downstream ?
+                        node->list_downstream_nodes() :
+                        node->list_upstream_nodes();
                     for (
                         node_map_t::iterator it = all_next_nodes.begin();
                         it != all_next_nodes.end();
                         ++it
                     ) {
-                        size_t our_port = it->first;
+                        size_t connected_port = it->first;
+                        // If port is given, limit traversal to only that port.
+                        if (our_port != ANY_PORT and our_port != connected_port)
+                        {
+                            continue;
+                        }
                         if (active_only
                             and not (downstream ? _tx_streamer_active[our_port] : _rx_streamer_active[our_port] )) {
                             continue;
@@ -60,7 +75,22 @@ namespace uhd {
                         if (next_node_sptr) {
                             results_s.insert(next_node_sptr);
                         } else {
-                            next_nodes.insert(one_next_node);
+                            size_t next_port = ANY_PORT;
+                            // FIXME:  Need proper mapping from input port
+                            // to output port.
+                            // The code below assumes that blocks with the same
+                            // number of connected upstream and downstream nodes
+                            // map each input port directly to the same output
+                            // port.  This limits the graph traversal to prevent
+                            // finding nodes that are not part of this chain.
+                            if (one_next_node->list_upstream_nodes().size() ==
+                                one_next_node->list_downstream_nodes().size())
+                            {
+                                next_port = (downstream ?
+                                    node->get_downstream_port(connected_port) :
+                                    node->get_upstream_port(connected_port));
+                            }
+                            next_nodes.insert(node_map_pair_t(next_port, it->second));
                         }
                     }
                 }