Fix downstream traversal in graph editor

source/MaterialXCore/Document.cpp

@@ -97,6 +97,18 @@ class Document::Cache
                         portElementMap.emplace(portElem->getQualifiedName(nodeName), portElem);
                     }
                 }
+                else
+                {
+                    const string& nodegraphName = elem->getAttribute(PortElement::NODE_GRAPH_ATTRIBUTE);
+                    if (!nodegraphName.empty())
+                    {
+                        PortElementPtr portElem = elem->asA<PortElement>();
+                        if (portElem)
+                        {
+                            portElementMap.emplace(portElem->getQualifiedName(nodegraphName), portElem);
+                        }
+                    }
+                }
                 if (!nodeString.empty())
                 {
                     NodeDefPtr nodeDef = elem->asA<NodeDef>();

source/MaterialXCore/Node.cpp

@@ -733,6 +733,33 @@ InterfaceElementPtr NodeGraph::getImplementation() const
     return nodedef ? nodedef->getImplementation() : InterfaceElementPtr();
 }
 
+vector<PortElementPtr> NodeGraph::getDownstreamPorts() const
+{
+    vector<PortElementPtr> downstreamPorts;
+    for (PortElementPtr port : getDocument()->getMatchingPorts(getQualifiedName(getName())))
+    {
+        ConstGraphElementPtr graph = nullptr;
+        ElementPtr node = port->getParent();
+        if (node)
+        {
+            node = node->getParent();
+            if (node->isA<GraphElement>())
+            {
+                graph = node->asA<GraphElement>();
+            }
+        }
+        if (graph && graph->getChild(getName()) == getSelf())
+        {
+            downstreamPorts.push_back(port);
+        }
+    }
+    std::sort(downstreamPorts.begin(), downstreamPorts.end(), [](const ConstElementPtr& a, const ConstElementPtr& b)
+    {
+        return a->getName() > b->getName();
+    });
+    return downstreamPorts;
+}
+
 bool NodeGraph::validate(string* message) const
 {
     bool res = true;

source/MaterialXCore/Node.h

@@ -358,6 +358,14 @@ class MX_CORE_API NodeGraph : public GraphElement
     ///    none was found.
     InterfaceElementPtr getImplementation() const;
 
+    /// @}
+    /// @name Traversal
+    /// @{
+
+    /// Return a vector of all downstream ports that connect to this graph, ordered by
+    /// the names of the port elements.
+    vector<PortElementPtr> getDownstreamPorts() const;
+
     /// @}
     /// @name Utility
     /// @{

source/MaterialXGraphEditor/Graph.cpp

@@ -681,8 +681,13 @@ void Graph::selectMaterial(UiNodePtr uiNode)
 // set the node to display in render veiw based off the selected node or nodegraph
 void Graph::setRenderMaterial(UiNodePtr node)
 {
+    // For now surface shaders and materials are considered renderable.
+    // This can be adjusted as desired to includ being able to use outputs,
+    // and / a sub-graph in the nodegraph.
+    mx::StringSet renderableTypes = { "material", "surfaceshader" };
+
     // set render node right away is node is a material
-    if (node->getNode() && node->getNode()->getType() == "material")
+    if (node->getNode() && renderableTypes.count(node->getNode()->getType()))
     {
         // only set new render node if different material has been selected
         if (_currRenderNode != node)
@@ -694,92 +699,124 @@ void Graph::setRenderMaterial(UiNodePtr node)
     }
     else
     {
-        // continue downstream using output connections until a material node is found
-        std::vector<UiNodePtr> outNodes = node->getOutputConnections();
-        if (outNodes.size() > 0)
+        // Traverse downstream looking for the first renderable item found.
+        mx::NodePtr mtlxNode = node->getNode();
+        mx::NodeGraphPtr mtlxNodeGraph = node->getNodeGraph();
+        mx::OutputPtr mtlxOutput = node->getOutput();
+        if (mtlxOutput)
+        {
+            mx::ElementPtr parent = mtlxOutput->getParent();
+            if (parent->isA<mx::NodeGraph>())
+                mtlxNodeGraph = parent->asA<mx::NodeGraph>();
+            else if (parent->isA<mx::Node>())
+                mtlxNode = parent->asA<mx::Node>();
+        }
+        mx::StringSet testPaths;
+        if (mtlxNode)
+        {
+            testPaths.insert(mtlxNode->getNamePath());
+        }
+        else if (mtlxNodeGraph)
+        {
+            testPaths.insert(mtlxNodeGraph->getNamePath());
+        }
+
+        mx::NodePtr foundNode = nullptr;
+        while (!testPaths.empty() && !foundNode)
         {
-            if (outNodes[0]->getNode())
+            mx::StringSet nextPaths;
+            for (const std::string& testPath : testPaths)
             {
-                if (outNodes[0]->getNode()->getType() == mx::SURFACE_SHADER_TYPE_STRING)
+                mx::ElementPtr testElem = _graphDoc->getDescendant(testPath);
+                mx::NodePtr testNode = testElem->asA<mx::Node>();
+                std::vector<mx::PortElementPtr> downstreamPorts;
+                if (testNode)
+                {
+                    downstreamPorts = testNode->getDownstreamPorts();
+                }
+                else
                 {
-                    std::vector<UiNodePtr> shaderOut = outNodes[0]->getOutputConnections();
-                    if (shaderOut.size() > 0)
+                    mx::NodeGraphPtr testGraph = testElem->asA<mx::NodeGraph>();
+                    if (testGraph)
                     {
-                        if (shaderOut[0])
+                        downstreamPorts = testGraph->getDownstreamPorts();
+                    }
+                }
+                // Test all downstream ports. If the port's node is renderable
+                // then stop searching.
+                for (mx::PortElementPtr downstreamPort : downstreamPorts)
+                {
+                    mx::ElementPtr parent = downstreamPort->getParent();
+                    if (parent)
+                    {
+                        mx::NodePtr downstreamNode = parent->asA<mx::Node>();
+                        if (downstreamNode)
                         {
-                            if (shaderOut[0]->getNode()->getType() == "material")
+                            mx::NodeDefPtr nodeDef = downstreamNode->getNodeDef();
+                            if (nodeDef)
                             {
-                                if (_currRenderNode != shaderOut[0])
+                                if (renderableTypes.count(nodeDef->getType()))
                                 {
-                                    _currRenderNode = shaderOut[0];
-                                    _frameCount = ImGui::GetFrameCount();
-                                    _renderer->setMaterialCompilation(true);
+                                    foundNode = downstreamNode;
+                                    break;
                                 }
                             }
                         }
-                    }
-                    else
-                    {
-                        _currRenderNode = nullptr;
+                        if (!foundNode)
+                        {
+                            nextPaths.insert(parent->getNamePath());
+                        }
                     }
                 }
-                else if (outNodes[0]->getNode()->getType() == mx::MATERIAL_TYPE_STRING)
+                if (foundNode)
                 {
-                    if (_currRenderNode != outNodes[0])
+                    break;
+                }
+            }
+
+            // Set up next set of nodes to search downstream
+            testPaths = nextPaths;
+        }
+
+        // Update rendering. If found use that node, otherwise
+        // use the current fallback of using the first renderable node.
+        if (foundNode)
+        {
+            for (auto uiNode : _graphNodes)
+            {
+                if (uiNode->getNode() == foundNode)
+                {
+                    if (_currRenderNode != uiNode)
                     {
-                        _currRenderNode = outNodes[0];
+                        _currRenderNode = uiNode;
                         _frameCount = ImGui::GetFrameCount();
                         _renderer->setMaterialCompilation(true);
                     }
+                    break;
                 }
             }
-            else
-            {
-                _currRenderNode = nullptr;
-            }
         }
         else
         {
             _currRenderNode = nullptr;
+            _frameCount = ImGui::GetFrameCount();
+            _renderer->setMaterialCompilation(true);
         }
     }
 }
 
 void Graph::updateMaterials(mx::InputPtr input, mx::ValuePtr value)
 {
     std::string renderablePath;
-    mx::TypedElementPtr renderableElem;
-    std::vector<mx::TypedElementPtr> elems = mx::findRenderableElements(_graphDoc);
-
-    size_t num = 0;
-    int num2 = 0;
-    for (mx::TypedElementPtr elem : elems)
+    if (_currRenderNode)
     {
-        renderableElem = elem;
-        mx::NodePtr node = elem->asA<mx::Node>();
-        if (node)
+        if (_currRenderNode->getNode())
         {
-            if (_currRenderNode)
-            {
-                if (node->getName() == _currRenderNode->getName())
-                {
-                    renderablePath = renderableElem->getNamePath();
-                    break;
-                }
-            }
-            else
-            {
-                renderablePath = renderableElem->getNamePath();
-            }
+            renderablePath = _currRenderNode->getNode()->getNamePath();
         }
-        else
+        else if (_currRenderNode->getOutput())
         {
-            renderablePath = renderableElem->getNamePath();
-            if (num2 == 2)
-            {
-                break;
-            }
-            num2++;
+            renderablePath = _currRenderNode->getOutput()->getNamePath();
         }
     }
 
@@ -810,7 +847,7 @@ void Graph::updateMaterials(mx::InputPtr input, mx::ValuePtr value)
             // Note that if there is a topogical change due to
             // this value change or a transparency change, then
             // this is not currently caught here.
-            _renderer->getMaterials()[num]->modifyUniform(name, value);
+            _renderer->getMaterials()[0]->modifyUniform(name, value);
         }
     }
 }
@@ -3619,10 +3656,7 @@ void Graph::drawGraph(ImVec2 mousePos)
                     // update render material if needed
                     if (_currUiNode->getNode())
                     {
-                        if (_currUiNode->getNode()->getType() == mx::SURFACE_SHADER_TYPE_STRING || _currUiNode->getNode()->getType() == mx::MATERIAL_TYPE_STRING)
-                        {
-                            setRenderMaterial(_currUiNode);
-                        }
+                        setRenderMaterial(_currUiNode);
                     }
                     else if (_currUiNode->getNodeGraph())
                     {

source/PyMaterialX/PyMaterialXCore/PyNode.cpp

@@ -59,6 +59,7 @@ void bindPyNode(py::module& mod)
         .def("addInterfaceName", &mx::NodeGraph::addInterfaceName)
         .def("removeInterfaceName", &mx::NodeGraph::removeInterfaceName)
         .def("modifyInterfaceName", &mx::NodeGraph::modifyInterfaceName)
+        .def("getDownstreamPorts", &mx::NodeGraph::getDownstreamPorts)
         .def_readonly_static("CATEGORY", &mx::NodeGraph::CATEGORY);
 
     py::class_<mx::Backdrop, mx::BackdropPtr, mx::Element>(mod, "Backdrop")