Solidify addTransitionType Semantics

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -25,6 +25,7 @@ import {
   enableComponentPerformanceTrack,
   enableProfilerTimer,
 } from 'shared/ReactFeatureFlags';
+import {clearEntangledAsyncTransitionTypes} from './ReactFiberTransitionTypes';
 
 // If there are multiple, concurrent async actions, they are entangled. All
 // transition updates that occur while the async action is still in progress
@@ -84,6 +85,7 @@ function pingEngtangledActionScope() {
         clearAsyncTransitionTimer();
       }
     }
+    clearEntangledAsyncTransitionTypes();
     if (currentEntangledListeners !== null) {
       // All the actions have finished. Close the entangled async action scope
       // and notify all the listeners.

packages/react-reconciler/src/ReactFiberHooks.js

@@ -42,6 +42,7 @@ import {
   enableLegacyCache,
   disableLegacyMode,
   enableNoCloningMemoCache,
+  enableViewTransition,
   enableGestureTransition,
 } from 'shared/ReactFeatureFlags';
 import {
@@ -2159,6 +2160,17 @@ function runActionStateAction<S, P>(
     // This is a fork of startTransition
     const prevTransition = ReactSharedInternals.T;
     const currentTransition: Transition = ({}: any);
+    if (enableViewTransition) {
+      currentTransition.types =
+        prevTransition !== null
+          ? // If we're a nested transition, we should use the same set as the parent
+            // since we're conceptually always joined into the same entangled transition.
+            // In practice, this only matters if we add transition types in the inner
+            // without setting state. In that case, the inner transition can finish
+            // without waiting for the outer.
+            prevTransition.types
+          : null;
+    }
     if (enableGestureTransition) {
       currentTransition.gesture = null;
     }
@@ -2180,6 +2192,24 @@ function runActionStateAction<S, P>(
     } catch (error) {
       onActionError(actionQueue, node, error);
     } finally {
+      if (prevTransition !== null && currentTransition.types !== null) {
+        // If we created a new types set in the inner transition, we transfer it to the parent
+        // since they should share the same set. They're conceptually entangled.
+        if (__DEV__) {
+          if (
+            prevTransition.types !== null &&
+            prevTransition.types !== currentTransition.types
+          ) {
+            // Just assert that assumption holds that we're not overriding anything.
+            console.error(
+              'We expected inner Transitions to have transferred the outer types set and ' +
+                'that you cannot add to the outer Transition while inside the inner.' +
+                'This is a bug in React.',
+            );
+          }
+        }
+        prevTransition.types = currentTransition.types;
+      }
       ReactSharedInternals.T = prevTransition;
 
       if (__DEV__) {
@@ -3052,6 +3082,17 @@ function startTransition<S>(
 
   const prevTransition = ReactSharedInternals.T;
   const currentTransition: Transition = ({}: any);
+  if (enableViewTransition) {
+    currentTransition.types =
+      prevTransition !== null
+        ? // If we're a nested transition, we should use the same set as the parent
+          // since we're conceptually always joined into the same entangled transition.
+          // In practice, this only matters if we add transition types in the inner
+          // without setting state. In that case, the inner transition can finish
+          // without waiting for the outer.
+          prevTransition.types
+        : null;
+  }
   if (enableGestureTransition) {
     currentTransition.gesture = null;
   }
@@ -3137,6 +3178,24 @@ function startTransition<S>(
   } finally {
     setCurrentUpdatePriority(previousPriority);
 
+    if (prevTransition !== null && currentTransition.types !== null) {
+      // If we created a new types set in the inner transition, we transfer it to the parent
+      // since they should share the same set. They're conceptually entangled.
+      if (__DEV__) {
+        if (
+          prevTransition.types !== null &&
+          prevTransition.types !== currentTransition.types
+        ) {
+          // Just assert that assumption holds that we're not overriding anything.
+          console.error(
+            'We expected inner Transitions to have transferred the outer types set and ' +
+              'that you cannot add to the outer Transition while inside the inner.' +
+              'This is a bug in React.',
+          );
+        }
+      }
+      prevTransition.types = currentTransition.types;
+    }
     ReactSharedInternals.T = prevTransition;
 
     if (__DEV__) {

packages/react-reconciler/src/ReactFiberRoot.js

@@ -33,6 +33,7 @@ import {
   enableUpdaterTracking,
   enableTransitionTracing,
   disableLegacyMode,
+  enableViewTransition,
   enableGestureTransition,
 } from 'shared/ReactFeatureFlags';
 import {initializeUpdateQueue} from './ReactFiberClassUpdateQueue';
@@ -98,6 +99,10 @@ function FiberRootNode(
 
   this.formState = formState;
 
+  if (enableViewTransition) {
+    this.transitionTypes = null;
+  }
+
   if (enableGestureTransition) {
     this.pendingGestures = null;
     this.stoppingGestures = null;

packages/react-reconciler/src/ReactFiberTransition.js

@@ -12,15 +12,15 @@ import type {
   GestureProvider,
   GestureOptions,
 } from 'shared/ReactTypes';
-import type {Lanes} from './ReactFiberLane';
+import {NoLane, type Lanes} from './ReactFiberLane';
 import type {StackCursor} from './ReactFiberStack';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent';
 import type {Transition} from 'react/src/ReactStartTransition';
 import type {ScheduledGesture} from './ReactFiberGestureScheduler';
-import type {TransitionTypes} from 'react/src/ReactTransitionType';
 
 import {
   enableTransitionTracing,
+  enableViewTransition,
   enableGestureTransition,
 } from 'shared/ReactFeatureFlags';
 import {isPrimaryRenderer} from './ReactFiberConfig';
@@ -34,9 +34,17 @@ import {
   retainCache,
   CacheContext,
 } from './ReactFiberCacheComponent';
+import {
+  queueTransitionTypes,
+  entangleAsyncTransitionTypes,
+  entangledTransitionTypes,
+} from './ReactFiberTransitionTypes';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
-import {entangleAsyncAction} from './ReactFiberAsyncAction';
+import {
+  entangleAsyncAction,
+  peekEntangledActionLane,
+} from './ReactFiberAsyncAction';
 import {startAsyncTransitionTimer} from './ReactProfilerTimer';
 import {firstScheduledRoot} from './ReactFiberRootScheduler';
 import {
@@ -87,6 +95,33 @@ ReactSharedInternals.S = function onStartTransitionFinishForReconciler(
     const thenable: Thenable<mixed> = (returnValue: any);
     entangleAsyncAction(transition, thenable);
   }
+  if (enableViewTransition) {
+    if (entangledTransitionTypes !== null) {
+      // If we scheduled work on any new roots, we need to add any entangled async
+      // transition types to those roots too.
+      let root = firstScheduledRoot;
+      while (root !== null) {
+        queueTransitionTypes(root, entangledTransitionTypes);
+        root = root.next;
+      }
+    }
+    const transitionTypes = transition.types;
+    if (transitionTypes !== null) {
+      // Within this Transition we should've now scheduled any roots we have updates
+      // to work on. If there are no updates on a root, then the Transition type won't
+      // be applied to that root.
+      let root = firstScheduledRoot;
+      while (root !== null) {
+        queueTransitionTypes(root, transitionTypes);
+        root = root.next;
+      }
+      if (peekEntangledActionLane() !== NoLane) {
+        // If we have entangled, async actions going on, the update associated with
+        // these types might come later. We need to save them for later.
+        entangleAsyncTransitionTypes(transitionTypes);
+      }
+    }
+  }
   if (prevOnStartTransitionFinish !== null) {
     prevOnStartTransitionFinish(transition, returnValue);
   }
@@ -113,15 +148,13 @@ if (enableGestureTransition) {
     transition: Transition,
     provider: GestureProvider,
     options: ?GestureOptions,
-    transitionTypes: null | TransitionTypes,
   ): () => void {
     let cancel = null;
     if (prevOnStartGestureTransitionFinish !== null) {
       cancel = prevOnStartGestureTransitionFinish(
         transition,
         provider,
         options,
-        transitionTypes,
       );
     }
     // For every root that has work scheduled, check if there's a ScheduledGesture
@@ -138,7 +171,7 @@ if (enableGestureTransition) {
         root,
         provider,
         options,
-        transitionTypes,
+        transition.types,
       );
       if (scheduledGesture !== null) {
         cancel = chainGestureCancellation(root, scheduledGesture, cancel);

packages/react-reconciler/src/ReactFiberTransitionTypes.js

@@ -0,0 +1,70 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {FiberRoot} from './ReactInternalTypes';
+import type {TransitionTypes} from 'react/src/ReactTransitionType';
+
+import {enableViewTransition} from 'shared/ReactFeatureFlags';
+import {includesTransitionLane} from './ReactFiberLane';
+
+export function queueTransitionTypes(
+  root: FiberRoot,
+  transitionTypes: TransitionTypes,
+): void {
+  if (enableViewTransition) {
+    // TODO: We should really store transitionTypes per lane in a LaneMap on
+    // the root. Then merge it when we commit. We currently assume that all
+    // Transitions are entangled.
+    if (includesTransitionLane(root.pendingLanes)) {
+      let queued = root.transitionTypes;
+      if (queued === null) {
+        queued = root.transitionTypes = [];
+      }
+      for (let i = 0; i < transitionTypes.length; i++) {
+        const transitionType = transitionTypes[i];
+        if (queued.indexOf(transitionType) === -1) {
+          queued.push(transitionType);
+        }
+      }
+    }
+  }
+}
+
+// Store all types while we're entangled with an async Transition.
+export let entangledTransitionTypes: null | TransitionTypes = null;
+
+export function entangleAsyncTransitionTypes(
+  transitionTypes: TransitionTypes,
+): void {
+  if (enableViewTransition) {
+    let queued = entangledTransitionTypes;
+    if (queued === null) {
+      queued = entangledTransitionTypes = [];
+    }
+    for (let i = 0; i < transitionTypes.length; i++) {
+      const transitionType = transitionTypes[i];
+      if (queued.indexOf(transitionType) === -1) {
+        queued.push(transitionType);
+      }
+    }
+  }
+}
+
+export function clearEntangledAsyncTransitionTypes() {
+  // Called when all Async Actions are done.
+  entangledTransitionTypes = null;
+}
+
+export function claimQueuedTransitionTypes(
+  root: FiberRoot,
+): null | TransitionTypes {
+  const claimed = root.transitionTypes;
+  root.transitionTypes = null;
+  return claimed;
+}

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -358,6 +358,7 @@ import {
   deleteScheduledGesture,
   stopCompletedGestures,
 } from './ReactFiberGestureScheduler';
+import {claimQueuedTransitionTypes} from './ReactFiberTransitionTypes';
 
 const PossiblyWeakMap = typeof WeakMap === 'function' ? WeakMap : Map;
 
@@ -3404,11 +3405,7 @@ function commitRoot(
     pendingViewTransitionEvents = null;
     if (includesOnlyViewTransitionEligibleLanes(lanes)) {
       // Claim any pending Transition Types for this commit.
-      // This means that multiple roots committing independent View Transitions
-      // 1) end up staggered because we can only have one at a time.
-      // 2) only the first one gets all the Transition Types.
-      pendingTransitionTypes = ReactSharedInternals.V;
-      ReactSharedInternals.V = null;
+      pendingTransitionTypes = claimQueuedTransitionTypes(root);
       passiveSubtreeMask = PassiveTransitionMask;
     } else {
       pendingTransitionTypes = null;

packages/react-reconciler/src/ReactInternalTypes.js

@@ -18,6 +18,7 @@ import type {
   ReactComponentInfo,
   ReactDebugInfo,
 } from 'shared/ReactTypes';
+import type {TransitionTypes} from 'react/src/ReactTransitionType';
 import type {WorkTag} from './ReactWorkTags';
 import type {TypeOfMode} from './ReactTypeOfMode';
 import type {Flags} from './ReactFiberFlags';
@@ -280,6 +281,8 @@ type BaseFiberRootProperties = {
 
   formState: ReactFormState<any, any> | null,
 
+  // enableViewTransition only
+  transitionTypes: null | TransitionTypes, // TODO: Make this a LaneMap.
   // enableGestureTransition only
   pendingGestures: null | ScheduledGesture,
   stoppingGestures: null | ScheduledGesture,

packages/react/src/ReactSharedInternalsClient.js

@@ -10,20 +10,15 @@
 import type {Dispatcher} from 'react-reconciler/src/ReactInternalTypes';
 import type {AsyncDispatcher} from 'react-reconciler/src/ReactInternalTypes';
 import type {Transition} from './ReactStartTransition';
-import type {TransitionTypes} from './ReactTransitionType';
 import type {GestureProvider, GestureOptions} from 'shared/ReactTypes';
 
-import {
-  enableViewTransition,
-  enableGestureTransition,
-} from 'shared/ReactFeatureFlags';
+import {enableGestureTransition} from 'shared/ReactFeatureFlags';
 
 type onStartTransitionFinish = (Transition, mixed) => void;
 type onStartGestureTransitionFinish = (
   Transition,
   GestureProvider,
   ?GestureOptions,
-  transitionTypes: null | TransitionTypes,
 ) => () => void;
 
 export type SharedStateClient = {
@@ -32,7 +27,6 @@ export type SharedStateClient = {
   T: null | Transition, // ReactCurrentBatchConfig for Transitions
   S: null | onStartTransitionFinish,
   G: null | onStartGestureTransitionFinish,
-  V: null | TransitionTypes, // Pending Transition Types for the Next Transition
 
   // DEV-only
 
@@ -72,9 +66,6 @@ const ReactSharedInternals: SharedStateClient = ({
 if (enableGestureTransition) {
   ReactSharedInternals.G = null;
 }
-if (enableViewTransition) {
-  ReactSharedInternals.V = null;
-}
 
 if (__DEV__) {
   ReactSharedInternals.actQueue = null;