ShaderGraph.h

Go to the documentation of this file.

//
// Copyright Contributors to the MaterialX Project
// SPDX-License-Identifier: Apache-2.0
//
 
#ifndef MATERIALX_SHADERGRAPH_H
#define MATERIALX_SHADERGRAPH_H
 
 
#include <MaterialXGenShader/Export.h>
 
#include <MaterialXGenShader/ColorManagementSystem.h>
#include <MaterialXGenShader/ShaderNode.h>
#include <MaterialXGenShader/Syntax.h>
#include <MaterialXGenShader/TypeDesc.h>
#include <MaterialXGenShader/UnitSystem.h>
 
#include <MaterialXCore/Document.h>
#include <MaterialXCore/Node.h>
 
MATERIALX_NAMESPACE_BEGIN
 
class Syntax;
class ShaderGraphEdge;
class ShaderGraphEdgeIterator;
class GenOptions;
 
using ShaderGraphInputSocket = ShaderOutput;
 
using ShaderGraphOutputSocket = ShaderInput;
 
using ShaderGraphPtr = shared_ptr<class ShaderGraph>;
 
class MX_GENSHADER_API ShaderGraph : public ShaderNode
{
  public:
    ShaderGraph(const ShaderGraph* parent, const string& name, ConstDocumentPtr document, const StringSet& reservedWords);
 
    virtual ~ShaderGraph() { }
 
    static ShaderGraphPtr create(const ShaderGraph* parent, const string& name, ElementPtr element,
                                 GenContext& context);
 
    static ShaderGraphPtr create(const ShaderGraph* parent, const NodeGraph& nodeGraph,
                                 GenContext& context);
 
    bool isAGraph() const override { return true; }
 
    ShaderNode* getNode(const string& name);
 
    const ShaderNode* getNode(const string& name) const;
 
    const vector<ShaderNode*>& getNodes() const { return _nodeOrder; }
 
    size_t numInputSockets() const { return numOutputs(); }
 
    size_t numOutputSockets() const { return numInputs(); }
 
    ShaderGraphInputSocket* getInputSocket(size_t index) { return getOutput(index); }
    ShaderGraphOutputSocket* getOutputSocket(size_t index = 0) { return getInput(index); }
    const ShaderGraphInputSocket* getInputSocket(size_t index) const { return getOutput(index); }
    const ShaderGraphOutputSocket* getOutputSocket(size_t index = 0) const { return getInput(index); }
 
    ShaderGraphInputSocket* getInputSocket(const string& name) { return getOutput(name); }
    ShaderGraphOutputSocket* getOutputSocket(const string& name) { return getInput(name); }
    const ShaderGraphInputSocket* getInputSocket(const string& name) const { return getOutput(name); }
    const ShaderGraphOutputSocket* getOutputSocket(const string& name) const { return getInput(name); }
 
    const vector<ShaderGraphInputSocket*>& getInputSockets() const { return _outputOrder; }
    const vector<ShaderGraphOutputSocket*>& getOutputSockets() const { return _inputOrder; }
 
    void applyInputTransforms(ConstNodePtr node, ShaderNodePtr shaderNode, GenContext& context);
 
    ShaderNode* createNode(ConstNodePtr node, GenContext& context);
 
    ShaderGraphInputSocket* addInputSocket(const string& name, TypeDesc type);
    [[deprecated]] ShaderGraphInputSocket* addInputSocket(const string& name, const TypeDesc* type) { return addInputSocket(name, *type); }
 
    ShaderGraphOutputSocket* addOutputSocket(const string& name, TypeDesc type);
    [[deprecated]] ShaderGraphOutputSocket* addOutputSocket(const string& name, const TypeDesc* type) { return addOutputSocket(name, *type); }
 
    void addDefaultGeomNode(ShaderInput* input, const GeomPropDef& geomprop, GenContext& context);
 
    void topologicalSort();
 
    static ShaderGraphEdgeIterator traverseUpstream(ShaderOutput* output);
 
    IdentifierMap& getIdentifierMap() { return _identifiers; }
 
  protected:
    void createConnectedNodes(const ElementPtr& downstreamElement,
                              const ElementPtr& upstreamElement,
                              ElementPtr connectingElement,
                              GenContext& context);
 
    void addNode(ShaderNodePtr node);
 
    void addInputSockets(const InterfaceElement& elem, GenContext& context);
 
    void addOutputSockets(const InterfaceElement& elem);
 
    void addUpstreamDependencies(const Element& root, GenContext& context);
 
    void addColorTransformNode(ShaderInput* input, const ColorSpaceTransform& transform, GenContext& context);
 
    void addColorTransformNode(ShaderOutput* output, const ColorSpaceTransform& transform, GenContext& context);
 
    void addUnitTransformNode(ShaderInput* input, const UnitTransform& transform, GenContext& context);
 
    void addUnitTransformNode(ShaderOutput* output, const UnitTransform& transform, GenContext& context);
 
    void finalize(GenContext& context);
 
    void optimize();
 
    void bypass(ShaderNode* node, size_t inputIndex, size_t outputIndex = 0);
 
    void setVariableNames(GenContext& context);
 
    void populateColorTransformMap(ColorManagementSystemPtr colorManagementSystem, ShaderPort* shaderPort,
                                   const string& sourceColorSpace, const string& targetColorSpace, bool asInput);
 
    void populateUnitTransformMap(UnitSystemPtr unitSystem, ShaderPort* shaderPort, ValueElementPtr element, const string& targetUnitSpace, bool asInput);
 
    void disconnect(ShaderNode* node) const;
 
    ConstDocumentPtr _document;
    std::unordered_map<string, ShaderNodePtr> _nodeMap;
    std::vector<ShaderNode*> _nodeOrder;
    IdentifierMap _identifiers;
 
    // Temporary storage for inputs that require color transformations
    std::unordered_map<ShaderInput*, ColorSpaceTransform> _inputColorTransformMap;
    // Temporary storage for inputs that require unit transformations
    std::unordered_map<ShaderInput*, UnitTransform> _inputUnitTransformMap;
 
    // Temporary storage for outputs that require color transformations
    std::unordered_map<ShaderOutput*, ColorSpaceTransform> _outputColorTransformMap;
    // Temporary storage for outputs that require unit transformations
    std::unordered_map<ShaderOutput*, UnitTransform> _outputUnitTransformMap;
};
 
class MX_GENSHADER_API ShaderGraphEdge
{
  public:
    ShaderGraphEdge(ShaderOutput* up, ShaderInput* down) :
        upstream(up),
        downstream(down)
    {
    }
    ShaderOutput* upstream;
    ShaderInput* downstream;
};
 
class MX_GENSHADER_API ShaderGraphEdgeIterator
{
  public:
    ShaderGraphEdgeIterator(ShaderOutput* output);
    ~ShaderGraphEdgeIterator() { }
 
    bool operator==(const ShaderGraphEdgeIterator& rhs) const
    {
        return _upstream == rhs._upstream &&
               _downstream == rhs._downstream &&
               _stack == rhs._stack;
    }
    bool operator!=(const ShaderGraphEdgeIterator& rhs) const
    {
        return !(*this == rhs);
    }
 
    ShaderGraphEdge operator*() const
    {
        return ShaderGraphEdge(_upstream, _downstream);
    }
 
    ShaderGraphEdgeIterator& operator++();
 
    ShaderGraphEdgeIterator& begin()
    {
        return *this;
    }
 
    static const ShaderGraphEdgeIterator& end();
 
  private:
    void extendPathUpstream(ShaderOutput* upstream, ShaderInput* downstream);
    void returnPathDownstream(ShaderOutput* upstream);
 
    ShaderOutput* _upstream;
    ShaderInput* _downstream;
    using StackFrame = std::pair<ShaderOutput*, size_t>;
    std::vector<StackFrame> _stack;
    std::set<ShaderOutput*> _path;
};
 
MATERIALX_NAMESPACE_END
 
#endif