MAYA-123720 - As a user, I'd like Display Layer membership to survive the 'Edit As Maya Data' workflow

lib/mayaUsd/fileio/primUpdaterContext.h

@@ -19,6 +19,7 @@
 
 #include <mayaUsd/base/api.h>
 #include <mayaUsd/fileio/primUpdaterArgs.h>
+#include <mayaUsd/ufe/Utils.h>
 
 #include <pxr/base/tf/hashmap.h>
 #include <pxr/pxr.h>
@@ -67,6 +68,9 @@ class UsdMayaPrimUpdaterContext
     MAYAUSD_CORE_PUBLIC
     MDagPath MapSdfPathToDagPath(const SdfPath& sdfPath) const;
 
+    mutable MayaUsd::ufe::ReplicateExtrasFromUSD _pullExtras;
+    mutable MayaUsd::ufe::ReplicateExtrasToUSD   _pushExtras;
+
 private:
     const UsdTimeCode&           _timeCode;
     const UsdStageRefPtr         _stage;

lib/mayaUsd/fileio/primUpdaterManager.cpp

@@ -502,6 +502,8 @@ PullImportPaths pullImport(
         Ufe::Path::Segments s { ps, Ufe::PathSegment(v.first, rtid, '/') };
         Ufe::Path           p(std::move(s));
         objToUfePath.insert(ObjToUfePath::value_type(MObjectHandle(v.second), p));
+
+        context._pullExtras.processItem(p, v.second);
     }
     progressBar.advance();
 
@@ -631,6 +633,7 @@ PushCustomizeSrc pushExport(
     auto usdPathToDagPathMap = std::make_shared<UsdPathToDagPathMap>();
     for (const auto& v : writeJob.GetDagPathToUsdPathMap()) {
         usdPathToDagPathMap->insert(UsdPathToDagPathMap::value_type(v.second, v.first));
+        context._pushExtras.processItem(v.first, v.second);
     }
 
     std::get<UsdPathToDagPathMapPtr>(pushCustomizeSrc) = usdPathToDagPathMap;
@@ -883,7 +886,7 @@ bool PrimUpdaterManager::mergeToUsd(
     MString progStr(
         VtDictionaryIsHolding<std::string>(userArgs, "rn_primName") ? "Caching to USD"
                                                                     : "Merging to USD");
-    MayaUsd::ProgressBarScope progressBar(9, progStr);
+    MayaUsd::ProgressBarScope progressBar(10, progStr);
     PushPullScope             scopeIt(_inPushPull);
 
     auto ctxArgs = VtDictionaryOver(userArgs, UsdMayaJobExportArgs::GetDefaultDictionary());
@@ -1020,6 +1023,9 @@ bool PrimUpdaterManager::mergeToUsd(
     }
     progressBar.advance();
 
+    context._pushExtras.finalize(MayaUsd::ufe::stagePath(context.GetUsdStage()));
+    progressBar.advance();
+
     discardPullSetIfEmpty();
 
     // Some updaters (like MayaReference) may be writing and changing the variant during merge.
@@ -1070,6 +1076,7 @@ bool PrimUpdaterManager::editAsMaya(const Ufe::Path& path, const VtDictionary& u
 
     auto& scene = Ufe::Scene::instance();
     auto  ufeItem = Ufe::Hierarchy::createItem(path);
+    context._pullExtras.initRecursive(ufeItem);
     if (!updaterArgs._copyOperation && TF_VERIFY(ufeItem))
         scene.notify(Ufe::ObjectPreDelete(ufeItem));
 
@@ -1363,6 +1370,7 @@ bool PrimUpdaterManager::duplicate(
 
         UsdMayaPrimUpdaterContext context(
             srcProxyShape->getTime(), srcProxyShape->getUsdStage(), ctxArgs);
+        context._pullExtras.initRecursive(Ufe::Hierarchy::createItem(srcPath));
         progressBar.advance();
 
         pullImport(srcPath, srcPrim, context);
@@ -1377,7 +1385,7 @@ bool PrimUpdaterManager::duplicate(
             return false;
         }
 
-        MayaUsd::ProgressBarScope progressBar(6, "Duplicating to USD");
+        MayaUsd::ProgressBarScope progressBar(7, "Duplicating to USD");
 
         auto ctxArgs = VtDictionaryOver(userArgs, UsdMayaJobExportArgs::GetDefaultDictionary());
 
@@ -1423,6 +1431,10 @@ bool PrimUpdaterManager::duplicate(
         }
         progressBar.advance();
 
+        std::string* rootRename = (dstChildName != srcRootPath.GetName()) ? &dstChildName : nullptr;
+        context._pushExtras.finalize(MayaUsd::ufe::stagePath(context.GetUsdStage()), rootRename);
+        progressBar.advance();
+
         auto ufeItem = Ufe::Hierarchy::createItem(dstPath);
         if (TF_VERIFY(ufeItem)) {
             Ufe::Scene::instance().notify(Ufe::SubtreeInvalidate(ufeItem));

lib/mayaUsd/render/vp2RenderDelegate/proxyRenderDelegate.cpp

@@ -818,16 +818,20 @@ void ProxyRenderDelegate::_DirtyUsdSubtree(const UsdPrim& prim)
         | HdChangeTracker::DirtyDisplayStyle | MayaUsdRPrim::DirtySelectionHighlight
         | HdChangeTracker::DirtyMaterialId;
 
-    if (prim.IsA<UsdGeomGprim>()) {
+    if (prim.IsA<UsdGeomGprim>() && prim.IsActive()) {
         auto indexPath = _sceneDelegate->ConvertCachePathToIndexPath(prim.GetPath());
-        changeTracker.MarkRprimDirty(indexPath, dirtyBits);
+        if (_renderIndex->HasRprim(indexPath)) {
+            changeTracker.MarkRprimDirty(indexPath, dirtyBits);
+        }
     }
 
     UsdPrimSubtreeRange range = prim.GetDescendants();
     for (auto iter = range.begin(); iter != range.end(); ++iter) {
         if (iter->IsA<UsdGeomGprim>()) {
             auto indexPath = _sceneDelegate->ConvertCachePathToIndexPath(iter->GetPath());
-            changeTracker.MarkRprimDirty(indexPath, dirtyBits);
+            if (_renderIndex->HasRprim(indexPath)) {
+                changeTracker.MarkRprimDirty(indexPath, dirtyBits);
+            }
         }
     }
 }

lib/mayaUsd/ufe/Utils.cpp

@@ -37,6 +37,11 @@
 #include <pxr/usd/usdGeom/tokens.h>
 #include <pxr/usdImaging/usdImaging/delegate.h>
 
+#ifdef MAYA_HAS_DISPLAY_LAYER_API
+#include <maya/MFnDisplayLayer.h>
+#include <maya/MFnDisplayLayerManager.h>
+#endif
+
 #include <maya/MFnDependencyNode.h>
 #include <maya/MGlobal.h>
 #include <maya/MObjectHandle.h>
@@ -895,5 +900,96 @@ std::vector<std::string> splitString(const std::string& str, const std::string&
     return split;
 }
 
+void ReplicateExtrasFromUSD::initRecursive(Ufe::SceneItem::Ptr ufeItem)
+{
+    auto node = Ufe::Hierarchy::hierarchy(ufeItem);
+    if (!node) {
+        return;
+    }
+
+    // Go through the entire hierarchy
+    for (auto child : node->children()) {
+        initRecursive(child);
+    }
+
+#ifdef MAYA_HAS_DISPLAY_LAYER_API
+    // Prepare _displayLayerMap
+    MFnDisplayLayerManager displayLayerManager(
+        MFnDisplayLayerManager::currentDisplayLayerManager());
+
+    MObject displayLayerObj
+        = displayLayerManager.getLayer(Ufe::PathString::string(ufeItem->path()).c_str());
+    if (displayLayerObj.hasFn(MFn::kDisplayLayer)) {
+        MFnDisplayLayer displayLayer(displayLayerObj);
+        if (displayLayer.name() != "defaultLayer") {
+            _displayLayerMap[ufeItem->path()] = displayLayerObj;
+        }
+    }
+#endif
+}
+
+void ReplicateExtrasFromUSD::processItem(const Ufe::Path& path, const MObject& mayaObject)
+{
+#ifdef MAYA_HAS_DISPLAY_LAYER_API
+    // Replicate display layer membership
+    auto it = _displayLayerMap.find(path);
+    if (it != _displayLayerMap.end() && it->second.hasFn(MFn::kDisplayLayer)) {
+        MDagPath dagPath;
+        if (MDagPath::getAPathTo(mayaObject, dagPath) == MStatus::kSuccess) {
+            MFnDisplayLayer displayLayer(it->second);
+            displayLayer.add(dagPath);
+
+            // In case display layer membership was removed from the USD prim that we are
+            // replicating, we want to restore it here to make sure that the prim will stay in its
+            // display layer on DiscardEdits
+            displayLayer.add(Ufe::PathString::string(path).c_str());
+        }
+    }
+#endif
+}
+
+void ReplicateExtrasToUSD::processItem(const MDagPath& dagPath, const SdfPath& usdPath)
+{
+#ifdef MAYA_HAS_DISPLAY_LAYER_API
+    // Populate display layer membership map
+
+    // Since multiple dag paths may lead to a single USD path (like transform and node),
+    // we have to make sure we don't overwrite a non-default layer with a default one
+    bool displayLayerAssigned = false;
+    auto entry = _primToLayerMap.find(usdPath);
+    if (entry != _primToLayerMap.end() && entry->second.hasFn(MFn::kDisplayLayer)) {
+        MFnDisplayLayer displayLayer(entry->second);
+        displayLayerAssigned = (displayLayer.name() != "defaultLayer");
+    }
+
+    if (!displayLayerAssigned) {
+        MFnDisplayLayerManager displayLayerManager(
+            MFnDisplayLayerManager::currentDisplayLayerManager());
+        _primToLayerMap[usdPath] = displayLayerManager.getLayer(dagPath);
+    }
+#endif
+}
+
+void ReplicateExtrasToUSD::finalize(const Ufe::Path& stagePath, const std::string* renameRoot)
+{
+#ifdef MAYA_HAS_DISPLAY_LAYER_API
+    // Replicate display layer membership
+    for (const auto& entry : _primToLayerMap) {
+        if (entry.second.hasFn(MFn::kDisplayLayer)) {
+            auto rootPath = MayaUsd::ufe::usdPathToUfePathSegment(entry.first);
+            if (renameRoot && !rootPath.empty()) {
+                *rootPath.begin() = Ufe::PathComponent(*renameRoot);
+            }
+
+            Ufe::Path::Segments segments { stagePath.getSegments()[0], rootPath };
+            Ufe::Path           ufePath(std::move(segments));
+
+            MFnDisplayLayer displayLayer(entry.second);
+            displayLayer.add(Ufe::PathString::string(ufePath).c_str());
+        }
+    }
+#endif
+}
+
 } // namespace ufe
 } // namespace MAYAUSD_NS_DEF

lib/mayaUsd/ufe/Utils.h

@@ -231,5 +231,31 @@ Ufe::Selection recreateDescendants(const Ufe::Selection& src, const Ufe::Path& f
 MAYAUSD_CORE_PUBLIC
 std::vector<std::string> splitString(const std::string& str, const std::string& separators);
 
+class ReplicateExtrasFromUSD
+{
+public:
+    // Prepares the replication operation for the subtree starting with the given scene item
+    void initRecursive(Ufe::SceneItem::Ptr);
+    // Replicates extra features from the USD item defined by 'path' to the maya object
+    void processItem(const Ufe::Path& path, const MObject& mayaObject);
+
+private:
+    std::unordered_map<Ufe::Path, MObject> _displayLayerMap;
+};
+
+class ReplicateExtrasToUSD
+{
+public:
+    // Processes replication from a maya object defined by 'dagPath'
+    // to the usd item defined by 'usdPath'
+    void processItem(const MDagPath& dagPath, const PXR_NS::SdfPath& usdPath);
+    // Finalizes the replication operation to the USD stage defined by 'stagePath'
+    // with a possibility to rename the usd root node name to 'renameRoot'
+    void finalize(const Ufe::Path& stagePath, const std::string* renameRoot = nullptr);
+
+private:
+    std::map<PXR_NS::SdfPath, MObject> _primToLayerMap;
+};
+
 } // namespace ufe
 } // namespace MAYAUSD_NS_DEF

test/lib/ufe/testDisplayLayer.py

@@ -31,6 +31,8 @@
 import mayaUsd.ufe as mayaUsdUfe
 import mayaUsd.lib as mayaUsdLib
 
+from usdUtils import createSimpleXformScene
+
 from pxr import Usd, Kind
 
 import unittest
@@ -343,6 +345,40 @@ def testDisplayLayerClear(self):
         self.assertFalse(layer1.contains(self.CUBE1))
         self.assertFalse(layer1.contains(self.INVALID_PRIM))
 
+    def testDisplayLayerEditAsMaya(self):
+        '''Display layer membership in Edit As Maya workflow.'''
+
+        (ps, xlateOp, xlation, aUsdUfePathStr, aUsdUfePath, aUsdItem,
+         _, _, _, _, _) = createSimpleXformScene()
+
+        # Add an item to a new display layer
+        cmds.createDisplayLayer(name='layer1', number=1, empty=True)
+        cmds.editDisplayLayerMembers('layer1', '|stage1|stageShape1,/A', noRecurse=True)
+
+        # Edit aPrim as Maya data.
+        with mayaUsd.lib.OpUndoItemList():
+            self.assertTrue(mayaUsd.lib.PrimUpdaterManager.canEditAsMaya(aUsdUfePathStr))
+            self.assertTrue(mayaUsd.lib.PrimUpdaterManager.editAsMaya(aUsdUfePathStr))
+
+        # Check display layer membership on the other side
+        layerObjs = cmds.editDisplayLayerMembers('layer1', query=True, **self.kwArgsEditDisplayLayerMembers)
+        self.assertTrue("|__mayaUsd__|AParent|A" in layerObjs)
+
+        # Create a new layer and put the item there
+        cmds.createDisplayLayer(name='layer2', number=1, empty=True)
+        cmds.editDisplayLayerMembers('layer2', "|__mayaUsd__|AParent|A", noRecurse=True)
+
+        # Merge edits back to USD.
+        aMayaItem = ufe.GlobalSelection.get().front()
+        aMayaPath = aMayaItem.path()
+        aMayaPathStr = ufe.PathString.string(aMayaPath)
+        with mayaUsd.lib.OpUndoItemList():
+            self.assertTrue(mayaUsd.lib.PrimUpdaterManager.mergeToUsd(aMayaPathStr))
+
+        # Check display layer membership back on the USD side
+        layerObjs = cmds.editDisplayLayerMembers('layer2', query=True, **self.kwArgsEditDisplayLayerMembers)
+        self.assertTrue('|stage1|stageShape1,/A' in layerObjs)
+
     @unittest.skipUnless(mayaUtils.ufeSupportFixLevel() >= 3, "Requires Display Layer Ufe subtree invalidate fix.")
     def testDisplayLayerSubtreeInvalidate(self):
         # Create a scene with two variants.