Store transition types per root

We store them on all roots that have been scheduled as part of the
startTransition call.

We shouldn't just let the first root that commits steal them all. Each root
gets its own set.

Author: sebmarkbage

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

@@ -20,6 +20,7 @@ import type {ScheduledGesture} from './ReactFiberGestureScheduler';
 
 import {
   enableTransitionTracing,
+  enableViewTransition,
   enableGestureTransition,
 } from 'shared/ReactFeatureFlags';
 import {isPrimaryRenderer} from './ReactFiberConfig';
@@ -87,7 +88,17 @@ ReactSharedInternals.S = function onStartTransitionFinishForReconciler(
     const thenable: Thenable<mixed> = (returnValue: any);
     entangleAsyncAction(transition, thenable);
   }
-  queueTransitionTypes(transition.types);
+  if (enableViewTransition && transition.types !== 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.
+    // TODO: The exception is if we're to an async action, the updates might come in later.
+    let root = firstScheduledRoot;
+    while (root !== null) {
+      queueTransitionTypes(root, transition.types);
+      root = root.next;
+    }
+  }
   if (prevOnStartTransitionFinish !== null) {
     prevOnStartTransitionFinish(transition, returnValue);
   }

packages/react-reconciler/src/ReactFiberTransitionTypes.js

@@ -7,20 +7,22 @@
  * @flow
  */
 
+import type {FiberRoot} from './ReactInternalTypes';
 import type {TransitionTypes} from 'react/src/ReactTransitionType';
 
 import {enableViewTransition} from 'shared/ReactFeatureFlags';
 
-let queuedTransitionTypes: null | TransitionTypes = null;
-
 export function queueTransitionTypes(
+  root: FiberRoot,
   transitionTypes: null | TransitionTypes,
 ): void {
-  // Currently, we assume that all Transitions are batched together into a global single commit.
   if (enableViewTransition && transitionTypes !== null) {
-    let queued = queuedTransitionTypes;
+    // 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.
+    let queued = root.transitionTypes;
     if (queued === null) {
-      queued = queuedTransitionTypes = [];
+      queued = root.transitionTypes = [];
     }
     for (let i = 0; i < transitionTypes.length; i++) {
       const transitionType = transitionTypes[i];
@@ -31,8 +33,10 @@ export function queueTransitionTypes(
   }
 }
 
-export function claimQueuedTransitionTypes(): null | TransitionTypes {
-  const claimed = queuedTransitionTypes;
-  queuedTransitionTypes = null;
+export function claimQueuedTransitionTypes(
+  root: FiberRoot,
+): null | TransitionTypes {
+  const claimed = root.transitionTypes;
+  root.transitionTypes = null;
   return claimed;
 }

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -3405,7 +3405,7 @@ function commitRoot(
     pendingViewTransitionEvents = null;
     if (includesOnlyViewTransitionEligibleLanes(lanes)) {
       // Claim any pending Transition Types for this commit.
-      pendingTransitionTypes = claimQueuedTransitionTypes();
+      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,