Handle in Fiber

We match new Fibers regardless of what their key is if the current key is
optimistic.

This means that the new Fiber might have a different key for the first time
so we need to handle updating it and restoring it properly.

For the complex cases of maps, I use a negative index to represent previous
optimistic keys so that we can look up to see if we match one of those.

Author: sebmarkbage

packages/react-reconciler/src/ReactChildFiber.js

@@ -16,6 +16,7 @@ import type {
   ReactComponentInfo,
   SuspenseListRevealOrder,
   ReactKey,
+  ReactOptimisticKey,
 } from 'shared/ReactTypes';
 import type {Fiber} from './ReactInternalTypes';
 import type {Lanes} from './ReactFiberLane';
@@ -38,6 +39,7 @@ import {
   REACT_LAZY_TYPE,
   REACT_CONTEXT_TYPE,
   REACT_LEGACY_ELEMENT_TYPE,
+  REACT_OPTIMISTIC_KEY,
 } from 'shared/ReactSymbols';
 import {
   HostRoot,
@@ -51,6 +53,7 @@ import {
   enableAsyncIterableChildren,
   disableLegacyMode,
   enableFragmentRefs,
+  enableOptimisticKey,
 } from 'shared/ReactFeatureFlags';
 
 import {
@@ -463,18 +466,33 @@ function createChildReconciler(
 
   function mapRemainingChildren(
     currentFirstChild: Fiber,
-  ): Map<string | number, Fiber> {
+  ): Map<string | number | ReactOptimisticKey, Fiber> {
     // Add the remaining children to a temporary map so that we can find them by
     // keys quickly. Implicit (null) keys get added to this set with their index
     // instead.
-    const existingChildren: Map<string | number, Fiber> = new Map();
+    const existingChildren: Map<
+      | string
+      | number
+      // This type is only here for the case when enableOptimisticKey is disabled.
+      // Remove it after it ships.
+      | ReactOptimisticKey,
+      Fiber,
+    > = new Map();
 
     let existingChild: null | Fiber = currentFirstChild;
     while (existingChild !== null) {
-      if (existingChild.key !== null) {
-        existingChildren.set(existingChild.key, existingChild);
-      } else {
+      if (existingChild.key === null) {
         existingChildren.set(existingChild.index, existingChild);
+      } else if (
+        enableOptimisticKey &&
+        existingChild.key === REACT_OPTIMISTIC_KEY
+      ) {
+        // For optimistic keys, we store the negative index (minus one) to differentiate
+        // them from the regular indices. We'll look this up regardless of what the new
+        // key is, if there's no other match.
+        existingChildren.set(-existingChild.index - 1, existingChild);
+      } else {
+        existingChildren.set(existingChild.key, existingChild);
       }
       existingChild = existingChild.sibling;
     }
@@ -637,6 +655,10 @@ function createChildReconciler(
     } else {
       // Update
       const existing = useFiber(current, portal.children || []);
+      if (enableOptimisticKey) {
+        // If the old key was optimistic we need to now save the real one.
+        existing.key = portal.key;
+      }
       existing.return = returnFiber;
       if (__DEV__) {
         existing._debugInfo = currentDebugInfo;
@@ -671,6 +693,10 @@ function createChildReconciler(
     } else {
       // Update
       const existing = useFiber(current, fragment);
+      if (enableOptimisticKey) {
+        // If the old key was optimistic we need to now save the real one.
+        existing.key = key;
+      }
       existing.return = returnFiber;
       if (__DEV__) {
         existing._debugInfo = currentDebugInfo;
@@ -841,7 +867,10 @@ function createChildReconciler(
     if (typeof newChild === 'object' && newChild !== null) {
       switch (newChild.$$typeof) {
         case REACT_ELEMENT_TYPE: {
-          if (newChild.key === key) {
+          if (
+            newChild.key === key ||
+            (enableOptimisticKey && key === REACT_OPTIMISTIC_KEY)
+          ) {
             const prevDebugInfo = pushDebugInfo(newChild._debugInfo);
             const updated = updateElement(
               returnFiber,
@@ -856,7 +885,10 @@ function createChildReconciler(
           }
         }
         case REACT_PORTAL_TYPE: {
-          if (newChild.key === key) {
+          if (
+            newChild.key === key ||
+            (enableOptimisticKey && key === REACT_OPTIMISTIC_KEY)
+          ) {
             return updatePortal(returnFiber, oldFiber, newChild, lanes);
           } else {
             return null;
@@ -940,7 +972,7 @@ function createChildReconciler(
   }
 
   function updateFromMap(
-    existingChildren: Map<string | number, Fiber>,
+    existingChildren: Map<string | number | ReactOptimisticKey, Fiber>,
     returnFiber: Fiber,
     newIdx: number,
     newChild: any,
@@ -969,7 +1001,11 @@ function createChildReconciler(
           const matchedFiber =
             existingChildren.get(
               newChild.key === null ? newIdx : newChild.key,
-            ) || null;
+            ) ||
+            (enableOptimisticKey &&
+              // If the existing child was an optimistic key, we may still match on the index.
+              existingChildren.get(-newIdx - 1)) ||
+            null;
           const prevDebugInfo = pushDebugInfo(newChild._debugInfo);
           const updated = updateElement(
             returnFiber,
@@ -984,7 +1020,11 @@ function createChildReconciler(
           const matchedFiber =
             existingChildren.get(
               newChild.key === null ? newIdx : newChild.key,
-            ) || null;
+            ) ||
+            (enableOptimisticKey &&
+              // If the existing child was an optimistic key, we may still match on the index.
+              existingChildren.get(-newIdx - 1)) ||
+            null;
           return updatePortal(returnFiber, matchedFiber, newChild, lanes);
         }
         case REACT_LAZY_TYPE: {
@@ -1275,14 +1315,22 @@ function createChildReconciler(
           );
         }
         if (shouldTrackSideEffects) {
-          if (newFiber.alternate !== null) {
+          const currentFiber = newFiber.alternate;
+          if (currentFiber !== null) {
             // The new fiber is a work in progress, but if there exists a
             // current, that means that we reused the fiber. We need to delete
             // it from the child list so that we don't add it to the deletion
             // list.
-            existingChildren.delete(
-              newFiber.key === null ? newIdx : newFiber.key,
-            );
+            if (
+              enableOptimisticKey &&
+              currentFiber.key === REACT_OPTIMISTIC_KEY
+            ) {
+              existingChildren.delete(-newIdx - 1);
+            } else {
+              existingChildren.delete(
+                currentFiber.key === null ? newIdx : currentFiber.key,
+              );
+            }
           }
         }
         lastPlacedIndex = placeChild(newFiber, lastPlacedIndex, newIdx);
@@ -1569,14 +1617,22 @@ function createChildReconciler(
           );
         }
         if (shouldTrackSideEffects) {
-          if (newFiber.alternate !== null) {
+          const currentFiber = newFiber.alternate;
+          if (currentFiber !== null) {
             // The new fiber is a work in progress, but if there exists a
             // current, that means that we reused the fiber. We need to delete
             // it from the child list so that we don't add it to the deletion
             // list.
-            existingChildren.delete(
-              newFiber.key === null ? newIdx : newFiber.key,
-            );
+            if (
+              enableOptimisticKey &&
+              currentFiber.key === REACT_OPTIMISTIC_KEY
+            ) {
+              existingChildren.delete(-newIdx - 1);
+            } else {
+              existingChildren.delete(
+                currentFiber.key === null ? newIdx : currentFiber.key,
+              );
+            }
           }
         }
         lastPlacedIndex = placeChild(newFiber, lastPlacedIndex, newIdx);
@@ -1643,12 +1699,19 @@ function createChildReconciler(
     while (child !== null) {
       // TODO: If key === null and child.key === null, then this only applies to
       // the first item in the list.
-      if (child.key === key) {
+      if (
+        child.key === key ||
+        (enableOptimisticKey && child.key === REACT_OPTIMISTIC_KEY)
+      ) {
         const elementType = element.type;
         if (elementType === REACT_FRAGMENT_TYPE) {
           if (child.tag === Fragment) {
             deleteRemainingChildren(returnFiber, child.sibling);
             const existing = useFiber(child, element.props.children);
+            if (enableOptimisticKey) {
+              // If the old key was optimistic we need to now save the real one.
+              existing.key = key;
+            }
             if (enableFragmentRefs) {
               coerceRef(existing, element);
             }
@@ -1678,6 +1741,10 @@ function createChildReconciler(
           ) {
             deleteRemainingChildren(returnFiber, child.sibling);
             const existing = useFiber(child, element.props);
+            if (enableOptimisticKey) {
+              // If the old key was optimistic we need to now save the real one.
+              existing.key = key;
+            }
             coerceRef(existing, element);
             existing.return = returnFiber;
             if (__DEV__) {
@@ -1737,14 +1804,21 @@ function createChildReconciler(
     while (child !== null) {
       // TODO: If key === null and child.key === null, then this only applies to
       // the first item in the list.
-      if (child.key === key) {
+      if (
+        child.key === key ||
+        (enableOptimisticKey && child.key === REACT_OPTIMISTIC_KEY)
+      ) {
         if (
           child.tag === HostPortal &&
           child.stateNode.containerInfo === portal.containerInfo &&
           child.stateNode.implementation === portal.implementation
         ) {
           deleteRemainingChildren(returnFiber, child.sibling);
           const existing = useFiber(child, portal.children || []);
+          if (enableOptimisticKey) {
+            // If the old key was optimistic we need to now save the real one.
+            existing.key = key;
+          }
           existing.return = returnFiber;
           return existing;
         } else {

packages/react-reconciler/src/ReactFiber.js

@@ -44,6 +44,7 @@ import {
   enableObjectFiber,
   enableViewTransition,
   enableSuspenseyImages,
+  enableOptimisticKey,
 } from 'shared/ReactFeatureFlags';
 import {NoFlags, Placement, StaticMask} from './ReactFiberFlags';
 import {ConcurrentRoot} from './ReactRootTags';
@@ -365,6 +366,12 @@ export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
     workInProgress.subtreeFlags = NoFlags;
     workInProgress.deletions = null;
 
+    if (enableOptimisticKey) {
+      // For optimistic keys, the Fibers can have different keys if one is optimistic
+      // and the other one is filled in.
+      workInProgress.key = current.key;
+    }
+
     if (enableProfilerTimer) {
       // We intentionally reset, rather than copy, actualDuration & actualStartTime.
       // This prevents time from endlessly accumulating in new commits.
@@ -489,8 +496,15 @@ export function resetWorkInProgress(
     workInProgress.memoizedState = current.memoizedState;
     workInProgress.updateQueue = current.updateQueue;
     // Needed because Blocks store data on type.
+    // TODO: Blocks don't exist anymore. Do we still need this?
     workInProgress.type = current.type;
 
+    if (enableOptimisticKey) {
+      // For optimistic keys, the Fibers can have different keys if one is optimistic
+      // and the other one is filled in.
+      workInProgress.key = current.key;
+    }
+
     // Clone the dependencies object. This is mutated during the render phase, so
     // it cannot be shared with the current fiber.
     const currentDependencies = current.dependencies;