memo supports Hooks

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -36,6 +36,8 @@ import {
   Profiler,
   SuspenseComponent,
   IncompleteClassComponent,
+  MemoComponent,
+  SimpleMemoComponent,
 } from 'shared/ReactWorkTags';
 import {
   invokeGuardedCallback,
@@ -216,7 +218,8 @@ function commitBeforeMutationLifeCycles(
 ): void {
   switch (finishedWork.tag) {
     case FunctionComponent:
-    case ForwardRef: {
+    case ForwardRef:
+    case SimpleMemoComponent: {
       commitHookEffectList(UnmountSnapshot, NoHookEffect, finishedWork);
       return;
     }
@@ -313,7 +316,8 @@ function commitLifeCycles(
 ): void {
   switch (finishedWork.tag) {
     case FunctionComponent:
-    case ForwardRef: {
+    case ForwardRef:
+    case SimpleMemoComponent: {
       commitHookEffectList(UnmountLayout, MountLayout, finishedWork);
       const newUpdateQueue: FunctionComponentUpdateQueue | null = (finishedWork.updateQueue: any);
       if (newUpdateQueue !== null) {
@@ -588,7 +592,9 @@ function commitUnmount(current: Fiber): void {
 
   switch (current.tag) {
     case FunctionComponent:
-    case ForwardRef: {
+    case ForwardRef:
+    case MemoComponent:
+    case SimpleMemoComponent: {
       const updateQueue: FunctionComponentUpdateQueue | null = (current.updateQueue: any);
       if (updateQueue !== null) {
         const lastEffect = updateQueue.lastEffect;
@@ -981,7 +987,9 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
   if (!supportsMutation) {
     switch (finishedWork.tag) {
       case FunctionComponent:
-      case ForwardRef: {
+      case ForwardRef:
+      case MemoComponent:
+      case SimpleMemoComponent: {
         commitHookEffectList(UnmountMutation, MountMutation, finishedWork);
         return;
       }
@@ -993,7 +1001,9 @@ function commitWork(current: Fiber | null, finishedWork: Fiber): void {
 
   switch (finishedWork.tag) {
     case FunctionComponent:
-    case ForwardRef: {
+    case ForwardRef:
+    case MemoComponent:
+    case SimpleMemoComponent: {
       commitHookEffectList(UnmountMutation, MountMutation, finishedWork);
       return;
     }

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -52,7 +52,8 @@ import {
   FunctionComponent,
   HostPortal,
   HostRoot,
-  PureComponent,
+  MemoComponent,
+  SimpleMemoComponent,
 } from 'shared/ReactWorkTags';
 import {
   enableSchedulerTracing,
@@ -270,8 +271,8 @@ let nextEffect: Fiber | null = null;
 
 let isCommitting: boolean = false;
 let rootWithPendingPassiveEffects: FiberRoot | null = null;
-let firstPassiveEffect: Fiber | null = null;
 let passiveEffectCallbackHandle: * = null;
+let passiveEffectCallback: * = null;
 
 let legacyErrorBoundariesThatAlreadyFailed: Set<mixed> | null = null;
 
@@ -521,6 +522,7 @@ function commitAllLifeCycles(
 function commitPassiveEffects(root: FiberRoot, firstEffect: Fiber): void {
   rootWithPendingPassiveEffects = null;
   passiveEffectCallbackHandle = null;
+  passiveEffectCallback = null;
 
   // Set this to true to prevent re-entrancy
   const previousIsRendering = isRendering;
@@ -600,15 +602,13 @@ function flushPassiveEffectsBeforeSchedulingUpdateOnFiber(fiber: Fiber) {
 
 function flushPassiveEffects(root: FiberRoot) {
   if (
-    passiveEffectCallbackHandle !== null &&
+    passiveEffectCallback !== null &&
     root === rootWithPendingPassiveEffects
   ) {
     Schedule_cancelCallback(passiveEffectCallbackHandle);
-    passiveEffectCallbackHandle = null;
-    rootWithPendingPassiveEffects = null;
-    if (firstPassiveEffect !== null) {
-      commitPassiveEffects(root, firstPassiveEffect);
-    }
+    // We call the scheduled callback instead of commitPassiveEffects directly
+    // to ensure tracing works correctly.
+    passiveEffectCallback();
   }
 }
 
@@ -807,26 +807,15 @@ function commitRoot(root: FiberRoot, finishedWork: Fiber): void {
     // after the next paint. Schedule an callback to fire them in an async
     // event. To ensure serial execution, the callback will be flushed early if
     // we enter rootWithPendingPassiveEffects commit phase before then.
-    const resolvedFirstEffect = (firstPassiveEffect = firstEffect);
-
-    let passiveEffectCallback;
+    let callback = commitPassiveEffects.bind(null, root, firstEffect);
     if (enableSchedulerTracing) {
       // TODO: Avoid this extra callback by mutating the tracing ref directly,
       // like we do at the beginning of commitRoot. I've opted not to do that
       // here because that code is still in flux.
-      passiveEffectCallback = Schedule_tracing_wrap(() => {
-        commitPassiveEffects(root, resolvedFirstEffect);
-      });
-    } else {
-      passiveEffectCallback = commitPassiveEffects.bind(
-        null,
-        root,
-        resolvedFirstEffect,
-      );
+      callback = Schedule_tracing_wrap(callback);
     }
-    passiveEffectCallbackHandle = Schedule_scheduleCallback(
-      passiveEffectCallback,
-    );
+    passiveEffectCallbackHandle = Schedule_scheduleCallback(callback);
+    passiveEffectCallback = callback;
   }
 
   isCommitting = false;
@@ -1798,7 +1787,8 @@ function scheduleWorkToRoot(fiber: Fiber, expirationTime): FiberRoot | null {
           break;
         case FunctionComponent:
         case ForwardRef:
-        case PureComponent:
+        case MemoComponent:
+        case SimpleMemoComponent:
           warnAboutUpdateOnUnmounted(fiber, false);
           break;
       }

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

@@ -15,6 +15,7 @@
 let React;
 let ReactFeatureFlags;
 let ReactNoop;
+let SchedulerTracing;
 let useState;
 let useReducer;
 let useEffect;
@@ -26,6 +27,7 @@ let useRef;
 let useAPI;
 let forwardRef;
 let flushPassiveEffects;
+let memo;
 
 describe('ReactHooks', () => {
   beforeEach(() => {
@@ -56,8 +58,10 @@ describe('ReactHooks', () => {
     ReactFeatureFlags = require('shared/ReactFeatureFlags');
     ReactFeatureFlags.debugRenderPhaseSideEffectsForStrictMode = false;
     ReactFeatureFlags.enableHooks = true;
+    ReactFeatureFlags.enableSchedulerTracing = true;
     React = require('react');
     ReactNoop = require('react-noop-renderer');
+    SchedulerTracing = require('scheduler/tracing');
     useState = React.useState;
     useReducer = React.useReducer;
     useEffect = React.useEffect;
@@ -68,6 +72,7 @@ describe('ReactHooks', () => {
     useRef = React.useRef;
     useAPI = React.useAPI;
     forwardRef = React.forwardRef;
+    memo = React.memo;
   });
 
   function span(prop) {
@@ -330,6 +335,28 @@ describe('ReactHooks', () => {
           '    in Counter (at **)',
       );
     });
+
+    it('works with memo', () => {
+      let _updateCount;
+      function Counter(props) {
+        const [count, updateCount] = useState(0);
+        _updateCount = updateCount;
+        return <Text text={'Count: ' + count} />;
+      }
+      Counter = memo(Counter);
+
+      ReactNoop.render(<Counter />);
+      expect(ReactNoop.flush()).toEqual(['Count: 0']);
+      expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
+
+      ReactNoop.render(<Counter />);
+      expect(ReactNoop.flush()).toEqual([]);
+      expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
+
+      _updateCount(1);
+      expect(ReactNoop.flush()).toEqual(['Count: 1']);
+      expect(ReactNoop.getChildren()).toEqual([span('Count: 1')]);
+    });
   });
 
   describe('updates during the render phase', () => {
@@ -797,6 +824,7 @@ describe('ReactHooks', () => {
         }, []);
         return <Text text={'Count: ' + count} />;
       }
+
       ReactNoop.render(<Counter count={0} />);
       expect(ReactNoop.flush()).toEqual(['Count: 0']);
       expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
@@ -808,6 +836,56 @@ describe('ReactHooks', () => {
       expect(ReactNoop.getChildren()).toEqual([span('Count: 2')]);
     });
 
+    it('flushes serial effects before enqueueing work (with tracing)', () => {
+      const onInteractionScheduledWorkCompleted = jest.fn();
+      const onWorkCanceled = jest.fn();
+      SchedulerTracing.unstable_subscribe({
+        onInteractionScheduledWorkCompleted,
+        onInteractionTraced: jest.fn(),
+        onWorkCanceled,
+        onWorkScheduled: jest.fn(),
+        onWorkStarted: jest.fn(),
+        onWorkStopped: jest.fn(),
+      });
+
+      let _updateCount;
+      function Counter(props) {
+        const [count, updateCount] = useState(0);
+        _updateCount = updateCount;
+        useEffect(() => {
+          expect(SchedulerTracing.unstable_getCurrent()).toMatchInteractions([
+            tracingEvent,
+          ]);
+          ReactNoop.yield(`Will set count to 1`);
+          updateCount(1);
+        }, []);
+        return <Text text={'Count: ' + count} />;
+      }
+
+      const tracingEvent = {id: 0, name: 'hello', timestamp: 0};
+      SchedulerTracing.unstable_trace(
+        tracingEvent.name,
+        tracingEvent.timestamp,
+        () => {
+          ReactNoop.render(<Counter count={0} />);
+        },
+      );
+      expect(ReactNoop.flush()).toEqual(['Count: 0']);
+      expect(ReactNoop.getChildren()).toEqual([span('Count: 0')]);
+
+      expect(onInteractionScheduledWorkCompleted).toHaveBeenCalledTimes(0);
+
+      // Enqueuing this update forces the passive effect to be flushed --
+      // updateCount(1) happens first, so 2 wins.
+      _updateCount(2);
+      expect(ReactNoop.flush()).toEqual(['Will set count to 1', 'Count: 2']);
+      expect(ReactNoop.getChildren()).toEqual([span('Count: 2')]);
+
+      flushPassiveEffects();
+      expect(onInteractionScheduledWorkCompleted).toHaveBeenCalledTimes(1);
+      expect(onWorkCanceled).toHaveBeenCalledTimes(0);
+    });
+
     it(
       'in sync mode, useEffect is deferred and updates finish synchronously ' +
         '(in a single batch)',