Revisit timing to capture old/new nodes (#56877)

Summary:

## Changelog:

[Internal] [Changed] - Revisit timing to capture old/new nodes

Previously, when `applyViewTransitionName` is called from react reconciler, if there is active viewtransition, we decide it's applying to new node. This was incorrect because `applyViewTransitionName(old)` for next transition can be called when the previous transition is still ongoing. In fact `applyViewTransitionName(old)` is guaranteed to be called before mutationCallback in a transition while `applyViewTransitionName(new)` is inside mutationCallback. So that is a better criteria to separate old and new cases. Given mutationCallback is synchronous on JS thread we can just use a boolean flag to tell the timing.

Given this interleaved situation, when we clean up resources (for new/old nodes) of a transition at its end, we should also avoid removing those for a pending transition. Introducing transitionId for this case.

Differential Revision: D105214322

Author: zeyap

packages/react-native/ReactCommon/react/renderer/viewtransition/ViewTransitionModule.cpp

@@ -86,14 +86,22 @@ void ViewTransitionModule::applyViewTransitionName(
       .size = layoutMetrics.frame.size,
       .pointScaleFactor = layoutMetrics.pointScaleFactor};
 
-  nameRegistry_[tag].insert(name);
-
-  // If applyViewTransitionName is called after transition started, this is the
-  // "new" state (end snapshot). Otherwise, this is the "old" state (start
-  // snapshot)
-  if (!transitionStarted_) {
+  // Calls outside mutationCallback are from the before-mutation phase (old
+  // state for an upcoming transition). Assign a provisional next transition ID
+  // so startViewTransitionEnd cleanup preserves these entries.
+  auto currentTransitionId =
+      insideMutationCallback_ ? activeTransitionId_ : activeTransitionId_ + 1;
+  nameRegistry_[tag].names.insert(name);
+  nameRegistry_[tag].transitionId = currentTransitionId;
+
+  // Old state: called outside mutationCallback (before-mutation phase).
+  // New state: called inside mutationCallback (after-mutation phase).
+  if (!insideMutationCallback_) {
     AnimationKeyFrameView oldView{
-        .layoutMetrics = keyframeMetrics, .tag = tag, .surfaceId = surfaceId};
+        .layoutMetrics = keyframeMetrics,
+        .tag = tag,
+        .surfaceId = surfaceId,
+        .transitionId = currentTransitionId};
     oldLayout_[name] = oldView;
 
     // Request the platform to capture a bitmap snapshot of the old view
@@ -147,7 +155,10 @@ void ViewTransitionModule::applyViewTransitionName(
 
   } else {
     AnimationKeyFrameView newView{
-        .layoutMetrics = keyframeMetrics, .tag = tag, .surfaceId = surfaceId};
+        .layoutMetrics = keyframeMetrics,
+        .tag = tag,
+        .surfaceId = surfaceId,
+        .transitionId = activeTransitionId_};
     newLayout_[name] = newView;
   }
 }
@@ -316,7 +327,7 @@ void ViewTransitionModule::cancelViewTransitionName(
 
 void ViewTransitionModule::restoreViewTransitionName(
     const ShadowNode& shadowNode) {
-  nameRegistry_[shadowNode.getTag()].merge(
+  nameRegistry_[shadowNode.getTag()].names.merge(
       cancelledNameRegistry_[shadowNode.getTag()]);
   cancelledNameRegistry_.erase(shadowNode.getTag());
 }
@@ -387,13 +398,16 @@ void ViewTransitionModule::startViewTransition(
 
   // Mark transition as started
   transitionStarted_ = true;
+  activeTransitionId_ = ++transitionIdCounter_;
   pendingAnimationIds_.clear();
   onCompleteCallback_ = onCompleteCallback;
 
   // Call mutation callback (including commitRoot, measureInstance,
   // applyViewTransitionName, createViewTransitionInstance for old & new)
   if (mutationCallback) {
+    insideMutationCallback_ = true;
     mutationCallback();
+    insideMutationCallback_ = false;
   }
 
   applySnapshotsOnPseudoElementShadowNodes();
@@ -442,13 +456,28 @@ void ViewTransitionModule::suspendOnActiveViewTransition() {
 }
 
 void ViewTransitionModule::startViewTransitionEnd() {
-  for (const auto& [tag, names] : nameRegistry_) {
-    for (const auto& name : names) {
-      oldLayout_.erase(name);
-      newLayout_.erase(name);
+  auto finishedId = activeTransitionId_;
+
+  // Only clear layout and registry entries belonging to the finished
+  // transition. A suspended transition's before-mutation phase may have
+  // already written entries with a newer transitionId — preserve those.
+  for (auto it = nameRegistry_.begin(); it != nameRegistry_.end();) {
+    if (it->second.transitionId == finishedId) {
+      for (const auto& name : it->second.names) {
+        if (auto oit = oldLayout_.find(name);
+            oit != oldLayout_.end() && oit->second.transitionId == finishedId) {
+          oldLayout_.erase(oit);
+        }
+        if (auto nit = newLayout_.find(name);
+            nit != newLayout_.end() && nit->second.transitionId == finishedId) {
+          newLayout_.erase(nit);
+        }
+      }
+      it = nameRegistry_.erase(it);
+    } else {
+      ++it;
     }
   }
-  nameRegistry_.clear();
   oldPseudoElementNodes_.clear();
 
   // Clear any pending bitmap snapshots that were captured but never consumed.

packages/react-native/ReactCommon/react/renderer/viewtransition/ViewTransitionModule.h

@@ -98,17 +98,25 @@ class ViewTransitionModule : public UIManagerViewTransitionDelegate,
     AnimationKeyFrameViewLayoutMetrics layoutMetrics;
     Tag tag{0};
     SurfaceId surfaceId{0};
+    uint32_t transitionId{0};
   };
 
  private:
+  uint32_t transitionIdCounter_{0};
+  uint32_t activeTransitionId_{0};
+
   // registry of layout of old/new views
   std::unordered_map<std::string, AnimationKeyFrameView> oldLayout_{};
   std::unordered_map<std::string, AnimationKeyFrameView> newLayout_{};
-  // tag -> names registry, populated during applyViewTransitionName
+  // tag -> (names, transitionId) registry, populated during applyViewTransitionName
   // Note that tag and name are not 1:1 mapping
   // - In some nested composition 2 names are mappped to the same tag
   // - tags of old and new views are mapped to the same name(s)
-  std::unordered_map<Tag, std::unordered_set<std::string>> nameRegistry_{};
+  struct NameRegistryEntry {
+    std::unordered_set<std::string> names;
+    uint32_t transitionId{0};
+  };
+  std::unordered_map<Tag, NameRegistryEntry> nameRegistry_{};
 
   // used for cancel/restore viewTransitionName
   std::unordered_map<Tag, std::unordered_set<std::string>> cancelledNameRegistry_{};
@@ -138,6 +146,8 @@ class ViewTransitionModule : public UIManagerViewTransitionDelegate,
 
   bool transitionStarted_{false};
 
+  bool insideMutationCallback_{false};
+
   bool transitionReadyFinished_{false};
 
   // When suspendNextTransition_ is true and a transition is active, the next