Refactor code

Author: lunaruan

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

@@ -89,6 +89,7 @@ import {
   StaticMask,
   ShouldCapture,
   ForceClientRender,
+  Passive,
 } from './ReactFiberFlags';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {
@@ -979,11 +980,13 @@ function updateTracingMarkerComponent(
       const markerInstance: TracingMarkerInstance = {
         tag: TransitionTracingMarker,
         transitions: new Set(currentTransitions),
-        pendingBoundaries: new Map(),
+        pendingBoundaries: null,
         name: workInProgress.pendingProps.name,
         aborts: null,
       };
       workInProgress.stateNode = markerInstance;
+
+      workInProgress.flags |= Passive;
     }
   } else {
     if (__DEV__) {

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

@@ -1140,68 +1140,95 @@ function commitLayoutEffectOnFiber(
   }
 }
 
-function abortParentMarkerTransitions(
+function abortTransitionMarker(
   deletedFiber: Fiber,
-  nearestMountedAncestor: Fiber,
+  currentFiber: Fiber,
   abort: TransitionAbort,
+  isInDeletedTree: boolean,
 ) {
-  const deletedFiberInstance: OffscreenInstance = deletedFiber.stateNode;
-
-  let fiber = deletedFiber;
-  let isInDeletedTree = true;
-  while (fiber !== null) {
-    switch (fiber.tag) {
-      case TracingMarkerComponent:
-        const transitions = deletedFiberInstance.transitions;
-
-        const markerInstance = fiber.stateNode;
-        const markerTransitions = markerInstance.transitions;
-        if (markerTransitions !== null && transitions !== null) {
-          let abortMarker = false;
-          transitions.forEach(transition => {
-            if (markerTransitions.has(transition)) {
-              abortMarker = true;
-            }
-          });
+  const markerInstance: TracingMarkerInstance = currentFiber.stateNode;
+  const deletedInstance: OffscreenInstance = deletedFiber.stateNode;
+  const deletedTransitions = deletedFiber.stateNode.transitions;
+  const pendingBoundaries = markerInstance.pendingBoundaries;
+
+  let aborts = markerInstance.aborts;
+  if (aborts === null) {
+    markerInstance.aborts = aborts = [];
+  }
 
-          if (abortMarker) {
-            if (markerInstance.aborts === null) {
-              markerInstance.aborts = new Set();
-            }
+  aborts.push(abort);
+  addMarkerIncompleteCallbackToPendingTransition(
+    currentFiber.memoizedProps.name,
+    deletedTransitions,
+    aborts,
+  );
 
-            markerInstance.aborts.add(abort);
-            addMarkerIncompleteCallbackToPendingTransition(
-              fiber.memoizedProps.name,
-              transitions,
-              markerInstance.aborts,
-            );
+  // We only want to call onTransitionProgress when the marker hasn't been
+  // deleted
+  if (
+    !isInDeletedTree &&
+    pendingBoundaries !== null &&
+    pendingBoundaries.has(deletedInstance)
+  ) {
+    pendingBoundaries.delete(deletedInstance);
 
-            // We only want to call onTransitionProgress when the marker hasn't been
-            // deleted
-            if (
-              !isInDeletedTree &&
-              markerInstance.pendingBoundaries !== null &&
-              markerInstance.pendingBoundaries.has(deletedFiberInstance)
-            ) {
-              markerInstance.pendingBoundaries.delete(deletedFiberInstance);
+    addMarkerProgressCallbackToPendingTransition(
+      currentFiber.memoizedProps.name,
+      deletedTransitions,
+      pendingBoundaries,
+    );
+  }
+}
 
-              addMarkerProgressCallbackToPendingTransition(
-                fiber.memoizedProps.name,
-                transitions,
-                markerInstance.pendingBoundaries,
-              );
-            }
+function abortParentMarkerTransitions(
+  deletedFiber: Fiber,
+  nearestMountedAncestor: Fiber,
+  abort: TransitionAbort,
+) {
+  // Find all pending markers that are waiting on child suspense boundaries in the
+  // aborted subtree and cancels them
+  const deletedFiberInstance: OffscreenInstance = deletedFiber.stateNode;
+  const abortedTransitions = deletedFiberInstance.transitions;
+  if (abortedTransitions !== null) {
+    let fiber = deletedFiber;
+    let isInDeletedTree = true;
+    while (fiber !== null) {
+      switch (fiber.tag) {
+        case TracingMarkerComponent:
+          const markerInstance: TracingMarkerInstance = fiber.stateNode;
+          const markerTransitions = markerInstance.transitions;
+          if (markerTransitions !== null) {
+            // TODO: Refactor this code. Is there a way to move this code to
+            // the deletions phase instead of calculating it here while making sure
+            // complete is called appropriately?
+            abortedTransitions.forEach(transition => {
+              // If one of the transitions on the tracing marker is a transition
+              // that was in an aborted subtree, we will abort that tracing marker
+              if (
+                fiber !== null &&
+                markerTransitions.has(transition) &&
+                (markerInstance.aborts === null ||
+                  !markerInstance.aborts.includes(abort))
+              ) {
+                abortTransitionMarker(
+                  deletedFiber,
+                  fiber,
+                  abort,
+                  isInDeletedTree,
+                );
+              }
+            });
           }
-        }
-        break;
-      case HostRoot:
-        const root = fiber.stateNode;
-        const incompleteTransitions = root.incompleteTransitions;
-
-        if (deletedFiberInstance.transitions !== null) {
-          deletedFiberInstance.transitions.forEach(transition => {
-            if (incompleteTransitions.has(transition)) {
-              const transitionInstance = incompleteTransitions.get(transition);
+          break;
+        case HostRoot:
+          const root = fiber.stateNode;
+          const rootTransitions = root.incompleteTransitions;
+
+          abortedTransitions.forEach(transition => {
+            if (rootTransitions.has(transition)) {
+              const transitionInstance: TracingMarkerInstance = rootTransitions.get(
+                transition,
+              );
               if (transitionInstance.aborts === null) {
                 transitionInstance.aborts = [];
               }
@@ -1217,20 +1244,20 @@ function abortParentMarkerTransitions(
               }
             }
           });
-        }
-        break;
-      default:
-        break;
-    }
+          break;
+        default:
+          break;
+      }
 
-    if (
-      nearestMountedAncestor.deletions !== null &&
-      nearestMountedAncestor.deletions.includes(fiber)
-    ) {
-      isInDeletedTree = false;
-      fiber = nearestMountedAncestor;
-    } else {
-      fiber = fiber.return;
+      if (
+        nearestMountedAncestor.deletions !== null &&
+        nearestMountedAncestor.deletions.includes(fiber)
+      ) {
+        isInDeletedTree = false;
+        fiber = nearestMountedAncestor;
+      } else {
+        fiber = fiber.return;
+      }
     }
   }
 }
@@ -1280,6 +1307,7 @@ function commitTransitionProgress(offscreenFiber: Fiber) {
         pendingMarkers.forEach(markerInstance => {
           const pendingBoundaries = markerInstance.pendingBoundaries;
           const transitions = markerInstance.transitions;
+          const markerName = markerInstance.name;
           if (
             pendingBoundaries !== null &&
             !pendingBoundaries.has(offscreenInstance)
@@ -1290,10 +1318,10 @@ function commitTransitionProgress(offscreenFiber: Fiber) {
             if (transitions !== null) {
               if (
                 markerInstance.tag === TransitionTracingMarker &&
-                markerInstance.name !== undefined
+                markerName !== null
               ) {
                 addMarkerProgressCallbackToPendingTransition(
-                  markerInstance.name,
+                  markerName,
                   transitions,
                   pendingBoundaries,
                 );
@@ -1317,6 +1345,7 @@ function commitTransitionProgress(offscreenFiber: Fiber) {
         pendingMarkers.forEach(markerInstance => {
           const pendingBoundaries = markerInstance.pendingBoundaries;
           const transitions = markerInstance.transitions;
+          const markerName = markerInstance.name;
           if (
             pendingBoundaries !== null &&
             pendingBoundaries.has(offscreenInstance)
@@ -1325,13 +1354,25 @@ function commitTransitionProgress(offscreenFiber: Fiber) {
             if (transitions !== null) {
               if (
                 markerInstance.tag === TransitionTracingMarker &&
-                markerInstance.name !== undefined
+                markerName !== null
               ) {
                 addMarkerProgressCallbackToPendingTransition(
-                  markerInstance.name,
+                  markerName,
                   transitions,
                   pendingBoundaries,
                 );
+
+                if (pendingBoundaries.size === 0) {
+                  if (markerInstance.aborts === null) {
+                    addMarkerCompleteCallbackToPendingTransition(
+                      markerName,
+                      transitions,
+                    );
+                  }
+                  markerInstance.transitions = null;
+                  markerInstance.pendingBoundaries = null;
+                  markerInstance.aborts = null;
+                }
               } else if (markerInstance.tag === TransitionRoot) {
                 transitions.forEach(transition => {
                   addTransitionProgressCallbackToPendingTransition(
@@ -2097,28 +2138,6 @@ function commitDeletionEffectsOnFiber(
         offscreenSubtreeWasHidden =
           prevOffscreenSubtreeWasHidden || deletedFiber.memoizedState !== null;
 
-        if (enableTransitionTracing) {
-          // We need to mark this fiber's parents as deleted
-          const instance: OffscreenInstance = deletedFiber.stateNode;
-          const transitions = instance.transitions;
-          if (transitions !== null) {
-            let name = null;
-            const parent = deletedFiber.return;
-            if (
-              parent !== null &&
-              parent.tag === SuspenseComponent &&
-              parent.memoizedProps.unstable_name
-            ) {
-              name = parent.memoizedProps.unstable_name;
-            }
-
-            abortParentMarkerTransitions(deletedFiber, nearestMountedAncestor, {
-              reason: 'suspense',
-              name,
-            });
-          }
-        }
-
         recursivelyTraverseDeletionEffects(
           finishedRoot,
           nearestMountedAncestor,
@@ -2134,21 +2153,36 @@ function commitDeletionEffectsOnFiber(
       }
       break;
     }
+    case SuspenseComponent: {
+      if (enableTransitionTracing) {
+        // We need to mark this fiber's parents as deleted
+        const offscreenFiber: Fiber = (deletedFiber.child: any);
+        const instance: OffscreenInstance = offscreenFiber.stateNode;
+        const transitions = instance.transitions;
+        if (transitions !== null) {
+          abortParentMarkerTransitions(offscreenFiber, nearestMountedAncestor, {
+            reason: 'suspense',
+            name: deletedFiber.memoizedProps.unstable_name || null,
+          });
+        }
+      }
+      recursivelyTraverseDeletionEffects(
+        finishedRoot,
+        nearestMountedAncestor,
+        deletedFiber,
+      );
+      return;
+    }
     case TracingMarkerComponent: {
       if (enableTransitionTracing) {
         // We need to mark this fiber's parents as deleted
         const instance: TracingMarkerInstance = deletedFiber.stateNode;
         const transitions = instance.transitions;
         if (transitions !== null) {
-          const abort = {
+          abortParentMarkerTransitions(deletedFiber, nearestMountedAncestor, {
             reason: 'marker',
             name: deletedFiber.memoizedProps.name,
-          };
-          abortParentMarkerTransitions(
-            deletedFiber,
-            nearestMountedAncestor,
-            abort,
-          );
+          });
         }
       }
       recursivelyTraverseDeletionEffects(
@@ -3134,6 +3168,7 @@ function commitOffscreenPassiveMountEffects(
 
     commitTransitionProgress(finishedWork);
 
+    // TODO: Refactor this into an if/else branch
     if (!isHidden) {
       instance.transitions = null;
       instance.pendingMarkers = null;
@@ -3168,18 +3203,14 @@ function commitCachePassiveMountEffect(
 function commitTracingMarkerPassiveMountEffect(finishedWork: Fiber) {
   // Get the transitions that were initiatized during the render
   // and add a start transition callback for each of them
+  // We will only call this on initial mount of the tracing marker
+  // only if there are no suspense children
   const instance = finishedWork.stateNode;
-  if (
-    instance.transitions !== null &&
-    (instance.pendingBoundaries === null ||
-      instance.pendingBoundaries.size === 0)
-  ) {
-    if (instance.aborts === null) {
-      addMarkerCompleteCallbackToPendingTransition(
-        finishedWork.memoizedProps.name,
-        instance.transitions,
-      );
-    }
+  if (instance.transitions !== null && instance.pendingBoundaries === null) {
+    addMarkerCompleteCallbackToPendingTransition(
+      finishedWork.memoizedProps.name,
+      instance.transitions,
+    );
     instance.transitions = null;
     instance.pendingBoundaries = null;
     instance.aborts = null;
@@ -3674,21 +3705,6 @@ function commitAtomicPassiveEffects(
       }
       break;
     }
-    case TracingMarkerComponent: {
-      if (enableTransitionTracing) {
-        recursivelyTraverseAtomicPassiveEffects(
-          finishedRoot,
-          finishedWork,
-          committedLanes,
-          committedTransitions,
-        );
-        if (flags & Passive) {
-          commitTracingMarkerPassiveMountEffect(finishedWork);
-        }
-        break;
-      }
-      // Intentional fallthrough to next branch
-    }
     // eslint-disable-next-line-no-fallthrough
     default: {
       recursivelyTraverseAtomicPassiveEffects(