[Transition Tracing] Add transition to OffscreenState and pendingSuspenseBoundaries to RootState

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -236,7 +236,6 @@ import {
   markSkippedUpdateLanes,
   getWorkInProgressRoot,
   pushRenderLanes,
-  getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.new';
 import {setWorkInProgressVersion} from './ReactMutableSource.new';
 import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent.new';
@@ -652,9 +651,11 @@ function updateOffscreenComponent(
     // Rendering a hidden tree.
     if ((workInProgress.mode & ConcurrentMode) === NoMode) {
       // In legacy sync mode, don't defer the subtree. Render it now.
+      // TODO: Consider how Offscreen should work with transitions in the future
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       if (enableCache) {
@@ -688,6 +689,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: nextBaseLanes,
         cachePool: spawnedCachePool,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       workInProgress.updateQueue = null;
@@ -723,6 +725,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       // Push the lanes that were skipped when we bailed out.
@@ -1343,13 +1346,6 @@ function updateHostRoot(current, workInProgress, renderLanes) {
     }
   }
 
-  if (enableTransitionTracing) {
-    // FIXME: Slipped past code review. This is not a safe mutation:
-    // workInProgress.memoizedState is a shared object. Need to fix before
-    // rolling out the Transition Tracing experiment.
-    workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-  }
-
   // Caution: React DevTools currently depends on this property
   // being called "element".
   const nextChildren = nextState.element;
@@ -1363,6 +1359,7 @@ function updateHostRoot(current, workInProgress, renderLanes) {
       element: nextChildren,
       isDehydrated: false,
       cache: nextState.cache,
+      pendingSuspenseBoundaries: nextState.pendingSuspenseBoundaries,
       transitions: nextState.transitions,
     };
     const updateQueue: UpdateQueue<RootState> = (workInProgress.updateQueue: any);
@@ -1980,6 +1977,7 @@ function mountSuspenseOffscreenState(renderLanes: Lanes): OffscreenState {
   return {
     baseLanes: renderLanes,
     cachePool: getSuspendedCache(),
+    transitions: null,
   };
 }
 
@@ -2014,6 +2012,7 @@ function updateSuspenseOffscreenState(
   return {
     baseLanes: mergeLanes(prevOffscreenState.baseLanes, renderLanes),
     cachePool,
+    transitions: prevOffscreenState.transitions,
   };
 }
 
@@ -3580,9 +3579,6 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         const cache: Cache = current.memoizedState.cache;
         pushCacheProvider(workInProgress, cache);
       }
-      if (enableTransitionTracing) {
-        workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-      }
       resetHydrationState();
       break;
     case HostComponent:

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -236,7 +236,6 @@ import {
   markSkippedUpdateLanes,
   getWorkInProgressRoot,
   pushRenderLanes,
-  getWorkInProgressTransitions,
 } from './ReactFiberWorkLoop.old';
 import {setWorkInProgressVersion} from './ReactMutableSource.old';
 import {pushCacheProvider, CacheContext} from './ReactFiberCacheComponent.old';
@@ -652,9 +651,11 @@ function updateOffscreenComponent(
     // Rendering a hidden tree.
     if ((workInProgress.mode & ConcurrentMode) === NoMode) {
       // In legacy sync mode, don't defer the subtree. Render it now.
+      // TODO: Consider how Offscreen should work with transitions in the future
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       if (enableCache) {
@@ -688,6 +689,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: nextBaseLanes,
         cachePool: spawnedCachePool,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       workInProgress.updateQueue = null;
@@ -723,6 +725,7 @@ function updateOffscreenComponent(
       const nextState: OffscreenState = {
         baseLanes: NoLanes,
         cachePool: null,
+        transitions: null,
       };
       workInProgress.memoizedState = nextState;
       // Push the lanes that were skipped when we bailed out.
@@ -1343,13 +1346,6 @@ function updateHostRoot(current, workInProgress, renderLanes) {
     }
   }
 
-  if (enableTransitionTracing) {
-    // FIXME: Slipped past code review. This is not a safe mutation:
-    // workInProgress.memoizedState is a shared object. Need to fix before
-    // rolling out the Transition Tracing experiment.
-    workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-  }
-
   // Caution: React DevTools currently depends on this property
   // being called "element".
   const nextChildren = nextState.element;
@@ -1363,6 +1359,7 @@ function updateHostRoot(current, workInProgress, renderLanes) {
       element: nextChildren,
       isDehydrated: false,
       cache: nextState.cache,
+      pendingSuspenseBoundaries: nextState.pendingSuspenseBoundaries,
       transitions: nextState.transitions,
     };
     const updateQueue: UpdateQueue<RootState> = (workInProgress.updateQueue: any);
@@ -1980,6 +1977,7 @@ function mountSuspenseOffscreenState(renderLanes: Lanes): OffscreenState {
   return {
     baseLanes: renderLanes,
     cachePool: getSuspendedCache(),
+    transitions: null,
   };
 }
 
@@ -2014,6 +2012,7 @@ function updateSuspenseOffscreenState(
   return {
     baseLanes: mergeLanes(prevOffscreenState.baseLanes, renderLanes),
     cachePool,
+    transitions: prevOffscreenState.transitions,
   };
 }
 
@@ -3580,9 +3579,6 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         const cache: Cache = current.memoizedState.cache;
         pushCacheProvider(workInProgress, cache);
       }
-      if (enableTransitionTracing) {
-        workInProgress.memoizedState.transitions = getWorkInProgressTransitions();
-      }
       resetHydrationState();
       break;
     case HostComponent:

packages/react-reconciler/src/ReactFiberCommitWork.new.js

@@ -2725,9 +2725,8 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
-        const transitions = finishedWork.memoizedState.transitions;
-        if (transitions !== null) {
-          transitions.forEach(transition => {
+        if (committedTransitions !== null) {
+          committedTransitions.forEach(transition => {
             // TODO(luna) Do we want to log TransitionStart in the startTransition callback instead?
             addTransitionStartCallbackToPendingTransition({
               transitionName: transition.name,

packages/react-reconciler/src/ReactFiberCommitWork.old.js

@@ -2725,9 +2725,8 @@ function commitPassiveMountOnFiber(
       }
 
       if (enableTransitionTracing) {
-        const transitions = finishedWork.memoizedState.transitions;
-        if (transitions !== null) {
-          transitions.forEach(transition => {
+        if (committedTransitions !== null) {
+          committedTransitions.forEach(transition => {
             // TODO(luna) Do we want to log TransitionStart in the startTransition callback instead?
             addTransitionStartCallbackToPendingTransition({
               transitionName: transition.name,

packages/react-reconciler/src/ReactFiberOffscreenComponent.js

@@ -10,7 +10,7 @@
 import type {ReactNodeList, OffscreenMode} from 'shared/ReactTypes';
 import type {Lanes} from './ReactFiberLane.old';
 import type {SpawnedCachePool} from './ReactFiberCacheComponent.new';
-
+import type {Transition} from './ReactFiberTracingMarkerComponent.new';
 export type OffscreenProps = {|
   // TODO: Pick an API before exposing the Offscreen type. I've chosen an enum
   // for now, since we might have multiple variants. For example, hiding the
@@ -30,6 +30,7 @@ export type OffscreenState = {|
   // order to unhide the component.
   baseLanes: Lanes,
   cachePool: SpawnedCachePool | null,
+  transitions: Set<Transition> | null,
 |};
 
 export type OffscreenInstance = {};

packages/react-reconciler/src/ReactFiberRoot.new.js

@@ -15,7 +15,10 @@ import type {
 } from './ReactInternalTypes';
 import type {RootTag} from './ReactRootTags';
 import type {Cache} from './ReactFiberCacheComponent.new';
-import type {Transition} from './ReactFiberTracingMarkerComponent.new';
+import type {
+  PendingSuspenseBoundaries,
+  Transition,
+} from './ReactFiberTracingMarkerComponent.new';
 
 import {noTimeout, supportsHydration} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber.new';
@@ -42,7 +45,8 @@ export type RootState = {
   element: any,
   isDehydrated: boolean,
   cache: Cache,
-  transitions: Array<Transition> | null,
+  pendingSuspenseBoundaries: PendingSuspenseBoundaries | null,
+  transitions: Set<Transition> | null,
 };
 
 function FiberRootNode(
@@ -184,6 +188,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: initialCache,
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   } else {
@@ -192,6 +197,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: (null: any), // not enabled yet
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   }

packages/react-reconciler/src/ReactFiberRoot.old.js

@@ -15,7 +15,10 @@ import type {
 } from './ReactInternalTypes';
 import type {RootTag} from './ReactRootTags';
 import type {Cache} from './ReactFiberCacheComponent.old';
-import type {Transition} from './ReactFiberTracingMarkerComponent.old';
+import type {
+  PendingSuspenseBoundaries,
+  Transition,
+} from './ReactFiberTracingMarkerComponent.old';
 
 import {noTimeout, supportsHydration} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber.old';
@@ -42,7 +45,8 @@ export type RootState = {
   element: any,
   isDehydrated: boolean,
   cache: Cache,
-  transitions: Array<Transition> | null,
+  pendingSuspenseBoundaries: PendingSuspenseBoundaries | null,
+  transitions: Set<Transition> | null,
 };
 
 function FiberRootNode(
@@ -184,6 +188,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: initialCache,
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   } else {
@@ -192,6 +197,7 @@ export function createFiberRoot(
       isDehydrated: hydrate,
       cache: (null: any), // not enabled yet
       transitions: null,
+      pendingSuspenseBoundaries: null,
     };
     uninitializedFiber.memoizedState = initialState;
   }

packages/react-reconciler/src/ReactFiberSuspenseComponent.new.js

@@ -29,6 +29,7 @@ export type SuspenseProps = {|
   suspenseCallback?: (Set<Wakeable> | null) => mixed,
 
   unstable_expectedLoadTime?: number,
+  unstable_name?: string,
 |};
 
 // A null SuspenseState represents an unsuspended normal Suspense boundary.

packages/react-reconciler/src/ReactFiberSuspenseComponent.old.js

@@ -29,6 +29,7 @@ export type SuspenseProps = {|
   suspenseCallback?: (Set<Wakeable> | null) => mixed,
 
   unstable_expectedLoadTime?: number,
+  unstable_name?: string,
 |};
 
 // A null SuspenseState represents an unsuspended normal Suspense boundary.

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.new.js

@@ -7,8 +7,9 @@
  * @flow
  */
 
-import type {TransitionTracingCallbacks} from './ReactInternalTypes';
-import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
+import type {TransitionTracingCallbacks, Fiber} from './ReactInternalTypes';
+import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+
 import {enableTransitionTracing} from 'shared/ReactFeatureFlags';
 
 export type SuspenseInfo = {name: string | null};
@@ -34,10 +35,7 @@ export type BatchConfigTransition = {
   _updatedFibers?: Set<Fiber>,
 };
 
-export type TransitionCallback = 0 | 1;
-
-export const TransitionStart = 0;
-export const TransitionComplete = 1;
+export type PendingSuspenseBoundaries = Map<OffscreenInstance, SuspenseInfo>;
 
 export function processTransitionCallbacks(
   pendingTransitions: PendingTransitionCallbacks,

packages/react-reconciler/src/ReactFiberTracingMarkerComponent.old.js

@@ -7,8 +7,9 @@
  * @flow
  */
 
-import type {TransitionTracingCallbacks} from './ReactInternalTypes';
-import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
+import type {TransitionTracingCallbacks, Fiber} from './ReactInternalTypes';
+import type {OffscreenInstance} from './ReactFiberOffscreenComponent';
+
 import {enableTransitionTracing} from 'shared/ReactFeatureFlags';
 
 export type SuspenseInfo = {name: string | null};
@@ -34,10 +35,7 @@ export type BatchConfigTransition = {
   _updatedFibers?: Set<Fiber>,
 };
 
-export type TransitionCallback = 0 | 1;
-
-export const TransitionStart = 0;
-export const TransitionComplete = 1;
+export type PendingSuspenseBoundaries = Map<OffscreenInstance, SuspenseInfo>;
 
 export function processTransitionCallbacks(
   pendingTransitions: PendingTransitionCallbacks,