Fix performance issues with graph edge iteration in ShaderGraph

In complex materials graph and shader graph edge iteration can be
extremely slow, because some edges may be visited unnecessarily
multiple times. This is especially noticable in two functions:
ShaderGraph::addUpstreamDependencies and ShaderGraph::optimize() .

GraphIterator and ShaderGraphEdgeIterator classes iterate over DAGs
without marking nodes as visited, which may lead to exponential
traversal time for some DAGs:
https://stackoverflow.com/a/69326676

This patch adds two functions which efficiently generate a unique list
of (shader) graph edges and uses those lists instead of graph iterators.

source/MaterialXCore/Element.cpp

@@ -344,6 +344,11 @@ GraphIterator Element::traverseGraph() const
     return GraphIterator(getSelfNonConst());
 }
 
+vector<Edge> Element::uniqueGraphEdges() const
+{
+    return GraphIterator::uniqueGraphEdges(getSelfNonConst());
+}
+
 Edge Element::getUpstreamEdge(size_t) const
 {
     return NULL_EDGE;

source/MaterialXCore/Element.h

@@ -661,6 +661,9 @@ class MX_CORE_API Element : public std::enable_shared_from_this<Element>
     /// @sa getUpstreamElement
     GraphIterator traverseGraph() const;
 
+    /// Returns a list of all unique graph edges
+    vector<Edge> uniqueGraphEdges() const;
+
     /// Return the Edge with the given index that lies directly upstream from
     /// this element in the dataflow graph.
     /// @param index An optional index of the edge to be returned, where the

source/MaterialXCore/Traversal.cpp

@@ -177,6 +177,32 @@ void GraphIterator::returnPathDownstream(ElementPtr upstreamElem)
     _connectingElem = ElementPtr();
 }
 
+static void uniqueTraverse(ElementPtr upstreamElem, std::set<Edge>& visited, vector<Edge>& out)
+{
+    int upstreamEdgeCount = upstreamElem->getUpstreamEdgeCount();
+    for (int i = 0; i < upstreamEdgeCount; i++)
+    {
+        Edge edge = upstreamElem->getUpstreamEdge(i);
+        if (visited.find(edge) != visited.end() || !edge.getUpstreamElement())
+            continue;
+
+        visited.emplace(edge);
+        out.emplace_back(edge);
+        uniqueTraverse(edge.getUpstreamElement(), visited, out);
+    }
+}
+
+vector<Edge> GraphIterator::uniqueGraphEdges(ElementPtr root)
+{
+    vector<Edge> out;
+    std::set<Edge> visited;
+    if (root)
+    {
+        uniqueTraverse(root, visited, out);
+    }
+    return out;
+}
+
 //
 // InheritanceIterator methods
 //

source/MaterialXCore/Traversal.h

@@ -313,6 +313,9 @@ class MX_CORE_API GraphIterator
 
     /// @}
 
+    /// Returns a list of all unique graph edges.
+    static vector<Edge> uniqueGraphEdges(ElementPtr);
+
   private:
     void extendPathUpstream(ElementPtr upstreamElem, ElementPtr connectingElem);
     void returnPathDownstream(ElementPtr upstreamElem);

source/MaterialXGenShader/ShaderGraph.cpp

@@ -152,7 +152,7 @@ void ShaderGraph::addUpstreamDependencies(const Element& root, GenContext& conte
 {
     std::set<ElementPtr> processedOutputs;
 
-    for (Edge edge : root.traverseGraph())
+    for (Edge edge : root.uniqueGraphEdges())
     {
         ElementPtr upstreamElement = edge.getUpstreamElement();
         if (!upstreamElement)
@@ -900,7 +900,7 @@ void ShaderGraph::optimize()
             ShaderOutput* upstreamPort = outputSocket->getConnection();
             if (upstreamPort && upstreamPort->getNode() != this)
             {
-                for (ShaderGraphEdge edge : ShaderGraph::traverseUpstream(upstreamPort))
+                for (ShaderGraphEdge edge : ShaderGraphEdgeIterator::uniqueGraphEdges(upstreamPort))
                 {
                     ShaderNode* node = edge.upstream->getNode();
                     if (usedNodesSet.count(node) == 0)
@@ -1259,4 +1259,32 @@ void ShaderGraphEdgeIterator::returnPathDownstream(ShaderOutput* upstream)
     _downstream = nullptr;
 }
 
+static void uniqueTraverse(ShaderNode* upstreamNode, std::set<ShaderGraphEdge>& visited, vector<ShaderGraphEdge>& out)
+{
+    for (size_t i = 0, numInputs = upstreamNode->numInputs(); i < numInputs; i++)
+    {
+        ShaderInput* input = upstreamNode->getInput(i);
+        ShaderOutput* output = input->getConnection();
+        ShaderGraphEdge edge(output, input);
+
+        if (visited.find(edge) != visited.end() || !output)
+            continue;
+
+        visited.emplace(edge);
+        out.emplace_back(edge);
+        uniqueTraverse(output->getNode(), visited, out);
+    }
+}
+
+vector<ShaderGraphEdge> ShaderGraphEdgeIterator::uniqueGraphEdges(ShaderOutput* root)
+{
+    vector<ShaderGraphEdge> out;
+    std::set<ShaderGraphEdge> visited;
+    if (root)
+    {
+        uniqueTraverse(root->getNode(), visited, out);
+    }
+    return out;
+}
+
 MATERIALX_NAMESPACE_END

source/MaterialXGenShader/ShaderGraph.h

@@ -209,6 +209,22 @@ class MX_GENSHADER_API ShaderGraphEdge
         downstream(down)
     {
     }
+
+    bool operator==(const ShaderGraphEdge& rhs) const
+    {
+        return upstream == rhs.upstream && downstream == rhs.downstream;
+    }
+
+    bool operator!=(const ShaderGraphEdge& rhs) const
+    {
+        return !(*this == rhs);
+    }
+
+    bool operator<(const ShaderGraphEdge& rhs) const
+    {
+        return std::tie(upstream, downstream) < std::tie(rhs.upstream, rhs.downstream);
+    }
+
     ShaderOutput* upstream;
     ShaderInput* downstream;
 };
@@ -251,6 +267,9 @@ class MX_GENSHADER_API ShaderGraphEdgeIterator
     /// Return the end iterator.
     static const ShaderGraphEdgeIterator& end();
 
+    /// Returns a list of all unique ShaderGraph edges
+    static vector<ShaderGraphEdge> uniqueGraphEdges(ShaderOutput*);
+
   private:
     void extendPathUpstream(ShaderOutput* upstream, ShaderInput* downstream);
     void returnPathDownstream(ShaderOutput* upstream);