Merge pull request #3164 from Autodesk/samuelliu-adsk/MAYA-128422/pri…

…mread-condition

Maya-128422 Override PrimReader

lib/mayaUsd/fileio/jobs/readJob.cpp

@@ -528,7 +528,7 @@ void UsdMaya_ReadJob::_DoImportPrimIt(
 
         TfToken typeName = prim.GetTypeName();
         if (UsdMayaPrimReaderRegistry::ReaderFactoryFn factoryFn
-            = UsdMayaPrimReaderRegistry::FindOrFallback(typeName)) {
+            = UsdMayaPrimReaderRegistry::FindOrFallback(typeName, mArgs, prim)) {
             UsdMayaPrimReaderSharedPtr primReader = factoryFn(args);
             if (primReader) {
                 TempNodeTrackerScope scope(readCtx);

lib/mayaUsd/fileio/primReader.cpp

@@ -22,6 +22,15 @@ UsdMayaPrimReader::UsdMayaPrimReader(const UsdMayaPrimReaderArgs& args)
 {
 }
 
+/* static */
+UsdMayaPrimReader::ContextSupport
+UsdMayaPrimReader::CanImport(const UsdMayaJobImportArgs&, const UsdPrim& importPrim)
+{
+    // Default value for all readers is Fallback. More specialized writers can
+    // override the base CanImport to report Supported/Unsupported as necessary.
+    return ContextSupport::Fallback;
+}
+
 bool UsdMayaPrimReader::HasPostReadSubtree() const { return false; }
 
 void UsdMayaPrimReader::PostReadSubtree(UsdMayaPrimReaderContext&) { }

lib/mayaUsd/fileio/primReader.h

@@ -31,6 +31,25 @@ class UsdMayaPrimReader
     UsdMayaPrimReader(const UsdMayaPrimReaderArgs&);
     virtual ~UsdMayaPrimReader() {};
 
+    /// The level of support a reader can offer for a given context
+    ///
+    /// A basic reader that gives correct results across most contexts should
+    /// report `Fallback`, while a specialized reader 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
+    };
+
+    /// This static function is expected for all prim readers and allows
+    /// declaring how well this class can support the current context:
+    MAYAUSD_CORE_PUBLIC
+    static ContextSupport
+    CanImport(const UsdMayaJobImportArgs& importArgs, const UsdPrim& importPrim);
+
     /// Reads the USD prim given by the prim reader args into a Maya shape,
     /// modifying the prim reader context as a result.
     /// Returns true if successful.

lib/mayaUsd/fileio/primReaderRegistry.cpp

@@ -43,8 +43,42 @@ TF_DEFINE_PRIVATE_TOKENS(
 );
 // clang-format on
 
-typedef std::map<TfToken, UsdMayaPrimReaderRegistry::ReaderFactoryFn> _Registry;
-static _Registry                                                      _reg;
+namespace {
+struct _RegistryEntry
+{
+    UsdMayaPrimReaderRegistry::ContextPredicateFn _pred;
+    UsdMayaPrimReaderRegistry::ReaderFactoryFn    _fn;
+    int                                           _index;
+};
+
+typedef std::unordered_multimap<TfToken, _RegistryEntry, TfToken::HashFunctor> _Registry;
+static _Registry                                                               _reg;
+static int                                                                     _indexCounter = 0;
+
+_Registry::const_iterator
+_Find(const TfToken& typeName, const UsdMayaJobImportArgs& importArgs, const UsdPrim& importPrim)
+{
+    using ContextSupport = UsdMayaPrimReader::ContextSupport;
+
+    _Registry::const_iterator ret = _reg.cend();
+    _Registry::const_iterator first, last;
+
+    std::tie(first, last) = _reg.equal_range(typeName);
+    while (first != last) {
+        ContextSupport support = first->second._pred(importArgs, importPrim);
+        // 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 UsdMayaPrimReaderRegistry::Register(
@@ -53,13 +87,63 @@ void UsdMayaPrimReaderRegistry::Register(
     bool                                       fromPython)
 {
     TfToken tfTypeName(t.GetTypeName());
+    int     index = _indexCounter++;
+    TF_DEBUG(PXRUSDMAYA_REGISTRY)
+        .Msg("Registering UsdMayaPrimReader for TfType %s.\n", tfTypeName.GetText());
+
+    // Use default ContextSupport if not specified
+    _reg.insert(std::make_pair(
+        tfTypeName,
+        _RegistryEntry { [](const UsdMayaJobImportArgs&, const UsdPrim) {
+                            return UsdMayaPrimReader::ContextSupport::Fallback;
+                        },
+                         fn,
+                         index }));
+
+    if (fn) {
+        UsdMaya_RegistryHelper::AddUnloader(
+            [tfTypeName, index]() {
+                _Registry::const_iterator it, itEnd;
+                std::tie(it, itEnd) = _reg.equal_range(tfTypeName);
+                for (; it != itEnd; ++it) {
+                    if (it->second._index == index) {
+                        _reg.erase(it);
+                        break;
+                    }
+                }
+            },
+            fromPython);
+    }
+}
+
+/* static */
+void UsdMayaPrimReaderRegistry::Register(
+    const TfType&                                 t,
+    UsdMayaPrimReaderRegistry::ContextPredicateFn pred,
+    UsdMayaPrimReaderRegistry::ReaderFactoryFn    fn,
+    bool                                          fromPython)
+{
+    TfToken tfTypeName(t.GetTypeName());
+    int     index = _indexCounter++;
     TF_DEBUG(PXRUSDMAYA_REGISTRY)
         .Msg("Registering UsdMayaPrimReader for TfType %s.\n", tfTypeName.GetText());
-    std::pair<_Registry::iterator, bool> insertStatus = _reg.insert(std::make_pair(tfTypeName, fn));
-    if (insertStatus.second) {
-        UsdMaya_RegistryHelper::AddUnloader([tfTypeName]() { _reg.erase(tfTypeName); }, fromPython);
-    } else {
-        TF_CODING_ERROR("Multiple readers for type %s", tfTypeName.GetText());
+    _reg.insert(std::make_pair(tfTypeName, _RegistryEntry { pred, fn, index }));
+
+    // The unloader uses the index to know which entry to erase when there are
+    // more than one for the same usdInfoId.
+    if (fn) {
+        UsdMaya_RegistryHelper::AddUnloader(
+            [tfTypeName, index]() {
+                _Registry::const_iterator it, itEnd;
+                std::tie(it, itEnd) = _reg.equal_range(tfTypeName);
+                for (; it != itEnd; ++it) {
+                    if (it->second._index == index) {
+                        _reg.erase(it);
+                        break;
+                    }
+                }
+            },
+            fromPython);
     }
 }
 
@@ -70,8 +154,10 @@ void UsdMayaPrimReaderRegistry::RegisterRaw(const TfType& t, UsdMayaPrimReaderRe
 }
 
 /* static */
-UsdMayaPrimReaderRegistry::ReaderFactoryFn
-UsdMayaPrimReaderRegistry::Find(const TfToken& usdTypeName)
+UsdMayaPrimReaderRegistry::ReaderFactoryFn UsdMayaPrimReaderRegistry::Find(
+    const TfToken&              usdTypeName,
+    const UsdMayaJobImportArgs& importArgs,
+    const UsdPrim&              importPrim)
 {
     TfRegistryManager::GetInstance().SubscribeTo<UsdMayaPrimReaderRegistry>();
 
@@ -81,28 +167,40 @@ UsdMayaPrimReaderRegistry::Find(const TfToken& usdTypeName)
     std::string     typeNameStr = tfType.GetTypeName();
     TfToken         typeName(typeNameStr);
     ReaderFactoryFn ret = nullptr;
-    if (TfMapLookup(_reg, typeName, &ret)) {
-        return ret;
+
+    // For that we are using unordered_multimap now, we can't use TfMapLookup anymore
+    _Registry::const_iterator it = _Find(typeName, importArgs, importPrim);
+
+    if (it != _reg.end()) {
+        return it->second._fn;
     }
 
     static const TfTokenVector SCOPE = { _tokens->UsdMaya, _tokens->PrimReader };
     UsdMaya_RegistryHelper::FindAndLoadMayaPlug(SCOPE, typeNameStr);
 
     // ideally something just registered itself.  if not, we at least put it in
     // the registry in case we encounter it again.
-    if (!TfMapLookup(_reg, typeName, &ret)) {
+    if (_reg.count(typeName) == 0) {
         TF_DEBUG(PXRUSDMAYA_REGISTRY)
             .Msg("No usdMaya reader plugin for TfType %s. No maya plugin.\n", typeName.GetText());
-        _reg[typeName] = nullptr;
+        Register(
+            tfType,
+            [](const UsdMayaJobImportArgs&, const UsdPrim&) {
+                return UsdMayaPrimReader::ContextSupport::Fallback;
+            },
+            nullptr);
     }
-    return ret;
+
+    return nullptr;
 }
 
 /* static */
-UsdMayaPrimReaderRegistry::ReaderFactoryFn
-UsdMayaPrimReaderRegistry::FindOrFallback(const TfToken& usdTypeName)
+UsdMayaPrimReaderRegistry::ReaderFactoryFn UsdMayaPrimReaderRegistry::FindOrFallback(
+    const TfToken&              usdTypeName,
+    const UsdMayaJobImportArgs& importArgs,
+    const UsdPrim&              importPrim)
 {
-    if (ReaderFactoryFn fn = Find(usdTypeName)) {
+    if (ReaderFactoryFn fn = Find(usdTypeName, importArgs, importPrim)) {
         return fn;
     }
 

lib/mayaUsd/fileio/primReaderRegistry.h

@@ -54,6 +54,11 @@ struct UsdMayaPrimReaderRegistry
     /// for the given prim reader args.
     typedef std::function<UsdMayaPrimReaderSharedPtr(const UsdMayaPrimReaderArgs&)> ReaderFactoryFn;
 
+    /// Predicate function, i.e. a function that can tell the level of support
+    /// the reader function will provide for a given context.
+    using ContextPredicateFn = std::function<
+        UsdMayaPrimReader::ContextSupport(const UsdMayaJobImportArgs&, const UsdPrim&)>;
+
     /// Reader function, i.e. a function that reads a prim. This is the
     /// signature of the function declared in the PXRUSDMAYA_DEFINE_READER
     /// macro.
@@ -63,7 +68,16 @@ struct UsdMayaPrimReaderRegistry
     MAYAUSD_CORE_PUBLIC
     static void Register(const TfType& type, ReaderFactoryFn fn, bool fromPython = false);
 
+    /// \brief Register \p fn as a reader provider for \p type and provide the supportability.
+    MAYAUSD_CORE_PUBLIC
+    static void Register(
+        const TfType&      type,
+        ContextPredicateFn pred,
+        ReaderFactoryFn    fn,
+        bool               fromPython = false);
+
     /// \brief Register \p fn as a reader provider for \p T.
+    /// Context support will be set to default "Fallback"
     ///
     /// Example for registering a reader factory in your custom plugin, assuming
     /// that MyType is registered with the TfType system:
@@ -85,6 +99,30 @@ struct UsdMayaPrimReaderRegistry
         }
     }
 
+    /// \brief Register \p fn as a reader provider for \p T and provide the supportability.
+    /// Use "Supported" to override default reader
+    ///
+    /// Example for registering a reader factory in your custom plugin, assuming
+    /// that MyType is registered with the TfType system:
+    /// \code{.cpp}
+    /// class MyReader : public UsdMayaPrimReader {
+    ///     static UsdMayaPrimReaderSharedPtr Create(
+    ///             const UsdMayaPrimReaderArgs&);
+    /// };
+    /// TF_REGISTRY_FUNCTION_WITH_TAG(UsdMayaPrimReaderRegistry, MyType) {
+    ///     UsdMayaPrimReaderRegistry::Register<MyType>(ContextSupport, MyReader::Create);
+    /// }
+    /// \endcode
+    template <typename T>
+    static void Register(ContextPredicateFn pred, ReaderFactoryFn fn, bool fromPython = false)
+    {
+        if (TfType t = TfType::Find<T>()) {
+            Register(t, pred, fn, fromPython);
+        } else {
+            TF_CODING_ERROR("Cannot register unknown TfType: %s.", ArchGetDemangled<T>().c_str());
+        }
+    }
+
     /// \brief Wraps \p fn in a ReaderFactoryFn and registers that factory
     /// function as a reader provider for \p T.
     /// This is a helper method for the macro PXRUSDMAYA_DEFINE_READER;
@@ -100,13 +138,19 @@ struct UsdMayaPrimReaderRegistry
     /// prim.GetTypeName()
     /// \endcode
     MAYAUSD_CORE_PUBLIC
-    static ReaderFactoryFn Find(const TfToken& usdTypeName);
+    static ReaderFactoryFn Find(
+        const TfToken&              usdTypeName,
+        const UsdMayaJobImportArgs& importArgs,
+        const UsdPrim&              importPrim);
 
     /// Similar to Find(), but returns a "fallback" prim reader factory if none
     /// can be found for \p usdTypeName. Thus, this always returns a valid
     /// reader factory.
     MAYAUSD_CORE_PUBLIC
-    static ReaderFactoryFn FindOrFallback(const TfToken& usdTypeName);
+    static ReaderFactoryFn FindOrFallback(
+        const TfToken&              usdTypeName,
+        const UsdMayaJobImportArgs& importArgs,
+        const UsdPrim&              importPrim);
 };
 
 // Lookup TfType by name instead of static C++ type when

lib/mayaUsd/fileio/primUpdater.cpp

@@ -74,7 +74,12 @@ bool UsdMayaPrimUpdater::canEditAsMaya() const
     // Invalid prim cannot be edited.
     if (!prim)
         return false;
-    return (UsdMayaPrimReaderRegistry::Find(prim.GetTypeName()) != nullptr);
+
+    UsdMayaJobImportArgs jobArgs = UsdMayaJobImportArgs::CreateFromDictionary(
+        _context->GetUserArgs(),
+        /* importWithProxyShapes = */ false,
+        GfInterval::GetFullInterval());
+    return (UsdMayaPrimReaderRegistry::Find(prim.GetTypeName(), jobArgs, prim) != nullptr);
 }
 
 bool UsdMayaPrimUpdater::editAsMaya() { return true; }

lib/mayaUsd/fileio/shaderReader.cpp

@@ -36,7 +36,7 @@ UsdMayaShaderReader::UsdMayaShaderReader(const UsdMayaPrimReaderArgs& readArgs)
 }
 
 /* static */
-UsdMayaShaderReader::ContextSupport UsdMayaShaderReader::CanImport(const UsdMayaJobImportArgs&)
+UsdMayaPrimReader::ContextSupport UsdMayaShaderReader::CanImport(const UsdMayaJobImportArgs&)
 {
     // Default value for all readers is Fallback. More specialized writers can
     // override the base CanImport to report Supported/Unsupported as necessary.

lib/mayaUsd/fileio/shaderReader.h

@@ -41,19 +41,6 @@ class UsdMayaShaderReader : public UsdMayaPrimReader
     MAYAUSD_CORE_PUBLIC
     UsdMayaShaderReader(const UsdMayaPrimReaderArgs&);
 
-    /// The level of support a reader can offer for a given context
-    ///
-    /// A basic reader that gives correct results across most contexts should
-    /// report `Fallback`, while a specialized reader 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
-    };
-
     /// This static function is expected for all shader readers and allows
     /// declaring how well this class can support the current context:
     MAYAUSD_CORE_PUBLIC

lib/mayaUsd/fileio/shaderReaderRegistry.cpp

@@ -57,7 +57,7 @@ static int                                                                     _
 
 _Registry::const_iterator _Find(const TfToken& usdInfoId, const UsdMayaJobImportArgs& importArgs)
 {
-    using ContextSupport = UsdMayaShaderReader::ContextSupport;
+    using ContextSupport = UsdMayaPrimReader::ContextSupport;
 
     _Registry::const_iterator ret = _reg.cend();
     _Registry::const_iterator first, last;
@@ -120,7 +120,7 @@ void UsdMayaShaderReaderRegistry::Register(
 UsdMayaShaderReaderRegistry::ReaderFactoryFn
 UsdMayaShaderReaderRegistry::Find(const TfToken& usdInfoId, const UsdMayaJobImportArgs& importArgs)
 {
-    using ContextSupport = UsdMayaShaderReader::ContextSupport;
+    using ContextSupport = UsdMayaPrimReader::ContextSupport;
     TfRegistryManager::GetInstance().SubscribeTo<UsdMayaShaderReaderRegistry>();
 
     _Registry::const_iterator it = _Find(usdInfoId, importArgs);

lib/mayaUsd/fileio/shaderReaderRegistry.h

@@ -75,7 +75,7 @@ struct UsdMayaShaderReaderRegistry
     /// Predicate function, i.e. a function that can tell the level of support
     /// the reader function will provide for a given context.
     using ContextPredicateFn
-        = std::function<UsdMayaShaderReader::ContextSupport(const UsdMayaJobImportArgs&)>;
+        = std::function<UsdMayaPrimReader::ContextSupport(const UsdMayaJobImportArgs&)>;
 
     /// Reader factory function, i.e. a function that creates a prim reader
     /// for the given prim reader args.

lib/mayaUsd/fileio/shading/symmetricShaderReader.cpp

@@ -74,7 +74,7 @@ void UsdMayaSymmetricShaderReader::RegisterReader(
 }
 
 /* static */
-UsdMayaShaderReader::ContextSupport UsdMayaSymmetricShaderReader::CanImport(
+UsdMayaPrimReader::ContextSupport UsdMayaSymmetricShaderReader::CanImport(
     const UsdMayaJobImportArgs& importArgs,
     const TfToken&              materialConversion)
 {