useMutableSource: "Entangle" instead of expiring

A lane is said to be entangled with another when it's not allowed to
render in a batch that does not also include the other lane.

This commit implements entanglement for `useMutableSource`. If a source
is mutated in between when it's read in the render phase, but before
it's subscribed to in the commit phase, we must account for whether the
same source has pending mutations elsewhere. The old subscriptions must
not be allowed to re-render without also including the new subscription
(and vice versa), to prevent tearing.

In the old reconciler, we did this by synchronously flushing all the
pending subscription updates. This works, but isn't ideal. The new
reconciler can entangle the updates without de-opting to sync.

In the future, we plan to use this same mechanism for other features,
like skipping over intermediate useTransition states.

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

@@ -32,7 +32,7 @@ import {
   isSubsetOfLanes,
   mergeLanes,
   removeLanes,
-  markRootExpired,
+  markRootEntangled,
   markRootMutableRead,
 } from './ReactFiberLane';
 import {readContext} from './ReactFiberNewContext.new';
@@ -1004,7 +1004,7 @@ function useMutableSource<Source, Snapshot>(
         // There is no mechanism currently to associate these updates though,
         // so for now we fall back to synchronously flushing all pending updates.
         // TODO: This should entangle the lanes instead of expiring everything.
-        markRootExpired(root, root.mutableReadLanes);
+        markRootEntangled(root, root.mutableReadLanes);
       }
     }
   }, [getSnapshot, source, subscribe]);

packages/react-reconciler/src/ReactFiberLane.js

@@ -365,6 +365,37 @@ export function getNextLanes(root: FiberRoot, wipLanes: Lanes): Lanes {
     }
   }
 
+  // Check for entangled lanes and add them to the batch.
+  //
+  // A lane is said to be entangled with another when it's not allowed to render
+  // in a batch that does not also include the other lane. Typically we do this
+  // when multiple updates have the same source, and we only want to respond to
+  // the most recent event from that source.
+  //
+  // Note that we apply entanglements *after* checking for partial work above.
+  // This means that if a lane is entangled during an interleaved event while
+  // it's already rendering, we won't interrupt it. This is intentional, since
+  // entanglement is usually "best effort": we'll try our best to render the
+  // lanes in the same batch, but it's not worth throwing out partially
+  // completed work in order to do it.
+  //
+  // For those exceptions where entanglement is semantically important, like
+  // useMutableSource, we should ensure that there is no partial work at the
+  // time we apply the entanglement.
+  const entangledLanes = root.entangledLanes;
+  if (entangledLanes !== NoLanes) {
+    const entanglements = root.entanglements;
+    let lanes = nextLanes & entangledLanes;
+    while (lanes > 0) {
+      const index = ctrz(lanes);
+      const lane = 1 << index;
+
+      nextLanes |= entanglements[index];
+
+      lanes &= ~lane;
+    }
+  }
+
   return nextLanes;
 }
 
@@ -692,6 +723,8 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
   root.expiredLanes &= remainingLanes;
   root.mutableReadLanes &= remainingLanes;
 
+  root.entangledLanes &= remainingLanes;
+
   const expirationTimes = root.expirationTimes;
   let lanes = noLongerPendingLanes;
   while (lanes > 0) {
@@ -705,6 +738,21 @@ export function markRootFinished(root: FiberRoot, remainingLanes: Lanes) {
   }
 }
 
+export function markRootEntangled(root: FiberRoot, entangledLanes: Lanes) {
+  root.entangledLanes |= entangledLanes;
+
+  const entanglements = root.entanglements;
+  let lanes = entangledLanes;
+  while (lanes > 0) {
+    const index = ctrz(lanes);
+    const lane = 1 << index;
+
+    entanglements[index] |= entangledLanes;
+
+    lanes &= ~lane;
+  }
+}
+
 export function getBumpedLaneForHydration(
   root: FiberRoot,
   renderLanes: Lanes,

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

@@ -49,6 +49,9 @@ function FiberRootNode(containerInfo, tag, hydrate) {
 
   this.finishedLanes = NoLanes;
 
+  this.entangledLanes = NoLanes;
+  this.entanglements = createLaneMap(NoLanes);
+
   if (enableSchedulerTracing) {
     this.interactionThreadID = unstable_getThreadID();
     this.memoizedInteractions = new Set();

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

@@ -673,6 +673,15 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
   currentEventWipLanes = NoLanes;
   currentEventPendingLanes = NoLanes;
 
+  invariant(
+    (executionContext & (RenderContext | CommitContext)) === NoContext,
+    'Should not already be working.',
+  );
+
+  // Flush any pending passive effects before deciding which lanes to work on,
+  // in case they schedule additional work.
+  flushPassiveEffects();
+
   // Determine the next expiration time to work on, using the fields stored
   // on the root.
   let lanes = getNextLanes(
@@ -697,12 +706,6 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
   }
 
   const originalCallbackNode = root.callbackNode;
-  invariant(
-    (executionContext & (RenderContext | CommitContext)) === NoContext,
-    'Should not already be working.',
-  );
-
-  flushPassiveEffects();
 
   let exitStatus = renderRootConcurrent(root, lanes);
 

packages/react-reconciler/src/ReactInternalTypes.js

@@ -262,6 +262,9 @@ type BaseFiberRootProperties = {|
   mutableReadLanes: Lanes,
 
   finishedLanes: Lanes,
+
+  entangledLanes: Lanes,
+  entanglements: LaneMap<Lanes>,
 |};
 
 // The following attributes are only used by interaction tracing builds.

packages/react-reconciler/src/__tests__/useMutableSource-test.internal.js

@@ -1481,21 +1481,39 @@ describe('useMutableSource', () => {
           source.valueB = '3';
         },
       );
-    });
-
-    expect(Scheduler).toHaveYielded([
-      // The partial render completes
-      'Child: 2',
-      'Commit: 2, 2',
 
-      // Then we start rendering the low priority mutation
-      'Parent: 3',
+      expect(Scheduler).toFlushAndYieldThrough([
+        // The partial render completes
+        'Child: 2',
+        'Commit: 2, 2',
+      ]);
 
-      // Eventually the child corrects itself, because of the check that
-      // occurs when re-subscribing.
-      'Child: 3',
-      'Commit: 3, 3',
-    ]);
+      // Now there are two pending mutations at different priorities. But they
+      // both read the same verion of the mutable source, so we must render
+      // them simultaneously.
+      //
+      if (gate(flags => flags.new)) {
+        // In the new reconciler, we can do this with entanglement: when the
+        // high priority render starts, we'll also include the low pri work.
+        expect(Scheduler).toFlushAndYieldThrough([
+          'Parent: 3',
+          // Demonstrates that we can yield here
+        ]);
+        expect(Scheduler).toFlushAndYield([
+          // Now finish the rest of the update
+          'Child: 3',
+          'Commit: 3, 3',
+        ]);
+      } else {
+        // In the old reconciler, we don't have an entanglement mechanism. The
+        // best we can do is synchronously flush both updates.
+        expect(Scheduler).toFlushAndYield([
+          'Parent: 3',
+          'Child: 3',
+          'Commit: 3, 3',
+        ]);
+      }
+    });
   });
 
   // @gate experimental