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;
class ShaderGraphRefactor;

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,
                GenContext& context);

    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& uniqueId);

    const ShaderNode* getNode(const string& uniqueId) 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, ShaderNode* shaderNode, GenContext& context);

    ShaderNode* createNode(ConstNodePtr node, GenContext& context);
 
    ShaderNode* inlineNodeBeforeOutput(ShaderGraphOutputSocket* output,
                                        const std::string& newNodeName,
                                        const std::string& nodeDefName,
                                        const std::string& inputName,
                                        const std::string& outputName,
                                        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; }

    ConstDocumentPtr getDocument() const { return _document; }

    ShaderNode* createNode(const string& name, const string& uniqueId, ConstNodeDefPtr nodeDef, GenContext& context);

    void bypass(ShaderNode* node, size_t inputIndex, size_t outputIndex = 0);

    void removeUnusedNodes();

    void replaceOutput(ShaderOutput* oldOutput, ShaderOutput* newOutput);
 
  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, GenContext& context);

    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 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::vector<std::pair<ShaderInput*, ColorSpaceTransform>> _inputColorTransformMap;
    // Temporary storage for inputs that require unit transformations
    std::vector<std::pair<ShaderInput*, UnitTransform>> _inputUnitTransformMap;
 
    // Temporary storage for outputs that require color transformations
    std::vector<std::pair<ShaderOutput*, ColorSpaceTransform>> _outputColorTransformMap;
    // Temporary storage for outputs that require unit transformations
    std::vector<std::pair<ShaderOutput*, UnitTransform>> _outputUnitTransformMap;
};

class MX_GENSHADER_API ShaderGraphEdge
{
  public:
    ShaderGraphEdge(ShaderOutput* up, ShaderInput* down) :
        upstream(up),
        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;
};

class MX_GENSHADER_API ShaderGraphEdgeIterator
{
  public:
    ShaderGraphEdgeIterator(ShaderOutput* output);
    ~ShaderGraphEdgeIterator() = default;
 
    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);
    bool skipOrMarkAsVisited(ShaderGraphEdge);
 
    ShaderOutput* _upstream;
    ShaderInput* _downstream;
    using StackFrame = std::pair<ShaderOutput*, size_t>;
    std::vector<StackFrame> _stack;
    std::set<ShaderOutput*> _path;
    std::set<ShaderGraphEdge> _visitedEdges;
};
 
MATERIALX_NAMESPACE_END
 
#endif