Merge pull request #3189 from Autodesk/samuelliu-adsk/EMSUSD-101/over…

…ride-primwriter

EMSUSD-101 Override primWriter

lib/mayaUsd/fileio/primWriter.h

@@ -61,6 +61,25 @@ class UsdMayaPrimWriter
         const SdfPath&           usdPath,
         UsdMayaWriteJobContext&  jobCtx);
 
+    /// The level of support a writer can offer for a given context
+    ///
+    /// A basic writer that gives correct results across most contexts should
+    /// report `Fallback`, while a specialized writer that really shines in a
+    /// given context should report `Supported` when the context is right and
+    /// `Unsupported` if the context is not as expected.
+    enum class ContextSupport
+    {
+        Supported,
+        Fallback,
+        Unsupported
+    };
+
+    /// A static function is expected for all writers and allows
+    /// declaring how well this class can support the current context.
+    MAYAUSD_CORE_PUBLIC
+    static ContextSupport
+    CanExport(const UsdMayaJobExportArgs& exportArgs, const MObject& exportObj);
+
     MAYAUSD_CORE_PUBLIC
     virtual ~UsdMayaPrimWriter();
 

lib/mayaUsd/fileio/primWriterRegistry.cpp

@@ -42,9 +42,75 @@ TF_DEFINE_PRIVATE_TOKENS(
 );
 // clang-format on
 
-typedef std::map<std::string, UsdMayaPrimWriterRegistry::WriterFactoryFn> _Registry;
-static _Registry                                                          _reg;
-static std::set<std::string> _mayaTypesThatDoNotCreatePrims;
+namespace {
+struct _RegistryEntry
+{
+    UsdMayaPrimWriterRegistry::ContextPredicateFn _pred;
+    UsdMayaPrimWriterRegistry::WriterFactoryFn    _writer;
+    int                                           _index;
+};
+
+// typedef std::map<std::string, UsdMayaPrimWriterRegistry::WriterFactoryFn> _Registry;
+typedef std::unordered_multimap<std::string, _RegistryEntry> _Registry;
+static _Registry                                             _reg;
+static std::set<std::string>                                 _mayaTypesThatDoNotCreatePrims;
+static int                                                   _indexCounter = 0;
+
+_Registry::const_iterator _Find(
+    const std::string&          mayaTypeName,
+    const UsdMayaJobExportArgs& exportArgs,
+    const MObject&              exportObj)
+{
+    using ContextSupport = UsdMayaPrimWriter::ContextSupport;
+
+    _Registry::const_iterator ret = _reg.cend();
+    _Registry::const_iterator first, last;
+    std::tie(first, last) = _reg.equal_range(mayaTypeName);
+    while (first != last) {
+        ContextSupport support = first->second._pred(exportArgs, exportObj);
+        // Look for a "Supported" reader. If no "Supported" reader is found, use a "Fallback" reader
+        if (support == ContextSupport::Supported) {
+            ret = first;
+            break;
+        } else if (support == ContextSupport::Fallback && ret == _reg.end()) {
+            ret = first;
+        }
+        ++first;
+    }
+
+    return ret;
+}
+
+} // namespace
+
+/* static */
+void UsdMayaPrimWriterRegistry::Register(
+    const std::string&                            mayaTypeName,
+    UsdMayaPrimWriterRegistry::ContextPredicateFn pred,
+    UsdMayaPrimWriterRegistry::WriterFactoryFn    fn,
+    bool                                          fromPython)
+{
+    TF_DEBUG(PXRUSDMAYA_REGISTRY)
+        .Msg("Registering UsdMayaPrimWriter for maya type %s.\n", mayaTypeName.c_str());
+
+    int index = _indexCounter++;
+    _reg.insert(std::make_pair(mayaTypeName, _RegistryEntry { pred, fn, index }));
+
+    // The unloader uses the index to know which entry to erase when there are
+    // more than one for the same mayaTypeName.
+    UsdMaya_RegistryHelper::AddUnloader(
+        [mayaTypeName, index]() {
+            _Registry::const_iterator it, itEnd;
+            std::tie(it, itEnd) = _reg.equal_range(mayaTypeName);
+            for (; it != itEnd; ++it) {
+                if (it->second._index == index) {
+                    _reg.erase(it);
+                    break;
+                }
+            }
+        },
+        fromPython);
+}
 
 /* static */
 void UsdMayaPrimWriterRegistry::Register(
@@ -55,14 +121,31 @@ void UsdMayaPrimWriterRegistry::Register(
     TF_DEBUG(PXRUSDMAYA_REGISTRY)
         .Msg("Registering UsdMayaPrimWriter for maya type %s.\n", mayaTypeName.c_str());
 
-    std::pair<_Registry::iterator, bool> insertStatus
-        = _reg.insert(std::make_pair(mayaTypeName, fn));
-    if (insertStatus.second) {
-        UsdMaya_RegistryHelper::AddUnloader(
-            [mayaTypeName]() { _reg.erase(mayaTypeName); }, fromPython);
-    } else {
-        TF_CODING_ERROR("Multiple writers for type %s", mayaTypeName.c_str());
-    }
+    int index = _indexCounter++;
+
+    // Use default ContextSupport if not specified
+    _reg.insert(std::make_pair(
+        mayaTypeName,
+        _RegistryEntry { [](const UsdMayaJobExportArgs&, const MObject&) {
+                            return UsdMayaPrimWriter::ContextSupport::Fallback;
+                        },
+                         fn,
+                         index }));
+
+    // The unloader uses the index to know which entry to erase when there are
+    // more than one for the same mayaTypeName.
+    UsdMaya_RegistryHelper::AddUnloader(
+        [mayaTypeName, index]() {
+            _Registry::const_iterator it, itEnd;
+            std::tie(it, itEnd) = _reg.equal_range(mayaTypeName);
+            for (; it != itEnd; ++it) {
+                if (it->second._index == index) {
+                    _reg.erase(it);
+                    break;
+                }
+            }
+        },
+        fromPython);
 }
 
 /* static */
@@ -74,32 +157,59 @@ void UsdMayaPrimWriterRegistry::RegisterRaw(
 }
 
 /* static */
-UsdMayaPrimWriterRegistry::WriterFactoryFn
-UsdMayaPrimWriterRegistry::Find(const std::string& mayaTypeName)
+UsdMayaPrimWriterRegistry::WriterFactoryFn UsdMayaPrimWriterRegistry::Find(
+    const std::string&          mayaTypeName,
+    const UsdMayaJobExportArgs& exportArgs,
+    const MObject&              exportObj)
 {
     TfRegistryManager::GetInstance().SubscribeTo<UsdMayaPrimWriterRegistry>();
 
-    // unfortunately, usdTypeName is diff from the tfTypeName which we use to
-    // register.  do the conversion here.
-    WriterFactoryFn ret = nullptr;
-    if (TfMapLookup(_reg, mayaTypeName, &ret)) {
-        return ret;
+    _Registry::const_iterator it = _Find(mayaTypeName, exportArgs, exportObj);
+
+    if (it != _reg.end()) {
+        return it->second._writer;
     }
 
     static const TfTokenVector SCOPE = { _tokens->UsdMaya, _tokens->PrimWriter };
     UsdMaya_RegistryHelper::FindAndLoadMayaPlug(SCOPE, mayaTypeName);
 
-    // ideally something just registered itself.  if not, we at least put it in
+    it = _Find(mayaTypeName, exportArgs, exportObj);
+
+    if (it != _reg.end()) {
+        return it->second._writer;
+    }
+
+    // if nothing was found and nothing was registered, we at least put it in
     // the registry in case we encounter it again.
-    if (!TfMapLookup(_reg, mayaTypeName, &ret)) {
-        TF_DEBUG(PXRUSDMAYA_REGISTRY)
-            .Msg(
-                "No usdMaya writer plugin for maya type %s. No maya plugin found.\n",
-                mayaTypeName.c_str());
-        _reg[mayaTypeName] = nullptr;
+    TF_DEBUG(PXRUSDMAYA_REGISTRY)
+        .Msg("No usdMaya writer plugin for TfType %s. No maya plugin.\n", mayaTypeName.c_str());
+    if (_reg.count(mayaTypeName) == 0) {
+        // Nothing registered at all, remember that:
+        _reg.insert(std::make_pair(
+            mayaTypeName,
+            _RegistryEntry { [](const UsdMayaJobExportArgs&, const MObject&) {
+                                return UsdMayaPrimWriter::ContextSupport::Fallback;
+                            },
+                             nullptr,
+                             -1 }));
     }
 
-    return ret;
+    return nullptr;
+}
+
+/* static */
+void UsdMayaPrimWriterRegistry::CheckForWriterPlugin(const std::string& mayaTypeName)
+{
+    TfRegistryManager::GetInstance().SubscribeTo<UsdMayaPrimWriterRegistry>();
+
+    _Registry::const_iterator first, last;
+    std::tie(first, last) = _reg.equal_range(mayaTypeName);
+
+    if (first == last) {
+        // If the type name is not currently in our registry, check for plugin registry
+        static const TfTokenVector SCOPE = { _tokens->UsdMaya, _tokens->PrimWriter };
+        UsdMaya_RegistryHelper::FindAndLoadMayaPlug(SCOPE, mayaTypeName);
+    }
 }
 
 /* static */

lib/mayaUsd/fileio/primWriterRegistry.h

@@ -77,6 +77,39 @@ struct UsdMayaPrimWriterRegistry
     /// macro.
     typedef std::function<bool(const UsdMayaPrimWriterArgs&, UsdMayaPrimWriterContext*)> WriterFn;
 
+    /// Predicate function, i.e. a function that can tell the level of support
+    /// the writer function will provide for a given set of export options.
+    using ContextPredicateFn = std::function<
+        UsdMayaPrimWriter::ContextSupport(const UsdMayaJobExportArgs&, const MObject&)>;
+
+    /// \brief Register \p fn as a factory function providing a
+    /// UsdMayaPrimWriter subclass that can be used to write \p mayaType.
+    /// Provide a supportability of the primWriter. Use "supported" to
+    /// override the default primWriter
+    ///
+    /// If you can't provide a valid UsdMayaPrimWriter for the given arguments,
+    /// return a null pointer from the factory function \p fn.
+    ///
+    /// Example for registering a writer factory in your custom plugin:
+    /// \code{.cpp}
+    /// class MyWriter : public UsdMayaPrimWriter {
+    ///     static UsdMayaPrimWriterSharedPtr Create(
+    ///             const MFnDependencyNode& depNodeFn,
+    ///             const SdfPath& usdPath,
+    ///             UsdMayaWriteJobContext& jobCtx);
+    /// };
+    /// TF_REGISTRY_FUNCTION_WITH_TAG(UsdMayaPrimWriterRegistry, MyWriter) {
+    ///     UsdMayaPrimWriterRegistry::Register("myCustomMayaNode",
+    ///             MyWriter::Create);
+    /// }
+    /// \endcode
+    MAYAUSD_CORE_PUBLIC
+    static void Register(
+        const std::string& mayaType,
+        ContextPredicateFn pred,
+        WriterFactoryFn    fn,
+        bool               fromPython = false);
+
     /// \brief Register \p fn as a factory function providing a
     /// UsdMayaPrimWriter subclass that can be used to write \p mayaType.
     /// If you can't provide a valid UsdMayaPrimWriter for the given arguments,
@@ -109,7 +142,14 @@ struct UsdMayaPrimWriterRegistry
     ///
     /// If there is no writer plugin for \p mayaTypeName, returns nullptr.
     MAYAUSD_CORE_PUBLIC
-    static WriterFactoryFn Find(const std::string& mayaTypeName);
+    static WriterFactoryFn Find(
+        const std::string&          mayaTypeName,
+        const UsdMayaJobExportArgs& exportArgs,
+        const MObject&              exportObj);
+
+    /// \brief Check for external primWriter for \p mayaTypeName.
+    MAYAUSD_CORE_PUBLIC
+    static void CheckForWriterPlugin(const std::string& mayaTypeName);
 
     /// \brief Registers a maya node type to *not* create a new prim.
     ///

lib/mayaUsd/fileio/shaderWriter.h

@@ -43,19 +43,6 @@ class UsdMayaShaderWriter : public UsdMayaPrimWriter
         const SdfPath&           usdPath,
         UsdMayaWriteJobContext&  jobCtx);
 
-    /// The level of support a writer can offer for a given context
-    ///
-    /// A basic writer that gives correct results across most contexts should
-    /// report `Fallback`, while a specialized writer that really shines in a
-    /// given context should report `Supported` when the context is right and
-    /// `Unsupported` if the context is not as expected.
-    enum class ContextSupport
-    {
-        Supported,
-        Fallback,
-        Unsupported
-    };
-
     /// A static function is expected for all shader writers and allows
     /// declaring how well this class can support the current context.
     ///

lib/mayaUsd/fileio/shaderWriterRegistry.cpp

@@ -57,7 +57,7 @@ static int       _indexCounter = 0;
 
 _Registry::const_iterator _Find(const TfToken& usdInfoId, const UsdMayaJobExportArgs& exportArgs)
 {
-    using ContextSupport = UsdMayaShaderWriter::ContextSupport;
+    using ContextSupport = UsdMayaPrimWriter::ContextSupport;
 
     TfToken    conversion = exportArgs.convertMaterialsTo;
     const bool noFallback
@@ -141,7 +141,7 @@ UsdMayaShaderWriterRegistry::WriterFactoryFn UsdMayaShaderWriterRegistry::Find(
         _reg.insert(std::make_pair(
             mayaTypeName,
             _RegistryEntry { [](const UsdMayaJobExportArgs&) {
-                                return UsdMayaShaderWriter::ContextSupport::Fallback;
+                                return UsdMayaPrimWriter::ContextSupport::Fallback;
                             },
                              nullptr,
                              -1 }));

lib/mayaUsd/fileio/shaderWriterRegistry.h

@@ -84,7 +84,7 @@ struct UsdMayaShaderWriterRegistry
     /// Predicate function, i.e. a function that can tell the level of support
     /// the writer function will provide for a given set of export options.
     using ContextPredicateFn
-        = std::function<UsdMayaShaderWriter::ContextSupport(const UsdMayaJobExportArgs&)>;
+        = std::function<UsdMayaPrimWriter::ContextSupport(const UsdMayaJobExportArgs&)>;
 
     /// \brief Register \p fn as a factory function providing a
     /// UsdMayaShaderWriter subclass that can be used to write \p mayaType.

lib/mayaUsd/fileio/shading/symmetricShaderWriter.cpp

@@ -67,7 +67,7 @@ void UsdMayaSymmetricShaderWriter::RegisterWriter(
 }
 
 /* static */
-UsdMayaShaderWriter::ContextSupport UsdMayaSymmetricShaderWriter::CanExport(
+UsdMayaPrimWriter::ContextSupport UsdMayaSymmetricShaderWriter::CanExport(
     const UsdMayaJobExportArgs& exportArgs,
     const TfToken&              materialConversionName)
 {

lib/mayaUsd/fileio/utils/adaptor.cpp

@@ -158,7 +158,7 @@ TfType UsdMayaAdaptor::GetUsdType() const
     // The adaptor type mapping might be registered externally in a prim writer
     // plugin. This simply pokes the prim writer registry to load the prim
     // writer plugin in order to pull in the adaptor mapping.
-    UsdMayaPrimWriterRegistry::Find(depNode.typeName().asChar());
+    UsdMayaPrimWriterRegistry::CheckForWriterPlugin(depNode.typeName().asChar());
     TfRegistryManager::GetInstance().SubscribeTo<UsdMayaAdaptor>();
 
     const auto iter = _schemaLookup.find(depNode.typeName().asChar());

lib/mayaUsd/fileio/writeJobContext.cpp

@@ -537,8 +537,7 @@ UsdMayaPrimWriterSharedPtr UsdMayaWriteJobContext::CreatePrimWriter(
     // This is either a DG node or a non-instanced DAG node, so try to look up
     // a writer plugin. We search through the node's type ancestors, working
     // backwards until we find a prim writer plugin.
-    const std::string mayaTypeName(depNodeFn.typeName().asChar());
-    if (UsdMayaPrimWriterRegistry::WriterFactoryFn primWriterFactory = _FindWriter(mayaTypeName)) {
+    if (UsdMayaPrimWriterRegistry::WriterFactoryFn primWriterFactory = _FindWriter(depNodeFn)) {
         if (UsdMayaPrimWriterSharedPtr primPtr = primWriterFactory(depNodeFn, writePath, *this)) {
             // We found a registered user prim writer that handles this node
             // type, so return now.
@@ -551,8 +550,10 @@ UsdMayaPrimWriterSharedPtr UsdMayaWriteJobContext::CreatePrimWriter(
 }
 
 UsdMayaPrimWriterRegistry::WriterFactoryFn
-UsdMayaWriteJobContext::_FindWriter(const std::string& mayaNodeType)
+UsdMayaWriteJobContext::_FindWriter(const MFnDependencyNode& mayaNode)
 {
+    const std::string mayaNodeType(mayaNode.typeName().asChar());
+
     // Check if type is already cached locally.
     auto iter = mWriterFactoryCache.find(mayaNodeType);
     if (iter != mWriterFactoryCache.end()) {
@@ -564,7 +565,7 @@ UsdMayaWriteJobContext::_FindWriter(const std::string& mayaNodeType)
         = UsdMayaUtil::GetAllAncestorMayaNodeTypes(mayaNodeType);
     for (auto i = ancestorTypes.rbegin(); i != ancestorTypes.rend(); ++i) {
         if (UsdMayaPrimWriterRegistry::WriterFactoryFn primWriterFactory
-            = UsdMayaPrimWriterRegistry::Find(*i)) {
+            = UsdMayaPrimWriterRegistry::Find(*i, mArgs, mayaNode.object())) {
             mWriterFactoryCache[mayaNodeType] = primWriterFactory;
             return primWriterFactory;
         }

lib/mayaUsd/fileio/writeJobContext.h

@@ -186,7 +186,7 @@ class UsdMayaWriteJobContext
         std::vector<UsdMayaPrimWriterSharedPtr>::const_iterator* end) const;
 
     /// Prim writer search with ancestor type resolution behavior.
-    UsdMayaPrimWriterRegistry::WriterFactoryFn _FindWriter(const std::string& mayaNodeType);
+    UsdMayaPrimWriterRegistry::WriterFactoryFn _FindWriter(const MFnDependencyNode& mayaNode);
 
     struct MObjectHandleComp
     {