Track which fibers scheduled the current render work

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

@@ -37,6 +37,7 @@ import {
   enableStrictEffects,
   deletedTreeCleanUpLevel,
   enableSuspenseLayoutEffectSemantics,
+  enableUpdaterTracking,
 } from 'shared/ReactFeatureFlags';
 import {
   FunctionComponent,
@@ -89,7 +90,7 @@ import {
   resetCurrentFiber as resetCurrentDebugFiberInDEV,
   setCurrentFiber as setCurrentDebugFiberInDEV,
 } from './ReactCurrentFiber';
-
+import {isDevToolsPresent} from './ReactFiberDevToolsHook.new';
 import {onCommitUnmount} from './ReactFiberDevToolsHook.new';
 import {resolveDefaultProps} from './ReactFiberLazyComponent.new';
 import {
@@ -137,6 +138,7 @@ import {
   resolveRetryWakeable,
   markCommitTimeOfFallback,
   enqueuePendingPassiveProfilerEffect,
+  restorePendingUpdaters,
 } from './ReactFiberWorkLoop.new';
 import {
   NoFlags as NoHookEffect,
@@ -162,6 +164,10 @@ const PossiblyWeakSet = typeof WeakSet === 'function' ? WeakSet : Set;
 
 let nextEffect: Fiber | null = null;
 
+// Used for Profiling builds to track updaters.
+let inProgressLanes: Lanes | null = null;
+let inProgressRoot: FiberRoot | null = null;
+
 const callComponentWillUnmountWithTimer = function(current, instance) {
   instance.props = current.memoizedProps;
   instance.state = current.memoizedState;
@@ -2094,6 +2100,20 @@ function attachSuspenseRetryListeners(finishedWork: Fiber) {
           }
         }
         retryCache.add(wakeable);
+
+        if (enableUpdaterTracking) {
+          if (isDevToolsPresent) {
+            if (inProgressLanes !== null && inProgressRoot !== null) {
+              // If we have pending work still, associate the original updaters with it.
+              restorePendingUpdaters(inProgressRoot, inProgressLanes);
+            } else {
+              throw Error(
+                'Expected finished root and lanes to be set. This is a bug in React.',
+              );
+            }
+          }
+        }
+
         wakeable.then(retry, retry);
       }
     });
@@ -2124,9 +2144,19 @@ function commitResetTextContent(current: Fiber) {
   resetTextContent(current.stateNode);
 }
 
-export function commitMutationEffects(root: FiberRoot, firstChild: Fiber) {
+export function commitMutationEffects(
+  root: FiberRoot,
+  firstChild: Fiber,
+  committedLanes: Lanes,
+) {
+  inProgressLanes = committedLanes;
+  inProgressRoot = root;
   nextEffect = firstChild;
+
   commitMutationEffects_begin(root);
+
+  inProgressLanes = null;
+  inProgressRoot = null;
 }
 
 function commitMutationEffects_begin(root: FiberRoot) {
@@ -2280,8 +2310,14 @@ export function commitLayoutEffects(
   root: FiberRoot,
   committedLanes: Lanes,
 ): void {
+  inProgressLanes = committedLanes;
+  inProgressRoot = root;
   nextEffect = finishedWork;
+
   commitLayoutEffects_begin(finishedWork, root, committedLanes);
+
+  inProgressLanes = null;
+  inProgressRoot = null;
 }
 
 function commitLayoutEffects_begin(

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

@@ -37,6 +37,7 @@ import {
   enableStrictEffects,
   deletedTreeCleanUpLevel,
   enableSuspenseLayoutEffectSemantics,
+  enableUpdaterTracking,
 } from 'shared/ReactFeatureFlags';
 import {
   FunctionComponent,
@@ -89,7 +90,7 @@ import {
   resetCurrentFiber as resetCurrentDebugFiberInDEV,
   setCurrentFiber as setCurrentDebugFiberInDEV,
 } from './ReactCurrentFiber';
-
+import {isDevToolsPresent} from './ReactFiberDevToolsHook.old';
 import {onCommitUnmount} from './ReactFiberDevToolsHook.old';
 import {resolveDefaultProps} from './ReactFiberLazyComponent.old';
 import {
@@ -137,6 +138,7 @@ import {
   resolveRetryWakeable,
   markCommitTimeOfFallback,
   enqueuePendingPassiveProfilerEffect,
+  restorePendingUpdaters,
 } from './ReactFiberWorkLoop.old';
 import {
   NoFlags as NoHookEffect,
@@ -162,6 +164,10 @@ const PossiblyWeakSet = typeof WeakSet === 'function' ? WeakSet : Set;
 
 let nextEffect: Fiber | null = null;
 
+// Used for Profiling builds to track updaters.
+let inProgressLanes: Lanes | null = null;
+let inProgressRoot: FiberRoot | null = null;
+
 const callComponentWillUnmountWithTimer = function(current, instance) {
   instance.props = current.memoizedProps;
   instance.state = current.memoizedState;
@@ -2094,6 +2100,20 @@ function attachSuspenseRetryListeners(finishedWork: Fiber) {
           }
         }
         retryCache.add(wakeable);
+
+        if (enableUpdaterTracking) {
+          if (isDevToolsPresent) {
+            if (inProgressLanes !== null && inProgressRoot !== null) {
+              // If we have pending work still, associate the original updaters with it.
+              restorePendingUpdaters(inProgressRoot, inProgressLanes);
+            } else {
+              throw Error(
+                'Expected finished root and lanes to be set. This is a bug in React.',
+              );
+            }
+          }
+        }
+
         wakeable.then(retry, retry);
       }
     });
@@ -2124,9 +2144,19 @@ function commitResetTextContent(current: Fiber) {
   resetTextContent(current.stateNode);
 }
 
-export function commitMutationEffects(root: FiberRoot, firstChild: Fiber) {
+export function commitMutationEffects(
+  root: FiberRoot,
+  firstChild: Fiber,
+  committedLanes: Lanes,
+) {
+  inProgressLanes = committedLanes;
+  inProgressRoot = root;
   nextEffect = firstChild;
+
   commitMutationEffects_begin(root);
+
+  inProgressLanes = null;
+  inProgressRoot = null;
 }
 
 function commitMutationEffects_begin(root: FiberRoot) {
@@ -2280,8 +2310,14 @@ export function commitLayoutEffects(
   root: FiberRoot,
   committedLanes: Lanes,
 ): void {
+  inProgressLanes = committedLanes;
+  inProgressRoot = root;
   nextEffect = finishedWork;
+
   commitLayoutEffects_begin(finishedWork, root, committedLanes);
+
+  inProgressLanes = null;
+  inProgressRoot = null;
 }
 
 function commitLayoutEffects_begin(

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

@@ -35,7 +35,12 @@ export type Lanes = number;
 export type Lane = number;
 export type LaneMap<T> = Array<T>;
 
-import {enableCache, enableSchedulingProfiler} from 'shared/ReactFeatureFlags';
+import {
+  enableCache,
+  enableSchedulingProfiler,
+  enableUpdaterTracking,
+} from 'shared/ReactFeatureFlags';
+import {isDevToolsPresent} from './ReactFiberDevToolsHook.new';
 
 // Lane values below should be kept in sync with getLabelsForLanes(), used by react-devtools-scheduling-profiler.
 // If those values are changed that package should be rebuilt and redeployed.
@@ -742,6 +747,57 @@ export function getBumpedLaneForHydration(
   return lane;
 }
 
+export function addFiberToLanesMap(
+  root: FiberRoot,
+  fiber: Fiber,
+  lanes: Lanes | Lane,
+) {
+  if (!enableUpdaterTracking) {
+    return;
+  }
+  if (!isDevToolsPresent) {
+    return;
+  }
+  const pendingUpdatersLaneMap = root.pendingUpdatersLaneMap;
+  while (lanes > 0) {
+    const index = laneToIndex(lanes);
+    const lane = 1 << index;
+
+    const updaters = pendingUpdatersLaneMap[index];
+    updaters.add(fiber);
+
+    lanes &= ~lane;
+  }
+}
+
+export function movePendingFibersToMemoized(root: FiberRoot, lanes: Lanes) {
+  if (!enableUpdaterTracking) {
+    return;
+  }
+  if (!isDevToolsPresent) {
+    return;
+  }
+  const pendingUpdatersLaneMap = root.pendingUpdatersLaneMap;
+  const memoizedUpdaters = root.memoizedUpdaters;
+  while (lanes > 0) {
+    const index = laneToIndex(lanes);
+    const lane = 1 << index;
+
+    const updaters = pendingUpdatersLaneMap[index];
+    if (updaters.size > 0) {
+      updaters.forEach(fiber => {
+        const alternate = fiber.alternate;
+        if (alternate === null || !memoizedUpdaters.has(alternate)) {
+          memoizedUpdaters.add(fiber);
+        }
+      });
+      updaters.clear();
+    }
+
+    lanes &= ~lane;
+  }
+}
+
 const clz32 = Math.clz32 ? Math.clz32 : clz32Fallback;
 
 // Count leading zeros. Only used on lanes, so assume input is an integer.

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

@@ -35,7 +35,12 @@ export type Lanes = number;
 export type Lane = number;
 export type LaneMap<T> = Array<T>;
 
-import {enableCache, enableSchedulingProfiler} from 'shared/ReactFeatureFlags';
+import {
+  enableCache,
+  enableSchedulingProfiler,
+  enableUpdaterTracking,
+} from 'shared/ReactFeatureFlags';
+import {isDevToolsPresent} from './ReactFiberDevToolsHook.old';
 
 // Lane values below should be kept in sync with getLabelsForLanes(), used by react-devtools-scheduling-profiler.
 // If those values are changed that package should be rebuilt and redeployed.
@@ -742,6 +747,57 @@ export function getBumpedLaneForHydration(
   return lane;
 }
 
+export function addFiberToLanesMap(
+  root: FiberRoot,
+  fiber: Fiber,
+  lanes: Lanes | Lane,
+) {
+  if (!enableUpdaterTracking) {
+    return;
+  }
+  if (!isDevToolsPresent) {
+    return;
+  }
+  const pendingUpdatersLaneMap = root.pendingUpdatersLaneMap;
+  while (lanes > 0) {
+    const index = laneToIndex(lanes);
+    const lane = 1 << index;
+
+    const updaters = pendingUpdatersLaneMap[index];
+    updaters.add(fiber);
+
+    lanes &= ~lane;
+  }
+}
+
+export function movePendingFibersToMemoized(root: FiberRoot, lanes: Lanes) {
+  if (!enableUpdaterTracking) {
+    return;
+  }
+  if (!isDevToolsPresent) {
+    return;
+  }
+  const pendingUpdatersLaneMap = root.pendingUpdatersLaneMap;
+  const memoizedUpdaters = root.memoizedUpdaters;
+  while (lanes > 0) {
+    const index = laneToIndex(lanes);
+    const lane = 1 << index;
+
+    const updaters = pendingUpdatersLaneMap[index];
+    if (updaters.size > 0) {
+      updaters.forEach(fiber => {
+        const alternate = fiber.alternate;
+        if (alternate === null || !memoizedUpdaters.has(alternate)) {
+          memoizedUpdaters.add(fiber);
+        }
+      });
+      updaters.clear();
+    }
+
+    lanes &= ~lane;
+  }
+}
+
 const clz32 = Math.clz32 ? Math.clz32 : clz32Fallback;
 
 // Count leading zeros. Only used on lanes, so assume input is an integer.

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

@@ -16,6 +16,7 @@ import {
   NoLane,
   NoLanes,
   NoTimestamp,
+  TotalLanes,
   createLaneMap,
 } from './ReactFiberLane.new';
 import {
@@ -24,6 +25,7 @@ import {
   enableCache,
   enableProfilerCommitHooks,
   enableProfilerTimer,
+  enableUpdaterTracking,
 } from 'shared/ReactFeatureFlags';
 import {unstable_getThreadID} from 'scheduler/tracing';
 import {initializeUpdateQueue} from './ReactUpdateQueue.new';
@@ -77,6 +79,14 @@ function FiberRootNode(containerInfo, tag, hydrate) {
     this.passiveEffectDuration = 0;
   }
 
+  if (enableUpdaterTracking) {
+    this.memoizedUpdaters = new Set();
+    const pendingUpdatersLaneMap = (this.pendingUpdatersLaneMap = []);
+    for (let i = 0; i < TotalLanes; i++) {
+      pendingUpdatersLaneMap.push(new Set());
+    }
+  }
+
   if (__DEV__) {
     switch (tag) {
       case ConcurrentRoot:

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

@@ -16,6 +16,7 @@ import {
   NoLane,
   NoLanes,
   NoTimestamp,
+  TotalLanes,
   createLaneMap,
 } from './ReactFiberLane.old';
 import {
@@ -24,6 +25,7 @@ import {
   enableCache,
   enableProfilerCommitHooks,
   enableProfilerTimer,
+  enableUpdaterTracking,
 } from 'shared/ReactFeatureFlags';
 import {unstable_getThreadID} from 'scheduler/tracing';
 import {initializeUpdateQueue} from './ReactUpdateQueue.old';
@@ -77,6 +79,14 @@ function FiberRootNode(containerInfo, tag, hydrate) {
     this.passiveEffectDuration = 0;
   }
 
+  if (enableUpdaterTracking) {
+    this.memoizedUpdaters = new Set();
+    const pendingUpdatersLaneMap = (this.pendingUpdatersLaneMap = []);
+    for (let i = 0; i < TotalLanes; i++) {
+      pendingUpdatersLaneMap.push(new Set());
+    }
+  }
+
   if (__DEV__) {
     switch (tag) {
       case ConcurrentRoot: