add UFE contextOps for working set management (loading and unloading)

lib/mayaUsd/ufe/UsdContextOps.cpp

@@ -17,18 +17,25 @@
 
 #include "private/UfeNotifGuard.h"
 
+#include <algorithm>
 #include <cassert>
+#include <utility>
+#include <vector>
 
 #include <maya/MGlobal.h>
 
 #include <ufe/attributes.h>
 #include <ufe/attribute.h>
 #include <ufe/path.h>
 
+#include <pxr/base/tf/diagnostic.h>
+#include <pxr/usd/sdf/path.h>
 #include <pxr/usd/sdf/reference.h>
+#include <pxr/usd/usd/common.h>
+#include <pxr/usd/usd/prim.h>
+#include <pxr/usd/usd/stage.h>
 #include <pxr/usd/usd/references.h>
 #include <pxr/usd/usd/variantSets.h>
-#include <pxr/base/tf/diagnostic.h>
 #include <pxr/usd/usdGeom/tokens.h>
 
 #include <mayaUsd/utils/util.h>
@@ -48,6 +55,12 @@ namespace {
 static constexpr char kUSDLayerEditorItem[] = "USD Layer Editor";
 static constexpr char kUSDLayerEditorLabel[] = "USD Layer Editor...";
 static const std::string kUSDLayerEditorImage{"USD_generic.png"};
+static constexpr char kUSDLoadItem[] = "Load";
+static constexpr char kUSDLoadLabel[] = "Load";
+static constexpr char kUSDLoadWithDescendantsItem[] = "Load with Descendants";
+static constexpr char kUSDLoadWithDescendantsLabel[] = "Load with Descendants";
+static constexpr char kUSDUnloadItem[] = "Unload";
+static constexpr char kUSDUnloadLabel[] = "Unload";
 static constexpr char kUSDVariantSetsItem[] = "Variant Sets";
 static constexpr char kUSDVariantSetsLabel[] = "Variant Sets";
 static constexpr char kUSDToggleVisibilityItem[] = "Toggle Visibility";
@@ -83,6 +96,78 @@ static constexpr char kUSDSpherePrimItem[] = "Sphere";
 static constexpr char kUSDSpherePrimLabel[] = "Sphere";
 static const std::string kUSDSpherePrimImage{"out_USD_Sphere.png"};
 
+//! \brief Undoable command for loading a USD prim.
+class LoadUndoableCommand : public Ufe::UndoableCommand
+{
+public:
+    LoadUndoableCommand(const UsdPrim& prim, UsdLoadPolicy policy) :
+        _stage(prim.GetStage()),
+        _primPath(prim.GetPath()),
+        _oldLoadSet(prim.GetStage()->GetLoadSet()),
+        _policy(policy)
+    {
+    }
+
+    void undo() override
+    {
+        if (!_stage) {
+            return;
+        }
+
+        _stage->LoadAndUnload(_oldLoadSet, SdfPathSet({_primPath}));
+    }
+
+    void redo() override
+    {
+        if (!_stage) {
+            return;
+        }
+
+        _stage->Load(_primPath, _policy);
+    }
+
+private:
+    const UsdStageWeakPtr _stage;
+    const SdfPath _primPath;
+    const SdfPathSet _oldLoadSet;
+    const UsdLoadPolicy _policy;
+};
+
+//! \brief Undoable command for unloading a USD prim.
+class UnloadUndoableCommand : public Ufe::UndoableCommand
+{
+public:
+    UnloadUndoableCommand(const UsdPrim& prim) :
+        _stage(prim.GetStage()),
+        _primPath({prim.GetPath()}),
+        _oldLoadSet(prim.GetStage()->GetLoadSet())
+    {
+    }
+
+    void undo() override
+    {
+        if (!_stage) {
+            return;
+        }
+
+        _stage->LoadAndUnload(_oldLoadSet, SdfPathSet());
+    }
+
+    void redo() override
+    {
+        if (!_stage) {
+            return;
+        }
+
+        _stage->Unload(_primPath);
+    }
+
+private:
+    const UsdStageWeakPtr _stage;
+    const SdfPath _primPath;
+    const SdfPathSet _oldLoadSet;
+};
+
 //! \brief Undoable command for variant selection change
 class SetVariantSelectionUndoableCommand : public Ufe::UndoableCommand
 {
@@ -237,6 +322,79 @@ class ClearAllReferencesUndoableCommand : public Ufe::UndoableCommand
 const std::string ClearAllReferencesUndoableCommand::commandName("Clear All References");
 const MString ClearAllReferencesUndoableCommand::cancelRemoval("No");
 
+std::vector<std::pair<const char* const, const char* const>>
+_computeLoadAndUnloadItems(const UsdPrim& prim)
+{
+    std::vector<std::pair<const char* const, const char* const>> itemLabelPairs;
+
+    const bool isInPrototype =
+#if USD_VERSION_NUM >= 2011
+        prim.IsInPrototype();
+#else
+        prim.IsInMaster();
+#endif
+
+    if (!prim.IsActive() || isInPrototype) {
+        return itemLabelPairs;
+    }
+
+    UsdStageWeakPtr stage = prim.GetStage();
+    const SdfPathSet stageLoadSet = stage->GetLoadSet();
+    const SdfPathSet loadableSet = stage->FindLoadable(prim.GetPath());
+
+    // Intersect the set of what *can* be loaded at or below this prim path
+    // with the set of of what *is* loaded on the stage. The resulting set will
+    // contain all paths that are loaded at or below this prim path.
+    SdfPathSet loadedSet;
+    std::set_intersection(
+        loadableSet.cbegin(), loadableSet.cend(),
+        stageLoadSet.cbegin(), stageLoadSet.cend(),
+        std::inserter(loadedSet, loadedSet.end()));
+
+    // Subtract the set of what *is* loaded on the stage from the set of what
+    // *can* be loaded at or below this prim path. The resulting set will
+    // contain all paths that are loadable, but not currently loaded, at or
+    // below this prim path.
+    SdfPathSet unloadedSet;
+    std::set_difference(
+        loadableSet.cbegin(), loadableSet.cend(),
+        stageLoadSet.cbegin(), stageLoadSet.cend(),
+        std::inserter(unloadedSet, unloadedSet.end()));
+
+    if (!unloadedSet.empty()) {
+        // Loading without descendants is only meaningful for context ops when
+        // the current prim has an unloaded payload.
+        if (prim.HasPayload() && !prim.IsLoaded()) {
+            itemLabelPairs.emplace_back(
+                std::make_pair(kUSDLoadItem, kUSDLoadLabel));
+        }
+
+        // We always add an item for loading with descendants when there are
+        // unloaded paths at or below the current prim, since we may be in one
+        // of the following situations:
+        // - The current prim has a payload that is unloaded, and we don't know
+        //   whether loading it will introduce more payloads in descendants, so
+        //   we offer the choice to also load those or not.
+        // - The current prim has a payload that is loaded, so there must be
+        //   paths below it that are still unloaded.
+        // - The current prim does not have a payload, so there must be paths
+        //   below it that are still unloaded.
+        itemLabelPairs.emplace_back(
+            std::make_pair(
+                kUSDLoadWithDescendantsItem,
+                kUSDLoadWithDescendantsLabel));
+    }
+
+    // If anything is loaded at this prim path or any of its descendants, add
+    // an item for unload.
+    if (!loadedSet.empty()) {
+        itemLabelPairs.emplace_back(
+            std::make_pair(kUSDUnloadItem, kUSDUnloadLabel));
+    }
+
+    return itemLabelPairs;
+}
+
 }
 
 MAYAUSD_NS_DEF {
@@ -297,6 +455,12 @@ Ufe::ContextOps::Items UsdContextOps::getItems(
 
         // Top-level items (do not add for gateway type node):
         if (!fIsAGatewayType) {
+            // Working set management (load and unload):
+            const auto itemLabelPairs = _computeLoadAndUnloadItems(prim());
+            for (const auto& itemLabelPair : itemLabelPairs) {
+                items.emplace_back(itemLabelPair.first, itemLabelPair.second);
+            }
+
             // Variant sets:
             if (prim().HasVariantSets()) {
                 items.emplace_back(
@@ -397,7 +561,19 @@ Ufe::UndoableCommand::Ptr UsdContextOps::doOpCmd(const ItemPath& itemPath)
         return nullptr;
     }
 
-    if (itemPath[0] == kUSDVariantSetsItem) {
+    if (itemPath[0u] == kUSDLoadItem ||
+            itemPath[0u] == kUSDLoadWithDescendantsItem) {
+        const UsdLoadPolicy policy =
+            (itemPath[0u] == kUSDLoadWithDescendantsItem) ?
+                UsdLoadWithDescendants :
+                UsdLoadWithoutDescendants;
+
+        return std::make_shared<LoadUndoableCommand>(prim(), policy);
+    }
+    else if (itemPath[0u] == kUSDUnloadItem) {
+        return std::make_shared<UnloadUndoableCommand>(prim());
+    }
+    else if (itemPath[0] == kUSDVariantSetsItem) {
         // Operation is to set a variant in a variant set.  Need both the
         // variant set and the variant as arguments to the operation.
         if (itemPath.size() != 3u) {

test/lib/ufe/test-samples/ballset/StandaloneScene/Ball.usd

@@ -6,16 +6,68 @@
 
 def Xform "Ball" (
     assetInfo = {
-        asset identifier = @Ball/Ball.usd@
+        asset identifier = @Ball.usd@
         string name = "Ball"
     }
     kind = "component"
-    add references = [
-        @./Ball.shadingVariants.usda@,
-        @./Ball.maya.usd@
-    ]
+    payload = @./Ball_payload.usd@</Ball>
+    variants = {
+        string shadingVariant = "Cue"
+    }
+    add variantSets = ["shadingVariant"]
 )
 {
-}
+    variantSet "shadingVariant" = {
+        "Ball_1" {
+
+        }
+        "Ball_10" {
+
+        }
+        "Ball_11" {
+
+        }
+        "Ball_12" {
+
+        }
+        "Ball_13" {
+
+        }
+        "Ball_13" {
+
+        }
+        "Ball_14" {
+
+        }
+        "Ball_15" {
+
+        }
+        "Ball_2" {
 
+        }
+        "Ball_3" {
 
+        }
+        "Ball_4" {
+
+        }
+        "Ball_5" {
+
+        }
+        "Ball_6" {
+
+        }
+        "Ball_7" {
+
+        }
+        "Ball_8" {
+
+        }
+        "Ball_9" {
+
+        }
+        "Cue" {
+
+        }
+    }
+}

test/lib/ufe/test-samples/ballset/StandaloneScene/Ball_payload.usd

@@ -0,0 +1,18 @@
+#usda 1.0
+(
+    defaultPrim = "Ball"
+    upAxis = "Y"
+)
+
+def "Ball" (
+    assetInfo = {
+        asset identifier = @Ball.usd@
+        string name = "Ball"
+    }
+    add references = [
+        @./Ball.shadingVariants.usda@,
+        @./Ball.maya.usd@
+    ]
+)
+{
+}