MaterialX/source/MaterialXCore/Material.cpp at main · niklasharrysson/MaterialX · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
//
// TM & (c) 2017 Lucasfilm Entertainment Company Ltd. and Lucasfilm Ltd.
// All rights reserved.  See LICENSE.txt for license.
//

#include <MaterialXCore/Material.h>

MATERIALX_NAMESPACE_BEGIN

vector<NodePtr> getShaderNodes(NodePtr materialNode, const string& nodeType, const string& target)
{
    vector<NodePtr> shaderNodeVec;
    std::set<NodePtr> shaderNodeSet;

    vector<InputPtr> inputs = materialNode->getActiveInputs();
    for (InputPtr input : inputs)
    {
        // Scan for a node directly connected to the input.
        // Note that this will handle traversing through interfacename associations.
        NodePtr shaderNode = input->getConnectedNode();
        if (shaderNode && !shaderNodeSet.count(shaderNode))
        {
            if (!nodeType.empty() && shaderNode->getType() != nodeType)
            {
                continue;
            }

            if (!target.empty())
            {
                NodeDefPtr nodeDef = shaderNode->getNodeDef(target);
                if (!nodeDef)
                {
                    continue;
                }
            }

            shaderNodeVec.push_back(shaderNode);
            shaderNodeSet.insert(shaderNode);
        }
        else if (input->hasNodeGraphString())
        {
            // Check upstream nodegraph connected to the input.
            // If no explicit output name given then scan all outputs on the nodegraph.
            ElementPtr parent = materialNode->getParent();
            NodeGraphPtr nodeGraph = parent->getChildOfType<NodeGraph>(input->getNodeGraphString());
            if (!nodeGraph)
            {
                continue;
            }
            vector<OutputPtr> outputs;
            if (input->hasOutputString())
            {
                outputs.push_back(nodeGraph->getOutput(input->getOutputString()));
            }
            else
            {
                outputs = nodeGraph->getOutputs();
            }
            for (OutputPtr output : outputs)
            {
                NodePtr upstreamNode = output->getConnectedNode();
                if (upstreamNode && !shaderNodeSet.count(upstreamNode))
                {
                    if (!target.empty() && !upstreamNode->getNodeDef(target))
                    {
                        continue;
                    }
                    shaderNodeVec.push_back(upstreamNode);
                    shaderNodeSet.insert(upstreamNode);
                }
            }
        }
    }

    if (inputs.empty())
    {
        // Try to find material nodes in the implementation graph if any.
        // If a target is specified the nodedef for the given target is searched for.
        NodeDefPtr materialNodeDef = materialNode->getNodeDef(target);
        if (materialNodeDef)
        {
            InterfaceElementPtr impl = materialNodeDef->getImplementation();
            if (impl && impl->isA<NodeGraph>())
            {
                NodeGraphPtr implGraph = impl->asA<NodeGraph>();
                for (OutputPtr defOutput : materialNodeDef->getOutputs())
                {
                    if (defOutput->getType() == MATERIAL_TYPE_STRING)
                    {
                        OutputPtr implGraphOutput = implGraph->getOutput(defOutput->getName());
                        for (GraphIterator it = implGraphOutput->traverseGraph().begin(); it != GraphIterator::end(); ++it)
                        {
                            ElementPtr upstreamElem = it.getUpstreamElement();
                            if (!upstreamElem)
                            {
                                it.setPruneSubgraph(true);
                                continue;
                            }
                            NodePtr upstreamNode = upstreamElem->asA<Node>();
                            if (upstreamNode && upstreamNode->getType() == MATERIAL_TYPE_STRING)
                            {
                                for (NodePtr shaderNode : getShaderNodes(upstreamNode, nodeType, target))
                                {
                                    if (!shaderNodeSet.count(shaderNode))
                                    {
                                        shaderNodeVec.push_back(shaderNode);
                                        shaderNodeSet.insert(shaderNode);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    return shaderNodeVec;
}

vector<OutputPtr> getConnectedOutputs(NodePtr node)
{
    vector<OutputPtr> outputVec;
    std::set<OutputPtr> outputSet;
    for (InputPtr input : node->getInputs())
    {
        OutputPtr output = input->getConnectedOutput();
        if (output && !outputSet.count(output))
        {
            outputVec.push_back(output);
            outputSet.insert(output);
        }
    }
    return outputVec;
}

MATERIALX_NAMESPACE_END