Add support for hooks to ReactDOMServer

Co-authored-by: Alex Taylor <alexmckenley@gmail.com>
Co-authored-by: Andrew Clark <acdlite@fb.com>

packages/react-dom/src/__tests__/ReactDOMServerIntegrationHooks-test.internal.js

@@ -0,0 +1,650 @@
+/**
+ * Copyright (c) 2013-present, Facebook, Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+/* eslint-disable no-func-assign */
+
+'use strict';
+
+const ReactDOMServerIntegrationUtils = require('./utils/ReactDOMServerIntegrationTestUtils');
+
+let React;
+let ReactFeatureFlags;
+let ReactDOM;
+let ReactDOMServer;
+let useState;
+let useReducer;
+let useEffect;
+let useContext;
+let useCallback;
+let useMemo;
+let useRef;
+let useAPI;
+let useMutationEffect;
+let useLayoutEffect;
+let forwardRef;
+let yieldedValues;
+let yieldValue;
+let clearYields;
+
+function initModules() {
+  // Reset warning cache.
+  jest.resetModuleRegistry();
+
+  ReactFeatureFlags = require('shared/ReactFeatureFlags');
+  ReactFeatureFlags.debugRenderPhaseSideEffectsForStrictMode = false;
+  ReactFeatureFlags.enableHooks = true;
+  React = require('react');
+  ReactDOM = require('react-dom');
+  ReactDOMServer = require('react-dom/server');
+  useState = React.useState;
+  useReducer = React.useReducer;
+  useEffect = React.useEffect;
+  useContext = React.useContext;
+  useCallback = React.useCallback;
+  useMemo = React.useMemo;
+  useRef = React.useRef;
+  useAPI = React.useAPI;
+  useMutationEffect = React.useMutationEffect;
+  useLayoutEffect = React.useLayoutEffect;
+  forwardRef = React.forwardRef;
+
+  yieldedValues = [];
+  yieldValue = value => {
+    yieldedValues.push(value);
+  };
+  clearYields = () => {
+    const ret = yieldedValues;
+    yieldedValues = [];
+    return ret;
+  };
+
+  // Make them available to the helpers.
+  return {
+    ReactDOM,
+    ReactDOMServer,
+  };
+}
+
+const {
+  resetModules,
+  itRenders,
+  itThrowsWhenRendering,
+  serverRender,
+} = ReactDOMServerIntegrationUtils(initModules);
+
+describe('ReactDOMServerHooks', () => {
+  beforeEach(() => {
+    resetModules();
+  });
+
+  function Text(props) {
+    yieldValue(props.text);
+    return <span>{props.text}</span>;
+  }
+
+  describe('useState', () => {
+    itRenders('basic render', async render => {
+      function Counter(props) {
+        const [count] = useState(0);
+        return <span>Count: {count}</span>;
+      }
+
+      const domNode = await render(<Counter />);
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+
+    itRenders('lazy state initialization', async render => {
+      function Counter(props) {
+        const [count] = useState(() => {
+          return 0;
+        });
+        return <span>Count: {count}</span>;
+      }
+
+      const domNode = await render(<Counter />);
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+
+    it('does not trigger a re-renders when updater is invoked outside current render function', async () => {
+      function UpdateCount({setCount, count, children}) {
+        if (count < 3) {
+          setCount(c => c + 1);
+        }
+        return <span>{children}</span>;
+      }
+      function Counter() {
+        let [count, setCount] = useState(0);
+        return (
+          <div>
+            <UpdateCount setCount={setCount} count={count}>
+              Count: {count}
+            </UpdateCount>
+          </div>
+        );
+      }
+
+      const domNode = await serverRender(<Counter />);
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+
+    itThrowsWhenRendering(
+      'if used inside a class component',
+      async render => {
+        class Counter extends React.Component {
+          render() {
+            let [count] = useState(0);
+            return <Text text={count} />;
+          }
+        }
+
+        return render(<Counter />);
+      },
+      'Hooks can only be called inside the body of a functional component.',
+    );
+
+    itRenders('multiple times when an updater is called', async render => {
+      function Counter() {
+        let [count, setCount] = useState(0);
+        if (count < 12) {
+          setCount(c => c + 1);
+          setCount(c => c + 1);
+          setCount(c => c + 1);
+        }
+        return <Text text={'Count: ' + count} />;
+      }
+
+      const domNode = await render(<Counter />);
+      expect(domNode.textContent).toEqual('Count: 12');
+    });
+
+    itRenders('until there are no more new updates', async render => {
+      function Counter() {
+        let [count, setCount] = useState(0);
+        if (count < 3) {
+          setCount(count + 1);
+        }
+        return <span>Count: {count}</span>;
+      }
+
+      const domNode = await render(<Counter />);
+      expect(domNode.textContent).toEqual('Count: 3');
+    });
+
+    itThrowsWhenRendering(
+      'after too many iterations',
+      async render => {
+        function Counter() {
+          let [count, setCount] = useState(0);
+          setCount(count + 1);
+          return <span>{count}</span>;
+        }
+        return render(<Counter />);
+      },
+      'Too many re-renders. React limits the number of renders to prevent ' +
+        'an infinite loop.',
+    );
+  });
+
+  describe('useReducer', () => {
+    itRenders('with initial state', async render => {
+      function reducer(state, action) {
+        return action === 'increment' ? state + 1 : state;
+      }
+      function Counter() {
+        let [count] = useReducer(reducer, 0);
+        yieldValue('Render: ' + count);
+        return <Text text={count} />;
+      }
+
+      const domNode = await render(<Counter />);
+
+      expect(clearYields()).toEqual(['Render: 0', 0]);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('0');
+    });
+
+    itRenders('lazy initialization with initialAction', async render => {
+      function reducer(state, action) {
+        return action === 'increment' ? state + 1 : state;
+      }
+      function Counter() {
+        let [count] = useReducer(reducer, 0, 'increment');
+        yieldValue('Render: ' + count);
+        return <Text text={count} />;
+      }
+
+      const domNode = await render(<Counter />);
+
+      expect(clearYields()).toEqual(['Render: 1', 1]);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('1');
+    });
+
+    itRenders(
+      'multiple times when updates happen during the render phase',
+      async render => {
+        function reducer(state, action) {
+          return action === 'increment' ? state + 1 : state;
+        }
+        function Counter() {
+          let [count, dispatch] = useReducer(reducer, 0);
+          if (count < 3) {
+            dispatch('increment');
+          }
+          yieldValue('Render: ' + count);
+          return <Text text={count} />;
+        }
+
+        const domNode = await render(<Counter />);
+
+        expect(clearYields()).toEqual([
+          'Render: 0',
+          'Render: 1',
+          'Render: 2',
+          'Render: 3',
+          3,
+        ]);
+        expect(domNode.tagName).toEqual('SPAN');
+        expect(domNode.textContent).toEqual('3');
+      },
+    );
+
+    itRenders(
+      'using reducer passed at time of render, not time of dispatch',
+      async render => {
+        // This test is a bit contrived but it demonstrates a subtle edge case.
+
+        // Reducer A increments by 1. Reducer B increments by 10.
+        function reducerA(state, action) {
+          switch (action) {
+            case 'increment':
+              return state + 1;
+            case 'reset':
+              return 0;
+          }
+        }
+        function reducerB(state, action) {
+          switch (action) {
+            case 'increment':
+              return state + 10;
+            case 'reset':
+              return 0;
+          }
+        }
+
+        function Counter() {
+          let [reducer, setReducer] = useState(() => reducerA);
+          let [count, dispatch] = useReducer(reducer, 0);
+          if (count < 20) {
+            dispatch('increment');
+            // Swap reducers each time we increment
+            if (reducer === reducerA) {
+              setReducer(() => reducerB);
+            } else {
+              setReducer(() => reducerA);
+            }
+          }
+          yieldValue('Render: ' + count);
+          return <Text text={count} />;
+        }
+
+        const domNode = await render(<Counter />);
+
+        expect(clearYields()).toEqual([
+          // The count should increase by alternating amounts of 10 and 1
+          // until we reach 21.
+          'Render: 0',
+          'Render: 10',
+          'Render: 11',
+          'Render: 21',
+          21,
+        ]);
+        expect(domNode.tagName).toEqual('SPAN');
+        expect(domNode.textContent).toEqual('21');
+      },
+    );
+  });
+
+  describe('useMemo', () => {
+    itRenders('basic render', async render => {
+      function CapitalizedText(props) {
+        const text = props.text;
+        const capitalizedText = useMemo(
+          () => {
+            yieldValue(`Capitalize '${text}'`);
+            return text.toUpperCase();
+          },
+          [text],
+        );
+        return <Text text={capitalizedText} />;
+      }
+
+      const domNode = await render(<CapitalizedText text="hello" />);
+      expect(clearYields()).toEqual(["Capitalize 'hello'", 'HELLO']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('HELLO');
+    });
+
+    itRenders('if no inputs are provided', async render => {
+      function LazyCompute(props) {
+        const computed = useMemo(props.compute);
+        return <Text text={computed} />;
+      }
+
+      function computeA() {
+        yieldValue('compute A');
+        return 'A';
+      }
+
+      const domNode = await render(<LazyCompute compute={computeA} />);
+      expect(clearYields()).toEqual(['compute A', 'A']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('A');
+    });
+
+    itRenders(
+      'multiple times when updates happen during the render phase',
+      async render => {
+        function CapitalizedText(props) {
+          const [text, setText] = useState(props.text);
+          const capitalizedText = useMemo(
+            () => {
+              yieldValue(`Capitalize '${text}'`);
+              return text.toUpperCase();
+            },
+            [text],
+          );
+
+          if (text === 'hello') {
+            setText('hello, world.');
+          }
+          return <Text text={capitalizedText} />;
+        }
+
+        const domNode = await render(<CapitalizedText text="hello" />);
+        expect(clearYields()).toEqual([
+          "Capitalize 'hello'",
+          "Capitalize 'hello, world.'",
+          'HELLO, WORLD.',
+        ]);
+        expect(domNode.tagName).toEqual('SPAN');
+        expect(domNode.textContent).toEqual('HELLO, WORLD.');
+      },
+    );
+
+    itRenders(
+      'should only invoke the memoized function when the inputs change',
+      async render => {
+        function CapitalizedText(props) {
+          const [text, setText] = useState(props.text);
+          const [count, setCount] = useState(0);
+          const capitalizedText = useMemo(
+            () => {
+              yieldValue(`Capitalize '${text}'`);
+              return text.toUpperCase();
+            },
+            [text],
+          );
+
+          yieldValue(count);
+
+          if (count < 3) {
+            setCount(count + 1);
+          }
+
+          if (text === 'hello' && count === 2) {
+            setText('hello, world.');
+          }
+          return <Text text={capitalizedText} />;
+        }
+
+        const domNode = await render(<CapitalizedText text="hello" />);
+        expect(clearYields()).toEqual([
+          "Capitalize 'hello'",
+          0,
+          1,
+          2,
+          // `capitalizedText` only recomputes when the text has changed
+          "Capitalize 'hello, world.'",
+          3,
+          'HELLO, WORLD.',
+        ]);
+        expect(domNode.tagName).toEqual('SPAN');
+        expect(domNode.textContent).toEqual('HELLO, WORLD.');
+      },
+    );
+  });
+
+  describe('useRef', () => {
+    itRenders('basic render', async render => {
+      function Counter(props) {
+        const count = useRef(0);
+        return <span>Count: {count.current}</span>;
+      }
+
+      const domNode = await render(<Counter />);
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+
+    itRenders(
+      'multiple times when updates happen during the render phase',
+      async render => {
+        function Counter(props) {
+          const [count, setCount] = useState(0);
+          const ref = useRef(count);
+
+          if (count < 3) {
+            const newCount = count + 1;
+
+            ref.current = newCount;
+            setCount(newCount);
+          }
+
+          yieldValue(count);
+
+          return <span>Count: {ref.current}</span>;
+        }
+
+        const domNode = await render(<Counter />);
+        expect(clearYields()).toEqual([0, 1, 2, 3]);
+        expect(domNode.textContent).toEqual('Count: 3');
+      },
+    );
+
+    itRenders(
+      'always return the same reference through multiple renders',
+      async render => {
+        let firstRef = null;
+        function Counter(props) {
+          const [count, setCount] = useState(0);
+          const ref = useRef(count);
+          if (firstRef === null) {
+            firstRef = ref;
+          } else if (firstRef !== ref) {
+            throw new Error('should never change');
+          }
+
+          if (count < 3) {
+            setCount(count + 1);
+          } else {
+            firstRef = null;
+          }
+
+          yieldValue(count);
+
+          return <span>Count: {ref.current}</span>;
+        }
+
+        const domNode = await render(<Counter />);
+        expect(clearYields()).toEqual([0, 1, 2, 3]);
+        expect(domNode.textContent).toEqual('Count: 0');
+      },
+    );
+  });
+
+  describe('useEffect', () => {
+    itRenders('should ignore effects on the server', async render => {
+      function Counter(props) {
+        useEffect(() => {
+          yieldValue('should not be invoked');
+        });
+        return <Text text={'Count: ' + props.count} />;
+      }
+      const domNode = await render(<Counter count={0} />);
+      expect(clearYields()).toEqual(['Count: 0']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+  });
+
+  describe('useCallback', () => {
+    itRenders('should ignore callbacks on the server', async render => {
+      function Counter(props) {
+        useCallback(() => {
+          yieldValue('should not be invoked');
+        });
+        return <Text text={'Count: ' + props.count} />;
+      }
+      const domNode = await render(<Counter count={0} />);
+      expect(clearYields()).toEqual(['Count: 0']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+  });
+
+  describe('useAPI', () => {
+    it('should not be invoked on the server', async () => {
+      function Counter(props, ref) {
+        useAPI(ref, () => {
+          throw new Error('should not be invoked');
+        });
+        return <Text text={props.label + ': ' + ref.current} />;
+      }
+      Counter = forwardRef(Counter);
+      const counter = React.createRef();
+      counter.current = 0;
+      const domNode = await serverRender(
+        <Counter label="Count" ref={counter} />,
+      );
+      expect(clearYields()).toEqual(['Count: 0']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+  });
+
+  describe('useMutationEffect', () => {
+    it('should warn when invoked during render', async () => {
+      function Counter() {
+        useMutationEffect(() => {
+          throw new Error('should not be invoked');
+        });
+
+        return <Text text="Count: 0" />;
+      }
+      const domNode = await serverRender(<Counter />, 1);
+      expect(clearYields()).toEqual(['Count: 0']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+  });
+
+  describe('useLayoutEffect', () => {
+    it('should warn when invoked during render', async () => {
+      function Counter() {
+        useLayoutEffect(() => {
+          throw new Error('should not be invoked');
+        });
+
+        return <Text text="Count: 0" />;
+      }
+      const domNode = await serverRender(<Counter />, 1);
+      expect(clearYields()).toEqual(['Count: 0']);
+      expect(domNode.tagName).toEqual('SPAN');
+      expect(domNode.textContent).toEqual('Count: 0');
+    });
+  });
+
+  describe('useContext', () => {
+    itRenders(
+      'can use the same context multiple times in the same function',
+      async render => {
+        const Context = React.createContext(
+          {foo: 0, bar: 0, baz: 0},
+          (a, b) => {
+            let result = 0;
+            if (a.foo !== b.foo) {
+              result |= 0b001;
+            }
+            if (a.bar !== b.bar) {
+              result |= 0b010;
+            }
+            if (a.baz !== b.baz) {
+              result |= 0b100;
+            }
+            return result;
+          },
+        );
+
+        function Provider(props) {
+          return (
+            <Context.Provider
+              value={{foo: props.foo, bar: props.bar, baz: props.baz}}>
+              {props.children}
+            </Context.Provider>
+          );
+        }
+
+        function FooAndBar() {
+          const {foo} = useContext(Context, 0b001);
+          const {bar} = useContext(Context, 0b010);
+          return <Text text={`Foo: ${foo}, Bar: ${bar}`} />;
+        }
+
+        function Baz() {
+          const {baz} = useContext(Context, 0b100);
+          return <Text text={'Baz: ' + baz} />;
+        }
+
+        class Indirection extends React.Component {
+          shouldComponentUpdate() {
+            return false;
+          }
+          render() {
+            return this.props.children;
+          }
+        }
+
+        function App(props) {
+          return (
+            <div>
+              <Provider foo={props.foo} bar={props.bar} baz={props.baz}>
+                <Indirection>
+                  <Indirection>
+                    <FooAndBar />
+                  </Indirection>
+                  <Indirection>
+                    <Baz />
+                  </Indirection>
+                </Indirection>
+              </Provider>
+            </div>
+          );
+        }
+
+        const domNode = await render(<App foo={1} bar={3} baz={5} />);
+        expect(clearYields()).toEqual(['Foo: 1, Bar: 3', 'Baz: 5']);
+        expect(domNode.childNodes.length).toBe(2);
+        expect(domNode.firstChild.tagName).toEqual('SPAN');
+        expect(domNode.firstChild.textContent).toEqual('Foo: 1, Bar: 3');
+        expect(domNode.lastChild.tagName).toEqual('SPAN');
+        expect(domNode.lastChild.textContent).toEqual('Baz: 5');
+      },
+    );
+  });
+});

packages/react-dom/src/server/ReactPartialRenderer.js

@@ -48,6 +48,11 @@ import {
   createMarkupForRoot,
 } from './DOMMarkupOperations';
 import escapeTextForBrowser from './escapeTextForBrowser';
+import {
+  prepareToUseHooks,
+  finishHooks,
+  Dispatcher,
+} from './ReactPartialRendererHooks';
 import {
   Namespaces,
   getIntrinsicNamespace,
@@ -87,15 +92,6 @@ let pushCurrentDebugStack = (stack: Array<Frame>) => {};
 let pushElementToDebugStack = (element: ReactElement) => {};
 let popCurrentDebugStack = () => {};
 
-let Dispatcher = {
-  readContext<T>(
-    context: ReactContext<T>,
-    observedBits: void | number | boolean,
-  ): T {
-    return context._currentValue;
-  },
-};
-
 if (__DEV__) {
   ReactDebugCurrentFrame = ReactSharedInternals.ReactDebugCurrentFrame;
 
@@ -573,7 +569,11 @@ function resolve(
           }
         }
       }
+      const componentIdentity = {};
+      prepareToUseHooks(componentIdentity);
       inst = Component(element.props, publicContext, updater);
+      inst = finishHooks(Component, element.props, inst, publicContext);
+
       if (inst == null || inst.render == null) {
         child = inst;
         validateRenderResult(child, Component);
@@ -985,9 +985,17 @@ class ReactDOMServerRenderer {
         switch (elementType.$$typeof) {
           case REACT_FORWARD_REF_TYPE: {
             const element: ReactElement = ((nextChild: any): ReactElement);
-            const nextChildren = toArray(
-              elementType.render(element.props, element.ref),
+            let nextChildren;
+            const componentIdentity = {};
+            prepareToUseHooks(componentIdentity);
+            nextChildren = elementType.render(element.props, element.ref);
+            nextChildren = finishHooks(
+              elementType.render,
+              element.props,
+              nextChildren,
+              element.ref,
             );
+            nextChildren = toArray(nextChildren);
             const frame: Frame = {
               type: null,
               domNamespace: parentNamespace,

packages/react-dom/src/server/ReactPartialRendererHooks.js

@@ -0,0 +1,366 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+import type {ReactContext} from 'shared/ReactTypes';
+
+import invariant from 'shared/invariant';
+import warning from 'shared/warning';
+
+type BasicStateAction<S> = S | (S => S);
+type MaybeCallback<S> = void | null | (S => mixed);
+type Dispatch<S, A> = (A, MaybeCallback<S>) => void;
+
+type Update<S, A> = {
+  action: A,
+  next: Update<S, A> | null,
+};
+
+type UpdateQueue<S, A> = {
+  last: Update<S, A> | null,
+  dispatch: any,
+};
+
+type Hook = {
+  memoizedState: any,
+  queue: UpdateQueue<any, any> | null,
+  next: Hook | null,
+};
+
+let currentlyRenderingComponent: Object | null = null;
+let firstWorkInProgressHook: Hook | null = null;
+let workInProgressHook: Hook | null = null;
+// Whether the work-in-progress hook is a re-rendered hook
+let isReRender: boolean = false;
+// Whether an update was scheduled during the currently executing render pass.
+let didScheduleRenderPhaseUpdate: boolean = false;
+// Lazily created map of render-phase updates
+let renderPhaseUpdates: Map<
+  UpdateQueue<any, any>,
+  Update<any, any>,
+> | null = null;
+// Counter to prevent infinite loops.
+let numberOfReRenders: number = 0;
+const RE_RENDER_LIMIT = 25;
+
+function resolveCurrentlyRenderingComponent(): Object {
+  invariant(
+    currentlyRenderingComponent !== null,
+    'Hooks can only be called inside the body of a functional component.',
+  );
+  return currentlyRenderingComponent;
+}
+
+function createHook(): Hook {
+  return {
+    memoizedState: null,
+    queue: null,
+    next: null,
+  };
+}
+
+function createWorkInProgressHook(): Hook {
+  if (workInProgressHook === null) {
+    // This is the first hook in the list
+    if (firstWorkInProgressHook === null) {
+      isReRender = false;
+      firstWorkInProgressHook = workInProgressHook = createHook();
+    } else {
+      // There's already a work-in-progress. Reuse it.
+      isReRender = true;
+      workInProgressHook = firstWorkInProgressHook;
+    }
+  } else {
+    if (workInProgressHook.next === null) {
+      isReRender = false;
+      // Append to the end of the list
+      workInProgressHook = workInProgressHook.next = createHook();
+    } else {
+      // There's already a work-in-progress. Reuse it.
+      isReRender = true;
+      workInProgressHook = workInProgressHook.next;
+    }
+  }
+  return workInProgressHook;
+}
+
+export function prepareToUseHooks(componentIdentity: Object): void {
+  currentlyRenderingComponent = componentIdentity;
+
+  // The following should have already been reset
+  // didScheduleRenderPhaseUpdate = false;
+  // firstWorkInProgressHook = null;
+  // numberOfReRenders = 0;
+  // renderPhaseUpdates = null;
+  // workInProgressHook = null;
+}
+
+export function finishHooks(
+  Component: any,
+  props: any,
+  children: any,
+  refOrContext: any,
+): any {
+  // This must be called after every functional component to prevent hooks from
+  // being used in classes.
+
+  while (didScheduleRenderPhaseUpdate) {
+    // Updates were scheduled during the render phase. They are stored in
+    // the `renderPhaseUpdates` map. Call the component again, reusing the
+    // work-in-progress hooks and applying the additional updates on top. Keep
+    // restarting until no more updates are scheduled.
+    didScheduleRenderPhaseUpdate = false;
+    numberOfReRenders += 1;
+
+    // Start over from the beginning of the list
+    workInProgressHook = null;
+
+    children = Component(props, refOrContext);
+  }
+  currentlyRenderingComponent = null;
+  firstWorkInProgressHook = null;
+  numberOfReRenders = 0;
+  renderPhaseUpdates = null;
+  workInProgressHook = null;
+
+  // These were reset above
+  // currentlyRenderingComponent = null;
+  // didScheduleRenderPhaseUpdate = false;
+  // firstWorkInProgressHook = null;
+  // numberOfReRenders = 0;
+  // renderPhaseUpdates = null;
+  // workInProgressHook = null;
+
+  return children;
+}
+
+function useContext<T>(
+  context: ReactContext<T>,
+  observedBits: void | number | boolean,
+): T {
+  return context._currentValue;
+}
+
+function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
+  return typeof action === 'function' ? action(state) : action;
+}
+
+export function useState<S>(
+  initialState: S | (() => S),
+): [S, Dispatch<S, BasicStateAction<S>>] {
+  return useReducer(
+    basicStateReducer,
+    // useReducer has a special case to support lazy useState initializers
+    (initialState: any),
+  );
+}
+
+export function useReducer<S, A>(
+  reducer: (S, A) => S,
+  initialState: S,
+  initialAction: A | void | null,
+): [S, Dispatch<S, A>] {
+  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  workInProgressHook = createWorkInProgressHook();
+  if (isReRender) {
+    // This is a re-render. Apply the new render phase updates to the previous
+    // current hook.
+    const queue: UpdateQueue<S, A> = (workInProgressHook.queue: any);
+    const dispatch: Dispatch<S, A> = (queue.dispatch: any);
+    if (renderPhaseUpdates !== null) {
+      // Render phase updates are stored in a map of queue -> linked list
+      const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
+      if (firstRenderPhaseUpdate !== undefined) {
+        renderPhaseUpdates.delete(queue);
+        let newState = workInProgressHook.memoizedState;
+        let update = firstRenderPhaseUpdate;
+        do {
+          // Process this render phase update. We don't have to check the
+          // priority because it will always be the same as the current
+          // render's.
+          const action = update.action;
+          newState = reducer(newState, action);
+          update = update.next;
+        } while (update !== null);
+
+        workInProgressHook.memoizedState = newState;
+
+        return [newState, dispatch];
+      }
+    }
+    return [workInProgressHook.memoizedState, dispatch];
+  } else {
+    if (reducer === basicStateReducer) {
+      // Special case for `useState`.
+      if (typeof initialState === 'function') {
+        initialState = initialState();
+      }
+    } else if (initialAction !== undefined && initialAction !== null) {
+      initialState = reducer(initialState, initialAction);
+    }
+    workInProgressHook.memoizedState = initialState;
+    const queue: UpdateQueue<S, A> = (workInProgressHook.queue = {
+      last: null,
+      dispatch: null,
+    });
+    const dispatch: Dispatch<S, A> = (queue.dispatch = (dispatchAction.bind(
+      null,
+      currentlyRenderingComponent,
+      queue,
+    ): any));
+    return [workInProgressHook.memoizedState, dispatch];
+  }
+}
+
+function useMemo<T>(
+  nextCreate: () => T,
+  inputs: Array<mixed> | void | null,
+): T {
+  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  workInProgressHook = createWorkInProgressHook();
+
+  const nextInputs =
+    inputs !== undefined && inputs !== null ? inputs : [nextCreate];
+
+  if (
+    workInProgressHook !== null &&
+    workInProgressHook.memoizedState !== null
+  ) {
+    const prevState = workInProgressHook.memoizedState;
+    const prevInputs = prevState[1];
+    if (inputsAreEqual(nextInputs, prevInputs)) {
+      return prevState[0];
+    }
+  }
+
+  const nextValue = nextCreate();
+  workInProgressHook.memoizedState = [nextValue, nextInputs];
+  return nextValue;
+}
+
+function useRef<T>(initialValue: T): {current: T} {
+  currentlyRenderingComponent = resolveCurrentlyRenderingComponent();
+  workInProgressHook = createWorkInProgressHook();
+  const previousRef = workInProgressHook.memoizedState;
+  if (previousRef === null) {
+    const ref = {current: initialValue};
+    if (__DEV__) {
+      Object.seal(ref);
+    }
+    workInProgressHook.memoizedState = ref;
+    return ref;
+  } else {
+    return previousRef;
+  }
+}
+
+function useMutationEffect(
+  create: () => mixed,
+  inputs: Array<mixed> | void | null,
+) {
+  warning(
+    false,
+    'useMutationEffect does nothing on the server, because its effect cannot ' +
+      "be encoded into the server renderer's output format. This will lead " +
+      'to a mismatch between the initial, non-hydrated UI and the intended ' +
+      'UI. To avoid this, useMutationEffect should only be used in ' +
+      'components that render exclusively on the client.',
+  );
+}
+
+export function useLayoutEffect(
+  create: () => mixed,
+  inputs: Array<mixed> | void | null,
+) {
+  warning(
+    false,
+    'useLayoutEffect does nothing on the server, because its effect cannot ' +
+      "be encoded into the server renderer's output format. This will lead " +
+      'to a mismatch between the initial, non-hydrated UI and the intended ' +
+      'UI. To avoid this, useLayoutEffect should only be used in ' +
+      'components that render exclusively on the client.',
+  );
+}
+
+function dispatchAction<S, A>(
+  componentIdentity: Object,
+  queue: UpdateQueue<S, A>,
+  action: A,
+) {
+  invariant(
+    numberOfReRenders < RE_RENDER_LIMIT,
+    'Too many re-renders. React limits the number of renders to prevent ' +
+      'an infinite loop.',
+  );
+
+  if (componentIdentity === currentlyRenderingComponent) {
+    // This is a render phase update. Stash it in a lazily-created map of
+    // queue -> linked list of updates. After this render pass, we'll restart
+    // and apply the stashed updates on top of the work-in-progress hook.
+    didScheduleRenderPhaseUpdate = true;
+    const update: Update<S, A> = {
+      action,
+      next: null,
+    };
+    if (renderPhaseUpdates === null) {
+      renderPhaseUpdates = new Map();
+    }
+    const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
+    if (firstRenderPhaseUpdate === undefined) {
+      renderPhaseUpdates.set(queue, update);
+    } else {
+      // Append the update to the end of the list.
+      let lastRenderPhaseUpdate = firstRenderPhaseUpdate;
+      while (lastRenderPhaseUpdate.next !== null) {
+        lastRenderPhaseUpdate = lastRenderPhaseUpdate.next;
+      }
+      lastRenderPhaseUpdate.next = update;
+    }
+  } else {
+    // This means an update has happened after the functional component has
+    // returned. On the server this is a no-op. In React Fiber, the update
+    // would be scheduled for a future render.
+  }
+}
+
+function inputsAreEqual(arr1, arr2) {
+  // Don't bother comparing lengths because these arrays are always
+  // passed inline.
+  for (let i = 0; i < arr1.length; i++) {
+    // Inlined Object.is polyfill.
+    // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Object/is
+    const val1 = arr1[i];
+    const val2 = arr2[i];
+    if (
+      (val1 === val2 && (val1 !== 0 || 1 / val1 === 1 / (val2: any))) ||
+      (val1 !== val1 && val2 !== val2) // eslint-disable-line no-self-compare
+    ) {
+      continue;
+    }
+    return false;
+  }
+  return true;
+}
+
+function noop(): void {}
+
+export const Dispatcher = {
+  readContext: useContext,
+  useContext,
+  useMemo,
+  useReducer,
+  useRef,
+  useState,
+  useMutationEffect,
+  useLayoutEffect,
+  // useAPI is not run in the server environment
+  useAPI: noop,
+  // Callbacks are not run in the server environment.
+  useCallback: noop,
+  // Effects are not run in the server environment.
+  useEffect: noop,
+};

packages/react-reconciler/src/ReactFiberHooks.js

@@ -347,47 +347,48 @@ export function useReducer<S, A>(
 ): [S, Dispatch<S, A>] {
   currentlyRenderingFiber = resolveCurrentlyRenderingFiber();
   workInProgressHook = createWorkInProgressHook();
-  if (isReRender) {
-    // This is a re-render. Apply the new render phase updates to the previous
-    // work-in-progress hook.
-    const queue: UpdateQueue<S, A> = (workInProgressHook.queue: any);
-    const dispatch: Dispatch<S, A> = (queue.dispatch: any);
-    if (renderPhaseUpdates !== null) {
-      // Render phase updates are stored in a map of queue -> linked list
-      const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
-      if (firstRenderPhaseUpdate !== undefined) {
-        renderPhaseUpdates.delete(queue);
-        let newState = workInProgressHook.memoizedState;
-        let update = firstRenderPhaseUpdate;
-        do {
-          // Process this render phase update. We don't have to check the
-          // priority because it will always be the same as the current
-          // render's.
-          const action = update.action;
-          newState = reducer(newState, action);
-          const callback = update.callback;
-          if (callback !== null) {
-            pushCallback(currentlyRenderingFiber, update);
+  let queue: UpdateQueue<S, A> | null = (workInProgressHook.queue: any);
+  if (queue !== null) {
+    // Already have a queue, so this is an update.
+    if (isReRender) {
+      // This is a re-render. Apply the new render phase updates to the previous
+      // work-in-progress hook.
+      const dispatch: Dispatch<S, A> = (queue.dispatch: any);
+      if (renderPhaseUpdates !== null) {
+        // Render phase updates are stored in a map of queue -> linked list
+        const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
+        if (firstRenderPhaseUpdate !== undefined) {
+          renderPhaseUpdates.delete(queue);
+          let newState = workInProgressHook.memoizedState;
+          let update = firstRenderPhaseUpdate;
+          do {
+            // Process this render phase update. We don't have to check the
+            // priority because it will always be the same as the current
+            // render's.
+            const action = update.action;
+            newState = reducer(newState, action);
+            const callback = update.callback;
+            if (callback !== null) {
+              pushCallback(currentlyRenderingFiber, update);
+            }
+            update = update.next;
+          } while (update !== null);
+
+          workInProgressHook.memoizedState = newState;
+
+          // Don't persist the state accumlated from the render phase updates to
+          // the base state unless the queue is empty.
+          // TODO: Not sure if this is the desired semantics, but it's what we
+          // do for gDSFP. I can't remember why.
+          if (workInProgressHook.baseUpdate === queue.last) {
+            workInProgressHook.baseState = newState;
           }
-          update = update.next;
-        } while (update !== null);
 
-        workInProgressHook.memoizedState = newState;
-
-        // Don't persist the state accumlated from the render phase updates to
-        // the base state unless the queue is empty.
-        // TODO: Not sure if this is the desired semantics, but it's what we
-        // do for gDSFP. I can't remember why.
-        if (workInProgressHook.baseUpdate === queue.last) {
-          workInProgressHook.baseState = newState;
+          return [newState, dispatch];
         }
-
-        return [newState, dispatch];
       }
+      return [workInProgressHook.memoizedState, dispatch];
     }
-    return [workInProgressHook.memoizedState, dispatch];
-  } else if (currentHook !== null) {
-    const queue: UpdateQueue<S, A> = (workInProgressHook.queue: any);
 
     // The last update in the entire queue
     const last = queue.last;
@@ -457,27 +458,28 @@ export function useReducer<S, A>(
 
     const dispatch: Dispatch<S, A> = (queue.dispatch: any);
     return [workInProgressHook.memoizedState, dispatch];
-  } else {
-    if (reducer === basicStateReducer) {
-      // Special case for `useState`.
-      if (typeof initialState === 'function') {
-        initialState = initialState();
-      }
-    } else if (initialAction !== undefined && initialAction !== null) {
-      initialState = reducer(initialState, initialAction);
-    }
-    workInProgressHook.memoizedState = workInProgressHook.baseState = initialState;
-    const queue: UpdateQueue<S, A> = (workInProgressHook.queue = {
-      last: null,
-      dispatch: null,
-    });
-    const dispatch: Dispatch<S, A> = (queue.dispatch = (dispatchAction.bind(
-      null,
-      currentlyRenderingFiber,
-      queue,
-    ): any));
-    return [workInProgressHook.memoizedState, dispatch];
   }
+
+  // There's no existing queue, so this is the initial render.
+  if (reducer === basicStateReducer) {
+    // Special case for `useState`.
+    if (typeof initialState === 'function') {
+      initialState = initialState();
+    }
+  } else if (initialAction !== undefined && initialAction !== null) {
+    initialState = reducer(initialState, initialAction);
+  }
+  workInProgressHook.memoizedState = workInProgressHook.baseState = initialState;
+  queue = workInProgressHook.queue = {
+    last: null,
+    dispatch: null,
+  };
+  const dispatch: Dispatch<S, A> = (queue.dispatch = (dispatchAction.bind(
+    null,
+    currentlyRenderingFiber,
+    queue,
+  ): any));
+  return [workInProgressHook.memoizedState, dispatch];
 }
 
 function pushCallback(workInProgress: Fiber, update: Update<any, any>): void {
@@ -525,7 +527,8 @@ export function useRef<T>(initialValue: T): {current: T} {
   currentlyRenderingFiber = resolveCurrentlyRenderingFiber();
   workInProgressHook = createWorkInProgressHook();
   let ref;
-  if (currentHook === null) {
+
+  if (workInProgressHook.memoizedState === null) {
     ref = {current: initialValue};
     if (__DEV__) {
       Object.seal(ref);
@@ -637,14 +640,13 @@ export function useCallback<T>(
   const nextInputs =
     inputs !== undefined && inputs !== null ? inputs : [callback];
 
-  if (currentHook !== null) {
-    const prevState = currentHook.memoizedState;
+  const prevState = workInProgressHook.memoizedState;
+  if (prevState !== null) {
     const prevInputs = prevState[1];
     if (inputsAreEqual(nextInputs, prevInputs)) {
       return prevState[0];
     }
   }
-
   workInProgressHook.memoizedState = [callback, nextInputs];
   return callback;
 }
@@ -659,8 +661,8 @@ export function useMemo<T>(
   const nextInputs =
     inputs !== undefined && inputs !== null ? inputs : [nextCreate];
 
-  if (currentHook !== null) {
-    const prevState = currentHook.memoizedState;
+  const prevState = workInProgressHook.memoizedState;
+  if (prevState !== null) {
     const prevInputs = prevState[1];
     if (inputsAreEqual(nextInputs, prevInputs)) {
       return prevState[0];

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

@@ -1417,6 +1417,33 @@ describe('ReactHooks', () => {
       ReactNoop.render(<LazyCompute compute={computeB} />);
       expect(ReactNoop.flush()).toEqual(['compute B', 'B']);
     });
+
+    it('should not invoke memoized function during re-renders unless inputs change', () => {
+      function LazyCompute(props) {
+        const computed = useMemo(() => props.compute(props.input), [
+          props.input,
+        ]);
+        const [count, setCount] = useState(0);
+        if (count < 3) {
+          setCount(count + 1);
+        }
+        return <Text text={computed} />;
+      }
+
+      function compute(val) {
+        ReactNoop.yield('compute ' + val);
+        return val;
+      }
+
+      ReactNoop.render(<LazyCompute compute={compute} input="A" />);
+      expect(ReactNoop.flush()).toEqual(['compute A', 'A']);
+
+      ReactNoop.render(<LazyCompute compute={compute} input="A" />);
+      expect(ReactNoop.flush()).toEqual(['A']);
+
+      ReactNoop.render(<LazyCompute compute={compute} input="B" />);
+      expect(ReactNoop.flush()).toEqual(['compute B', 'B']);
+    });
   });
 
   describe('useRef', () => {
@@ -1476,6 +1503,30 @@ describe('ReactHooks', () => {
       jest.advanceTimersByTime(20);
       expect(ReactNoop.flush()).toEqual(['ping: 6']);
     });
+
+    it('should return the same ref during re-renders', () => {
+      function Counter() {
+        const ref = useRef('val');
+        const [count, setCount] = useState(0);
+        const [firstRef] = useState(ref);
+
+        if (firstRef !== ref) {
+          throw new Error('should never change');
+        }
+
+        if (count < 3) {
+          setCount(count + 1);
+        }
+
+        return <Text text={ref.current} />;
+      }
+
+      ReactNoop.render(<Counter />);
+      expect(ReactNoop.flush()).toEqual(['val']);
+
+      ReactNoop.render(<Counter />);
+      expect(ReactNoop.flush()).toEqual(['val']);
+    });
   });
 
   describe('progressive enhancement', () => {