act: Batch updates, even in legacy roots

In legacy roots, if an update originates outside of `batchedUpdates`,
check if it's inside an `act` scope; if so, treat it as if it were
batched. This is only necessary in legacy roots because in concurrent
roots, updates are batched by default.

With this change, the Test Utils and Test Renderer versions of `act` are
nothing more than aliases of the isomorphic API (still not exposed, but
will likely be the recommended API that replaces the others).

Author: acdlite

packages/react-dom/src/test-utils/ReactTestUtils.js

@@ -33,14 +33,8 @@ const getNodeFromInstance = EventInternals[1];
 const getFiberCurrentPropsFromNode = EventInternals[2];
 const enqueueStateRestore = EventInternals[3];
 const restoreStateIfNeeded = EventInternals[4];
-const batchedUpdates = EventInternals[5];
 
-const act_notBatchedInLegacyMode = React.unstable_act;
-function act(callback) {
-  return act_notBatchedInLegacyMode(() => {
-    return batchedUpdates(callback);
-  });
-}
+const act = React.unstable_act;
 
 function Event(suffix) {}
 

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

@@ -518,7 +518,9 @@ export function scheduleUpdateOnFiber(
   if (
     lane === SyncLane &&
     executionContext === NoContext &&
-    (fiber.mode & ConcurrentMode) === NoMode
+    (fiber.mode & ConcurrentMode) === NoMode &&
+    // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+    !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
   ) {
     // Flush the synchronous work now, unless we're already working or inside
     // a batch. This is intentionally inside scheduleUpdateOnFiber instead of
@@ -1049,7 +1051,11 @@ export function batchedUpdates<A, R>(fn: A => R, a: A): R {
     executionContext = prevExecutionContext;
     // If there were legacy sync updates, flush them at the end of the outer
     // most batchedUpdates-like method.
-    if (executionContext === NoContext) {
+    if (
+      executionContext === NoContext &&
+      // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+      !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
+    ) {
       resetRenderTimer();
       flushSyncCallbacksOnlyInLegacyMode();
     }

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

@@ -518,7 +518,9 @@ export function scheduleUpdateOnFiber(
   if (
     lane === SyncLane &&
     executionContext === NoContext &&
-    (fiber.mode & ConcurrentMode) === NoMode
+    (fiber.mode & ConcurrentMode) === NoMode &&
+    // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+    !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
   ) {
     // Flush the synchronous work now, unless we're already working or inside
     // a batch. This is intentionally inside scheduleUpdateOnFiber instead of
@@ -1049,7 +1051,11 @@ export function batchedUpdates<A, R>(fn: A => R, a: A): R {
     executionContext = prevExecutionContext;
     // If there were legacy sync updates, flush them at the end of the outer
     // most batchedUpdates-like method.
-    if (executionContext === NoContext) {
+    if (
+      executionContext === NoContext &&
+      // Treat `act` as if it's inside `batchedUpdates`, even in legacy mode.
+      !(__DEV__ && ReactCurrentActQueue.isBatchingLegacy)
+    ) {
       resetRenderTimer();
       flushSyncCallbacksOnlyInLegacyMode();
     }

packages/react-reconciler/src/__tests__/ReactIsomorphicAct-test.js

@@ -78,4 +78,32 @@ describe('isomorphic act()', () => {
     });
     expect(returnValue).toEqual('hi');
   });
+
+  // @gate __DEV__
+  test('in legacy mode, updates are batched', () => {
+    const root = ReactNoop.createLegacyRoot();
+
+    // Outside of `act`, legacy updates are flushed completely synchronously
+    root.render('A');
+    expect(root).toMatchRenderedOutput('A');
+
+    // `act` will batch the updates and flush them at the end
+    act(() => {
+      root.render('B');
+      // Hasn't flushed yet
+      expect(root).toMatchRenderedOutput('A');
+
+      // Confirm that a nested `batchedUpdates` call won't cause the updates
+      // to flush early.
+      ReactNoop.batchedUpdates(() => {
+        root.render('C');
+      });
+
+      // Still hasn't flushed
+      expect(root).toMatchRenderedOutput('A');
+    });
+
+    // Now everything renders in a single batch.
+    expect(root).toMatchRenderedOutput('C');
+  });
 });

packages/react-test-renderer/src/ReactTestRenderer.js

@@ -7,7 +7,6 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes';
 import type {Fiber} from 'react-reconciler/src/ReactInternalTypes';
 import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
 import type {Instance, TextInstance} from './ReactTestHostConfig';
@@ -50,12 +49,7 @@ import {getPublicInstance} from './ReactTestHostConfig';
 import {ConcurrentRoot, LegacyRoot} from 'react-reconciler/src/ReactRootTags';
 import {allowConcurrentByDefault} from 'shared/ReactFeatureFlags';
 
-const act_notBatchedInLegacyMode = React.unstable_act;
-function act<T>(callback: () => T): Thenable<T> {
-  return act_notBatchedInLegacyMode(() => {
-    return batchedUpdates(callback);
-  });
-}
+const act = React.unstable_act;
 
 // TODO: Remove from public bundle
 

packages/react/src/ReactAct.js

@@ -28,12 +28,20 @@ export function act<T>(callback: () => T | Thenable<T>): Thenable<T> {
       ReactCurrentActQueue.current = [];
     }
 
+    const prevIsBatchingLegacy = ReactCurrentActQueue.isBatchingLegacy;
     let result;
     try {
+      // Used to reproduce behavior of `batchedUpdates` in legacy mode. Only
+      // set to `true` while the given callback is executed, not for updates
+      // triggered during an async event, because this is how the legacy
+      // implementation of `act` behaved.
+      ReactCurrentActQueue.isBatchingLegacy = true;
       result = callback();
     } catch (error) {
       popActScope(prevActScopeDepth);
       throw error;
+    } finally {
+      ReactCurrentActQueue.isBatchingLegacy = prevIsBatchingLegacy;
     }
 
     if (

packages/react/src/ReactCurrentActQueue.js

@@ -17,6 +17,9 @@ const ReactCurrentActQueue = {
   // on at the testing frameworks layer? Instead of what we do now, which
   // is check if a `jest` global is defined.
   disableActWarning: (false: boolean),
+
+  // Used to reproduce behavior of `batchedUpdates` in legacy mode.
+  isBatchingLegacy: false,
 };
 
 export default ReactCurrentActQueue;