Solidify addTransitionType Semantics (facebook#32797)

Stacked on facebook#32793.

This is meant to model the intended semantics of `addTransitionType`
better. The previous hack just consumed all transition types when any
root committed so it could steal them from other roots. Really each root
should get its own set. Really each transition lane should get its own
set.

We can't implement the full ideal semantics yet because 1) we currently
entangle transition lanes 2) we lack `AsyncContext` on the client so for
async actions we can't associate a `addTransitionType` call to a
specific `startTransition`.

This starts by modeling Transition Types to be stored on the Transition
instance. Conceptually they belong to the Transition instance of that
`startTransition` they belong to. That instance is otherwise mostly just
used for Transition Tracing but it makes sense that those would be able
to be passed the Transition Types for that specific instance.

Nested `startTransition` need to get entangled. So that this
`addTransitionType` can be associated with the `setState`:

```js
startTransition(() => {
  startTransition(() => {
    addTransitionType(...)
  });
  setState(...);
});
```

Ideally we'd probably just use the same Transition instance itself since
these are conceptually all part of one entangled one. But transition
tracing uses multiple names and start times. Unclear what we want to do
with that. So I kept separate instances but shared `types` set.

Next I collect the types added during a `startTransition` to any root
scheduled with a Transition. This should really be collected one set per
Transition lane in a `LaneMap`. In fact, the information would already
be there if Transition Tracing was always enabled because it tracks all
Transition instances per lane. For now I just keep track of one set for
all Transition lanes. Maybe we should only add it if a `setState` was
done on this root in this particular `startTransition` call rather
having already scheduled any Transition earlier.

While async transitions are entangled, we don't know if there will be a
startTransition+setState on a new root in the future. Therefore, we
collect all transition types while this is happening and if a new root
gets startTransition+setState they get added to that root.

```js
startTransition(async () => {
  addTransitionType(...)
  await ...;
  setState(...);
});
```

Author: sebmarkbage

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;