EMSUSD-1613 clean session layer when undoing add prim

- Refactor the USD undo block and undo manager to support extending an undo item with new undos.
- Add a helper function to cleanup the session layer of a given prim data, capturing the removal in an undo item.
- Make the add-prim command use the session clean-up helper.
- Add unit test.

lib/usdUfe/ufe/UsdUndoAddNewPrimCommand.cpp

@@ -102,6 +102,7 @@ void UsdUndoAddNewPrimCommand::undo()
     UsdUfe::InAddOrDeleteOperation ad;
 
     _undoableItem.undo();
+    removeSessionLeftOvers(_stage, _primPath, &_undoableItem);
 }
 
 void UsdUndoAddNewPrimCommand::redo()

lib/usdUfe/ufe/Utils.cpp

@@ -21,6 +21,7 @@
 #include <usdUfe/ufe/UsdAttributes.h>
 #include <usdUfe/ufe/UsdSceneItem.h>
 #include <usdUfe/ufe/trf/XformOpUtils.h>
+#include <usdUfe/undo/UsdUndoBlock.h>
 #include <usdUfe/utils/editability.h>
 #include <usdUfe/utils/layers.h>
 #include <usdUfe/utils/loadRules.h>
@@ -34,6 +35,7 @@
 #include <pxr/usd/sdf/types.h>
 #include <pxr/usd/sdr/registry.h>
 #include <pxr/usd/sdr/shaderProperty.h>
+#include <pxr/usd/usd/editContext.h>
 #include <pxr/usd/usd/prim.h>
 #include <pxr/usd/usd/primCompositionQuery.h>
 #include <pxr/usd/usd/resolver.h>
@@ -1501,4 +1503,24 @@ bool isSessionLayerGroupMetadata(const std::string& groupName, std::string* adju
     return true;
 }
 
+void removeSessionLeftOvers(
+    const PXR_NS::UsdStageRefPtr& stage,
+    const PXR_NS::SdfPath&        primPath,
+    UsdUndoableItem*              undoableItem,
+    bool                          extraEdits)
+{
+    // Delete any information left in the session layer, adding any action taken
+    // to the undoable items. Note that if an undo/redo cycle already happened,
+    // the removal of the session data will already been done by the previous
+    // undo since this first undo captured removing the session data. In that
+    // case, the code below will do nothing and we won't capture double-removal
+    // of session data.
+    if (!stage)
+        return;
+
+    UsdEditContext editContext(stage, stage->GetSessionLayer());
+    UsdUndoBlock   undoBlock(undoableItem, extraEdits);
+    stage->RemovePrim(primPath);
+}
+
 } // namespace USDUFE_NS_DEF

lib/usdUfe/ufe/Utils.h

@@ -50,6 +50,7 @@ UFE_NS_DEF
 namespace USDUFE_NS_DEF {
 
 class UsdAttribute;
+class UsdUndoableItem;
 
 // DCC specific accessor functions.
 typedef PXR_NS::UsdStageWeakPtr (*StageAccessorFn)(const Ufe::Path&);
@@ -467,6 +468,15 @@ const char* getTransform3dMatrixOpName();
 USDUFE_PUBLIC
 bool isSessionLayerGroupMetadata(const std::string& groupName, std::string* adjustedGroupName);
 
+//! Remove data left behind in the session layer for the given prim path in the given stage
+//! and store the undos as extra undos in the given undo items.
+USDUFE_PUBLIC
+void removeSessionLeftOvers(
+    const PXR_NS::UsdStageRefPtr& stage,
+    const PXR_NS::SdfPath&        primPath,
+    UsdUndoableItem*              undoableItem,
+    bool                          extraEdits = true);
+
 } // namespace USDUFE_NS_DEF
 
 #endif // USDUFE_UFE_UTILS_H

lib/usdUfe/undo/UsdUndoBlock.cpp

@@ -23,8 +23,9 @@ USDUFE_VERIFY_CLASS_NOT_MOVE_OR_COPY(UsdUndoBlock);
 
 uint32_t UsdUndoBlock::_undoBlockDepth { 0 };
 
-UsdUndoBlock::UsdUndoBlock(UsdUndoableItem* undoItem)
+UsdUndoBlock::UsdUndoBlock(UsdUndoableItem* undoItem, bool extraEdits)
     : _undoItem(undoItem)
+    , _extraEdits(extraEdits)
 {
     // TfDebug::Enable(USDUFE_UNDOSTACK);
 
@@ -40,7 +41,7 @@ UsdUndoBlock::~UsdUndoBlock()
 
     if ((nullptr != _undoItem) && (_undoBlockDepth == 0)) {
         // transfer edits
-        UsdUfe::UsdUndoManagerAccessor::transferEdits(*_undoItem);
+        UsdUfe::UsdUndoManagerAccessor::transferEdits(*_undoItem, _extraEdits);
 
         TF_DEBUG_MSG(USDUFE_UNDOSTACK, "Undoable Item adopted the new edits.\n");
     }

lib/usdUfe/undo/UsdUndoBlock.h

@@ -35,7 +35,12 @@ TF_DECLARE_WEAK_AND_REF_PTRS(UsdUndoManager);
 class USDUFE_PUBLIC UsdUndoBlock
 {
 public:
-    UsdUndoBlock(UsdUndoableItem* undoItem);
+    /// @brief Create an undo block that will capture all undo into the given undo item.
+    /// @param undoItem the item to receive the undos.
+    /// @param extraEdits if true, the undos are added the item, even if the item already contained
+    /// undos.
+    ///                   Otherwise, any undos that were already in the items are discarded.
+    UsdUndoBlock(UsdUndoableItem* undoItem, bool extraEdits = false);
     virtual ~UsdUndoBlock();
 
     USDUFE_DISALLOW_COPY_MOVE_AND_ASSIGNMENT(UsdUndoBlock);
@@ -46,6 +51,7 @@ class USDUFE_PUBLIC UsdUndoBlock
     static uint32_t _undoBlockDepth;
 
     UsdUndoableItem* _undoItem;
+    bool             _extraEdits;
 };
 
 } // namespace USDUFE_NS_DEF

lib/usdUfe/undo/UsdUndoManager.cpp

@@ -53,10 +53,16 @@ void UsdUndoManager::addInverse(UsdUndoableItem::InvertFunc func)
     _invertFuncs.emplace_back(func);
 }
 
-void UsdUndoManager::transferEdits(UsdUndoableItem& undoableItem)
+void UsdUndoManager::transferEdits(UsdUndoableItem& undoableItem, bool extraEdits)
 {
     // transfer the edits
-    undoableItem._invertFuncs = std::move(_invertFuncs);
+    if (extraEdits) {
+        undoableItem._invertFuncs.insert(
+            undoableItem._invertFuncs.begin(), _invertFuncs.begin(), _invertFuncs.end());
+        _invertFuncs.clear();
+    } else {
+        undoableItem._invertFuncs = std::move(_invertFuncs);
+    }
 }
 
 } // namespace USDUFE_NS_DEF

lib/usdUfe/undo/UsdUndoManager.h

@@ -54,7 +54,7 @@ class USDUFE_PUBLIC UsdUndoManager
     ~UsdUndoManager() = default;
 
     void addInverse(UsdUndoableItem::InvertFunc func);
-    void transferEdits(UsdUndoableItem& undoableItem);
+    void transferEdits(UsdUndoableItem& undoableItem, bool extraEdits);
 
 private:
     UsdUndoableItem::InvertFuncs _invertFuncs;
@@ -76,10 +76,10 @@ class USDUFE_PUBLIC UsdUndoManagerAccessor
         auto& undoManager = UsdUfe::UsdUndoManager::instance();
         undoManager.addInverse(func);
     }
-    static void transferEdits(UsdUndoableItem& undoableItem)
+    static void transferEdits(UsdUndoableItem& undoableItem, bool extraEdits = false)
     {
         auto& undoManager = UsdUfe::UsdUndoManager::instance();
-        undoManager.transferEdits(undoableItem);
+        undoManager.transferEdits(undoableItem, extraEdits);
     }
 };
 

test/lib/ufe/testContextOps.py

@@ -748,6 +748,50 @@ def testAddNewPrim(self):
         self.assertEqual(ufeObs.nbAddNotif(), 2)
         self.assertEqual(ufeObs.nbDeleteNotif(), 2)
 
+    def testUndoAddNewPrimCleanSessionLayer(self):
+        cmds.file(new=True, force=True)
+
+        # Create a proxy shape with empty stage to start with.
+        proxyShape = mayaUsd_createStageWithNewLayer.createStageWithNewLayer()
+        stage = mayaUsd.lib.GetPrim(proxyShape).GetStage()
+
+        # Create a ContextOps interface for the proxy shape.
+        proxyShapePath = ufe.Path([mayaUtils.createUfePathSegment(proxyShape)])
+        proxyShapeItem = ufe.Hierarchy.createItem(proxyShapePath)
+        contextOps = ufe.ContextOps.contextOps(proxyShapeItem)
+
+        # Add a new prim.
+        cmd = contextOps.doOpCmd(['Add New Prim', 'Xform'])
+        self.assertIsNotNone(cmd)
+        ufeCmd.execute(cmd)
+
+        # The proxy shape should now have a single UFE child item.
+        proxyShapehier = ufe.Hierarchy.hierarchy(proxyShapeItem)
+        self.assertTrue(proxyShapehier.hasChildren())
+        self.assertEqual(len(proxyShapehier.children()), 1)
+
+        # Add a new prim to the prim we just added.
+        cmds.pickWalk(d='down')
+
+        # Get the scene item from the UFE selection.
+        snIter = iter(ufe.GlobalSelection.get())
+        xformItem = next(snIter)
+        xformPrim = usdUtils.getPrimFromSceneItem(xformItem)
+        xformPath = xformPrim.GetPath()
+
+        # Add data in the session layer.
+        metadataName = 'instanceable'
+        sessionLayer = stage.GetSessionLayer()
+        with Usd.EditContext(stage, sessionLayer):
+            xformPrim.SetMetadata(metadataName, True)
+
+        self.assertTrue(xformPrim.HasAuthoredMetadata(metadataName))
+        self.assertTrue(sessionLayer.GetPrimAtPath(xformPath))
+
+        # Verify that after undo the sessin layer got cleaned.
+        cmd.undo()
+        self.assertFalse(sessionLayer.GetPrimAtPath(xformPath))
+
     def testAddNewPrimInWeakerLayer(self):
         cmds.file(new=True, force=True)