From 51c2690c357a9264ff36ff8a0c11fb5afc3eba74 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <dg@domgan.com>
Date: Thu, 10 Jan 2019 15:58:58 +0000
Subject: [PATCH 1/6] Add hook support to ReactShallowRenderer

---
 .../src/ReactShallowRenderer.js               | 279 ++++++++++++++++++
 .../ReactShallowRendererHooks-test.js         | 273 +++++++++++++++++
 2 files changed, 552 insertions(+)
 create mode 100644 packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index 9bf373d033caf..ef1ff6d9ab3d7 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -13,6 +13,13 @@ import getComponentName from 'shared/getComponentName';
 import shallowEqual from 'shared/shallowEqual';
 import invariant from 'shared/invariant';
 import checkPropTypes from 'prop-types/checkPropTypes';
+import ReactSharedInternals from 'shared/ReactSharedInternals';
+import {enableHooks} from 'shared/ReactFeatureFlags';
+import areHookInputsEqual from 'shared/areHookInputsEqual';
+
+const {ReactCurrentDispatcher} = ReactSharedInternals;
+
+const RE_RENDER_LIMIT = 25;
 
 const emptyObject = {};
 if (__DEV__) {
@@ -85,6 +92,18 @@ class Updater {
   }
 }
 
+function createHook(): Hook {
+  return {
+    memoizedState: null,
+    queue: null,
+    next: null,
+  };
+}
+
+function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {
+  return typeof action === 'function' ? action(state) : action;
+}
+
 class ReactShallowRenderer {
   static createRenderer = function() {
     return new ReactShallowRenderer();
@@ -99,6 +118,259 @@ class ReactShallowRenderer {
     this._rendering = false;
     this._forcedUpdate = false;
     this._updater = new Updater(this);
+    if (enableHooks) {
+      this._dispatcher = ReactCurrentDispatcher.current = this._createDispatcher();
+      this._workInProgressHook = null;
+      this._firstWorkInProgressHook = null;
+      this._isReRender = false;
+      this._didScheduleRenderPhaseUpdate = false;
+      this._renderPhaseUpdates = null;
+      this._currentlyRenderingComponent = null;
+      this._numberOfReRenders = 0;
+    }
+  }
+
+  _validateCurrentlyRenderingComponent() {
+    invariant(
+      this._currentlyRenderingComponent !== null,
+      'Hooks can only be called inside the body of a function component.',
+    );
+  }
+
+  _createDispatcher() {
+    const useReducer = <S, A>(
+      reducer: (S, A) => S,
+      initialState: S,
+      initialAction: A | void | null,
+    ): [S, Dispatch<A>] => {
+      this._validateCurrentlyRenderingComponent();
+      this._createWorkInProgressHook();
+      if (this._isReRender) {
+        // This is a re-render. Apply the new render phase updates to the previous
+        // current hook.
+        const queue: UpdateQueue<A> = (this._workInProgressHook.queue: any);
+        const dispatch: Dispatch<A> = (queue.dispatch: any);
+        if (this._renderPhaseUpdates !== null) {
+          // Render phase updates are stored in a map of queue -> linked list
+          const firstRenderPhaseUpdate = this._renderPhaseUpdates.get(queue);
+          if (firstRenderPhaseUpdate !== undefined) {
+            this._renderPhaseUpdates.delete(queue);
+            let newState = this._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);
+
+            this._workInProgressHook.memoizedState = newState;
+
+            return [newState, dispatch];
+          }
+        }
+        return [this._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);
+        }
+        this._workInProgressHook.memoizedState = initialState;
+        const queue: UpdateQueue<A> = (this._workInProgressHook.queue = {
+          last: null,
+          dispatch: null,
+        });
+        const dispatch: Dispatch<
+          A,
+        > = (queue.dispatch = (this._dispatchAction.bind(
+          this,
+          this._currentlyRenderingComponent,
+          queue,
+        ): any));
+        return [this._workInProgressHook.memoizedState, dispatch];
+      }
+    };
+
+    const useState = <S>(
+      initialState: (() => S) | S,
+    ): [S, Dispatch<BasicStateAction<S>>] => {
+      return useReducer(
+        basicStateReducer,
+        // useReducer has a special case to support lazy useState initializers
+        (initialState: any),
+      );
+    };
+
+    const useMemo = <T>(
+      nextCreate: () => T,
+      inputs: Array<mixed> | void | null,
+    ): T => {
+      this._validateCurrentlyRenderingComponent();
+      this._createWorkInProgressHook();
+
+      const nextInputs =
+        inputs !== undefined && inputs !== null ? inputs : [nextCreate];
+
+      if (
+        this._workInProgressHook !== null &&
+        this._workInProgressHook.memoizedState !== null
+      ) {
+        const prevState = this._workInProgressHook.memoizedState;
+        const prevInputs = prevState[1];
+        if (areHookInputsEqual(nextInputs, prevInputs)) {
+          return prevState[0];
+        }
+      }
+
+      const nextValue = nextCreate();
+      this._workInProgressHook.memoizedState = [nextValue, nextInputs];
+      return nextValue;
+    };
+
+    const useRef = <T>(initialValue: T): {current: T} => {
+      this._validateCurrentlyRenderingComponent();
+      this._createWorkInProgressHook();
+      const previousRef = this._workInProgressHook.memoizedState;
+      if (previousRef === null) {
+        const ref = {current: initialValue};
+        if (__DEV__) {
+          Object.seal(ref);
+        }
+        this._workInProgressHook.memoizedState = ref;
+        return ref;
+      } else {
+        return previousRef;
+      }
+    };
+
+    const readContext = <T>(
+      context: ReactContext<T>,
+      observedBits: void | number | boolean,
+    ): T => {
+      return context._currentValue;
+    };
+
+    const noOp = () => {
+      this._validateCurrentlyRenderingComponent();
+    };
+
+    const identity = (fn: Function): Function => {
+      return fn;
+    };
+
+    return {
+      readContext,
+      useCallback: identity,
+      useContext: context => {
+        this._validateCurrentlyRenderingComponent();
+        return readContext(context);
+      },
+      useEffect: noOp,
+      useImperativeMethods: noOp,
+      useLayoutEffect: noOp,
+      useMemo,
+      useReducer,
+      useRef,
+      useState,
+    };
+  }
+
+  _dispatchAction<A>(
+    componentIdentity: Object,
+    queue: UpdateQueue<A>,
+    action: A,
+  ) {
+    invariant(
+      this._numberOfReRenders < RE_RENDER_LIMIT,
+      'Too many re-renders. React limits the number of renders to prevent ' +
+        'an infinite loop.',
+    );
+
+    if (componentIdentity === this._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.
+      this._didScheduleRenderPhaseUpdate = true;
+      const update: Update<A> = {
+        action,
+        next: null,
+      };
+      if (this._renderPhaseUpdates === null) {
+        this._renderPhaseUpdates = new Map();
+      }
+      const firstRenderPhaseUpdate = this._renderPhaseUpdates.get(queue);
+      if (firstRenderPhaseUpdate === undefined) {
+        this._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 function component has
+      // returned. On the server this is a no-op. In React Fiber, the update
+      // would be scheduled for a future render.
+    }
+  }
+
+  _createWorkInProgressHook(): Hook {
+    if (this._workInProgressHook === null) {
+      // This is the first hook in the list
+      if (this._firstWorkInProgressHook === null) {
+        this._isReRender = false;
+        this._firstWorkInProgressHook = this._workInProgressHook = createHook();
+      } else {
+        // There's already a work-in-progress. Reuse it.
+        this._isReRender = true;
+        this._workInProgressHook = this._firstWorkInProgressHook;
+      }
+    } else {
+      if (this._workInProgressHook.next === null) {
+        this._isReRender = false;
+        // Append to the end of the list
+        this._workInProgressHook = this._workInProgressHook.next = createHook();
+      } else {
+        // There's already a work-in-progress. Reuse it.
+        this._isReRender = true;
+        this._workInProgressHook = this._workInProgressHook.next;
+      }
+    }
+    return this._workInProgressHook;
+  }
+
+  _prepareToUseHooks(componentIdentity) {
+    this._currentlyRenderingComponent = componentIdentity;
+  }
+
+  _finishHooks(element, context) {
+    this._hooksEnabled = false;
+    if (this._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.
+      this._didScheduleRenderPhaseUpdate = false;
+      this._numberOfReRenders += 1;
+
+      // Start over from the beginning of the list
+      this._workInProgressHook = null;
+      this._rendering = false;
+      this.render(element, context);
+    } else {
+      this._currentlyRenderingComponent = null;
+      this._workInProgressHook = null;
+      this._renderPhaseUpdates = null;
+      this._numberOfReRenders = 0;
+    }
   }
 
   getMountedInstance() {
@@ -175,11 +447,17 @@ class ReactShallowRenderer {
 
         this._mountClassComponent(element, this._context);
       } else {
+        if (enableHooks) {
+          this._prepareToUseHooks(element.type);
+        }
         this._rendered = element.type.call(
           undefined,
           element.props,
           this._context,
         );
+        if (enableHooks) {
+          this._finishHooks(element, context);
+        }
       }
     }
 
@@ -196,6 +474,7 @@ class ReactShallowRenderer {
       }
     }
 
+    this._firstWorkInProgressHook = null;
     this._context = null;
     this._element = null;
     this._newState = null;
diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js b/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js
new file mode 100644
index 0000000000000..7f0dcc687af4e
--- /dev/null
+++ b/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js
@@ -0,0 +1,273 @@
+/**
+ * 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.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+
+'use strict';
+
+let createRenderer;
+let React;
+
+describe('ReactShallowRenderer with hooks', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    let ReactFeatureFlags = require('shared/ReactFeatureFlags');
+    // TODO: Switch this test to non-internal once the flag is on by default.
+    ReactFeatureFlags.enableHooks = true;
+    createRenderer = require('react-test-renderer/shallow').createRenderer;
+    React = require('react');
+  });
+
+  it('should work with useState', () => {
+    function SomeComponent({defaultName}) {
+      const [name] = React.useState(defaultName);
+
+      return (
+        <div>
+          <p>
+            Your name is: <span>{name}</span>
+          </p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(
+      <SomeComponent defaultName={'Dominic'} />,
+    );
+
+    expect(result).toEqual(
+      <div>
+        <p>
+          Your name is: <span>Dominic</span>
+        </p>
+      </div>,
+    );
+
+    result = shallowRenderer.render(
+      <SomeComponent defaultName={'Should not use this name'} />,
+    );
+
+    expect(result).toEqual(
+      <div>
+        <p>
+          Your name is: <span>Dominic</span>
+        </p>
+      </div>,
+    );
+  });
+
+  it('should work with updating a value from useState', () => {
+    function SomeComponent({defaultName}) {
+      const [name, updateName] = React.useState(defaultName);
+
+      if (name !== 'Dan') {
+        updateName('Dan');
+      }
+
+      return (
+        <div>
+          <p>
+            Your name is: <span>{name}</span>
+          </p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(
+      <SomeComponent defaultName={'Dominic'} />,
+    );
+
+    expect(result).toEqual(
+      <div>
+        <p>
+          Your name is: <span>Dan</span>
+        </p>
+      </div>,
+    );
+  });
+
+  it('should work with useReducer', () => {
+    const initialState = {count: 0};
+
+    function reducer(state, action) {
+      switch (action.type) {
+        case 'reset':
+          return initialState;
+        case 'increment':
+          return {count: state.count + 1};
+        case 'decrement':
+          return {count: state.count - 1};
+        default:
+          return state;
+      }
+    }
+
+    function SomeComponent({initialCount}) {
+      const [state] = React.useReducer(reducer, {count: initialCount});
+
+      return (
+        <div>
+          <p>
+            The counter is at: <span>{state.count.toString()}</span>
+          </p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SomeComponent initialCount={0} />);
+
+    expect(result).toEqual(
+      <div>
+        <p>
+          The counter is at: <span>0</span>
+        </p>
+      </div>,
+    );
+
+    result = shallowRenderer.render(<SomeComponent initialCount={10} />);
+
+    expect(result).toEqual(
+      <div>
+        <p>
+          The counter is at: <span>0</span>
+        </p>
+      </div>,
+    );
+  });
+
+  it('should work with a dispatched state change for a useReducer', () => {
+    const initialState = {count: 0};
+
+    function reducer(state, action) {
+      switch (action.type) {
+        case 'reset':
+          return initialState;
+        case 'increment':
+          return {count: state.count + 1};
+        case 'decrement':
+          return {count: state.count - 1};
+        default:
+          return state;
+      }
+    }
+
+    function SomeComponent({initialCount}) {
+      const [state, dispatch] = React.useReducer(reducer, {
+        count: initialCount,
+      });
+
+      if (state.count === 0) {
+        dispatch({type: 'increment'});
+      }
+
+      return (
+        <div>
+          <p>
+            The counter is at: <span>{state.count.toString()}</span>
+          </p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SomeComponent initialCount={0} />);
+
+    expect(result).toEqual(
+      <div>
+        <p>
+          The counter is at: <span>1</span>
+        </p>
+      </div>,
+    );
+  });
+
+  it('should not trigger effects', () => {
+    let effectsCalled = [];
+
+    function SomeComponent({defaultName}) {
+      React.useEffect(() => {
+        effectsCalled.push('useEffect');
+      });
+
+      React.useLayoutEffect(() => {
+        effectsCalled.push('useEffect');
+      });
+
+      return <div>Hello world</div>;
+    }
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent />);
+
+    expect(effectsCalled).toEqual([]);
+  });
+
+  it('should work with useRef', () => {
+    function SomeComponent() {
+      const randomNumberRef = React.useRef({number: Math.random()});
+
+      return (
+        <div>
+          <p>The random number is: {randomNumberRef.current.number}</p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let firstResult = shallowRenderer.render(<SomeComponent />);
+    let secondResult = shallowRenderer.render(<SomeComponent />);
+
+    expect(firstResult).toEqual(secondResult);
+  });
+
+  it('should work with useMemo', () => {
+    function SomeComponent() {
+      const randomNumber = React.useMemo(() => {
+        return {number: Math.random()};
+      }, []);
+
+      return (
+        <div>
+          <p>The random number is: {randomNumber.number}</p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let firstResult = shallowRenderer.render(<SomeComponent />);
+    let secondResult = shallowRenderer.render(<SomeComponent />);
+
+    expect(firstResult).toEqual(secondResult);
+  });
+
+  it('should work with useContext', () => {
+    const SomeContext = React.createContext('default');
+
+    function SomeComponent() {
+      const value = React.useContext(SomeContext);
+
+      return (
+        <div>
+          <p>{value}</p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SomeComponent />);
+
+    expect(result).toEqual(
+      <div>
+        <p>default</p>
+      </div>,
+    );
+  });
+});

From 7f56b048c83caae41d645c300cb66d5156667531 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <dg@domgan.com>
Date: Thu, 10 Jan 2019 16:20:19 +0000
Subject: [PATCH 2/6] Make the test internal

---
 ...erHooks-test.js => ReactShallowRendererHooks-test.internal.js} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename packages/react-test-renderer/src/__tests__/{ReactShallowRendererHooks-test.js => ReactShallowRendererHooks-test.internal.js} (100%)

diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js b/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.internal.js
similarity index 100%
rename from packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.js
rename to packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.internal.js

From 9c2832a72dccd53e00aafa8d0f67629e76b1d112 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <dg@domgan.com>
Date: Thu, 17 Jan 2019 11:40:07 +0000
Subject: [PATCH 3/6] Adds more Flow types + addresses feedback + bug fixes

---
 .../src/ReactShallowRenderer.js               | 189 ++++++++++++++----
 ...ReactShallowRendererHooks-test.internal.js |  47 +++++
 2 files changed, 198 insertions(+), 38 deletions(-)

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index ef1ff6d9ab3d7..fad1a87162693 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -4,6 +4,7 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  *
+ * @flow
  */
 
 import React from 'react';
@@ -15,7 +16,31 @@ import invariant from 'shared/invariant';
 import checkPropTypes from 'prop-types/checkPropTypes';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {enableHooks} from 'shared/ReactFeatureFlags';
-import areHookInputsEqual from 'shared/areHookInputsEqual';
+import warning from 'shared/warning';
+import is from 'shared/objectIs';
+
+import typeof {Dispatcher as DispatcherType} from 'react-reconciler/src/ReactFiberDispatcher';
+import type {ReactContext} from 'shared/ReactTypes';
+import type {ReactElement} from 'shared/ReactElementType';
+
+type BasicStateAction<S> = (S => S) | S;
+type Dispatch<A> = A => void;
+
+type Update<A> = {
+  action: A,
+  next: Update<A> | null,
+};
+
+type UpdateQueue<A> = {
+  last: Update<A> | null,
+  dispatch: any,
+};
+
+type Hook = {
+  memoizedState: any,
+  queue: UpdateQueue<any> | null,
+  next: Hook | null,
+};
 
 const {ReactCurrentDispatcher} = ReactSharedInternals;
 
@@ -26,12 +51,56 @@ if (__DEV__) {
   Object.freeze(emptyObject);
 }
 
+// In DEV, this is the name of the currently executing primitive hook
+let currentHookNameInDev: ?string;
+
+function areHookInputsEqual(
+  nextDeps: Array<mixed>,
+  prevDeps: Array<mixed> | null,
+) {
+  if (prevDeps === null) {
+    warning(
+      false,
+      '%s received a final argument during this render, but not during ' +
+        'the previous render. Even though the final argument is optional, ' +
+        'its type cannot change between renders.',
+      currentHookNameInDev,
+    );
+    return false;
+  }
+
+  // Don't bother comparing lengths in prod because these arrays should be
+  // passed inline.
+  if (nextDeps.length !== prevDeps.length) {
+    warning(
+      false,
+      'The final argument passed to %s changed size between renders. The ' +
+        'order and size of this array must remain constant.\n\n' +
+        'Previous: %s\n' +
+        'Incoming: %s',
+      currentHookNameInDev,
+      `[${nextDeps.join(', ')}]`,
+      `[${prevDeps.join(', ')}]`,
+    );
+  }
+  for (let i = 0; i < prevDeps.length && i < nextDeps.length; i++) {
+    if (is(nextDeps[i], prevDeps[i])) {
+      continue;
+    }
+    return false;
+  }
+  return true;
+}
+
 class Updater {
   constructor(renderer) {
     this._renderer = renderer;
     this._callbacks = [];
   }
 
+  _renderer: ReactShallowRenderer;
+  _callbacks: Array<any>;
+
   _enqueueCallback(callback, publicInstance) {
     if (typeof callback === 'function' && publicInstance) {
       this._callbacks.push({
@@ -119,7 +188,7 @@ class ReactShallowRenderer {
     this._forcedUpdate = false;
     this._updater = new Updater(this);
     if (enableHooks) {
-      this._dispatcher = ReactCurrentDispatcher.current = this._createDispatcher();
+      this._dispatcher = this._createDispatcher();
       this._workInProgressHook = null;
       this._firstWorkInProgressHook = null;
       this._isReRender = false;
@@ -127,9 +196,28 @@ class ReactShallowRenderer {
       this._renderPhaseUpdates = null;
       this._currentlyRenderingComponent = null;
       this._numberOfReRenders = 0;
+      this._previousComponentIdentity = null;
     }
   }
 
+  _context: null | Object;
+  _newState: null | Object;
+  _instance: any;
+  _element: null | ReactElement;
+  _rendered: null | mixed;
+  _updater: Updater;
+  _rendering: boolean;
+  _forcedUpdate: boolean;
+  _dispatcher: DispatcherType;
+  _workInProgressHook: null | Hook;
+  _firstWorkInProgressHook: null | Hook;
+  _currentlyRenderingComponent: null | Object;
+  _previousComponentIdentity: null | Object;
+  _renderPhaseUpdates: Map<UpdateQueue<any>, Update<any>> | null;
+  _isReRender: boolean;
+  _didScheduleRenderPhaseUpdate: boolean;
+  _numberOfReRenders: number;
+
   _validateCurrentlyRenderingComponent() {
     invariant(
       this._currentlyRenderingComponent !== null,
@@ -137,7 +225,7 @@ class ReactShallowRenderer {
     );
   }
 
-  _createDispatcher() {
+  _createDispatcher(): DispatcherType {
     const useReducer = <S, A>(
       reducer: (S, A) => S,
       initialState: S,
@@ -145,17 +233,18 @@ class ReactShallowRenderer {
     ): [S, Dispatch<A>] => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
+      const workInProgressHook: Hook = (this._workInProgressHook: any);
       if (this._isReRender) {
         // This is a re-render. Apply the new render phase updates to the previous
         // current hook.
-        const queue: UpdateQueue<A> = (this._workInProgressHook.queue: any);
+        const queue: UpdateQueue<A> = (workInProgressHook.queue: any);
         const dispatch: Dispatch<A> = (queue.dispatch: any);
         if (this._renderPhaseUpdates !== null) {
           // Render phase updates are stored in a map of queue -> linked list
           const firstRenderPhaseUpdate = this._renderPhaseUpdates.get(queue);
           if (firstRenderPhaseUpdate !== undefined) {
-            this._renderPhaseUpdates.delete(queue);
-            let newState = this._workInProgressHook.memoizedState;
+            (this._renderPhaseUpdates: any).delete(queue);
+            let newState = workInProgressHook.memoizedState;
             let update = firstRenderPhaseUpdate;
             do {
               // Process this render phase update. We don't have to check the
@@ -166,12 +255,12 @@ class ReactShallowRenderer {
               update = update.next;
             } while (update !== null);
 
-            this._workInProgressHook.memoizedState = newState;
+            workInProgressHook.memoizedState = newState;
 
             return [newState, dispatch];
           }
         }
-        return [this._workInProgressHook.memoizedState, dispatch];
+        return [workInProgressHook.memoizedState, dispatch];
       } else {
         if (reducer === basicStateReducer) {
           // Special case for `useState`.
@@ -181,8 +270,8 @@ class ReactShallowRenderer {
         } else if (initialAction !== undefined && initialAction !== null) {
           initialState = reducer(initialState, initialAction);
         }
-        this._workInProgressHook.memoizedState = initialState;
-        const queue: UpdateQueue<A> = (this._workInProgressHook.queue = {
+        workInProgressHook.memoizedState = initialState;
+        const queue: UpdateQueue<A> = (workInProgressHook.queue = {
           last: null,
           dispatch: null,
         });
@@ -190,10 +279,10 @@ class ReactShallowRenderer {
           A,
         > = (queue.dispatch = (this._dispatchAction.bind(
           this,
-          this._currentlyRenderingComponent,
+          (this._currentlyRenderingComponent: any),
           queue,
         ): any));
-        return [this._workInProgressHook.memoizedState, dispatch];
+        return [workInProgressHook.memoizedState, dispatch];
       }
     };
 
@@ -229,20 +318,20 @@ class ReactShallowRenderer {
       }
 
       const nextValue = nextCreate();
-      this._workInProgressHook.memoizedState = [nextValue, nextInputs];
+      (this._workInProgressHook: any).memoizedState = [nextValue, nextInputs];
       return nextValue;
     };
 
     const useRef = <T>(initialValue: T): {current: T} => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
-      const previousRef = this._workInProgressHook.memoizedState;
+      const previousRef = (this._workInProgressHook: any).memoizedState;
       if (previousRef === null) {
         const ref = {current: initialValue};
         if (__DEV__) {
           Object.seal(ref);
         }
-        this._workInProgressHook.memoizedState = ref;
+        (this._workInProgressHook: any).memoizedState = ref;
         return ref;
       } else {
         return previousRef;
@@ -266,13 +355,14 @@ class ReactShallowRenderer {
 
     return {
       readContext,
-      useCallback: identity,
-      useContext: context => {
+      useCallback: (identity: any),
+      useContext: <T>(context: ReactContext<T>): T => {
         this._validateCurrentlyRenderingComponent();
         return readContext(context);
       },
+      useDebugValue: noOp,
       useEffect: noOp,
-      useImperativeMethods: noOp,
+      useImperativeHandle: noOp,
       useLayoutEffect: noOp,
       useMemo,
       useReducer,
@@ -301,12 +391,13 @@ class ReactShallowRenderer {
         action,
         next: null,
       };
-      if (this._renderPhaseUpdates === null) {
-        this._renderPhaseUpdates = new Map();
+      let renderPhaseUpdates = this._renderPhaseUpdates;
+      if (renderPhaseUpdates === null) {
+        this._renderPhaseUpdates = renderPhaseUpdates = new Map();
       }
-      const firstRenderPhaseUpdate = this._renderPhaseUpdates.get(queue);
+      const firstRenderPhaseUpdate = renderPhaseUpdates.get(queue);
       if (firstRenderPhaseUpdate === undefined) {
-        this._renderPhaseUpdates.set(queue, update);
+        renderPhaseUpdates.set(queue, update);
       } else {
         // Append the update to the end of the list.
         let lastRenderPhaseUpdate = firstRenderPhaseUpdate;
@@ -337,7 +428,8 @@ class ReactShallowRenderer {
       if (this._workInProgressHook.next === null) {
         this._isReRender = false;
         // Append to the end of the list
-        this._workInProgressHook = this._workInProgressHook.next = createHook();
+        this._workInProgressHook = (this
+          ._workInProgressHook: any).next = createHook();
       } else {
         // There's already a work-in-progress. Reuse it.
         this._isReRender = true;
@@ -347,12 +439,18 @@ class ReactShallowRenderer {
     return this._workInProgressHook;
   }
 
-  _prepareToUseHooks(componentIdentity) {
+  _prepareToUseHooks(componentIdentity: Object): void {
+    if (
+      this._previousComponentIdentity !== null &&
+      this._previousComponentIdentity !== componentIdentity
+    ) {
+      this._firstWorkInProgressHook = null;
+    }
     this._currentlyRenderingComponent = componentIdentity;
+    this._previousComponentIdentity = componentIdentity;
   }
 
-  _finishHooks(element, context) {
-    this._hooksEnabled = false;
+  _finishHooks(element: ReactElement, context: null | Object) {
     if (this._didScheduleRenderPhaseUpdate) {
       // Updates were scheduled during the render phase. They are stored in
       // the `renderPhaseUpdates` map. Call the component again, reusing the
@@ -381,7 +479,7 @@ class ReactShallowRenderer {
     return this._rendered;
   }
 
-  render(element, context = emptyObject) {
+  render(element: ReactElement | null, context: null | Object = emptyObject) {
     invariant(
       React.isValidElement(element),
       'ReactShallowRenderer render(): Invalid component element.%s',
@@ -390,6 +488,7 @@ class ReactShallowRenderer {
           'it by passing it to React.createElement.'
         : '',
     );
+    element = ((element: any): ReactElement);
     // Show a special message for host elements since it's a common case.
     invariant(
       typeof element.type !== 'string',
@@ -448,15 +547,25 @@ class ReactShallowRenderer {
         this._mountClassComponent(element, this._context);
       } else {
         if (enableHooks) {
+          const prevDispatcher = ReactCurrentDispatcher.current;
+          ReactCurrentDispatcher.current = this._dispatcher;
           this._prepareToUseHooks(element.type);
-        }
-        this._rendered = element.type.call(
-          undefined,
-          element.props,
-          this._context,
-        );
-        if (enableHooks) {
+          try {
+            this._rendered = element.type.call(
+              undefined,
+              element.props,
+              this._context,
+            );
+          } finally {
+            ReactCurrentDispatcher.current = prevDispatcher;
+          }
           this._finishHooks(element, context);
+        } else {
+          this._rendered = element.type.call(
+            undefined,
+            element.props,
+            this._context,
+          );
         }
       }
     }
@@ -475,6 +584,7 @@ class ReactShallowRenderer {
     }
 
     this._firstWorkInProgressHook = null;
+    this._previousComponentIdentity = null;
     this._context = null;
     this._element = null;
     this._newState = null;
@@ -482,7 +592,7 @@ class ReactShallowRenderer {
     this._instance = null;
   }
 
-  _mountClassComponent(element, context) {
+  _mountClassComponent(element: ReactElement, context: null | Object) {
     this._instance.context = context;
     this._instance.props = element.props;
     this._instance.state = this._instance.state || null;
@@ -519,7 +629,7 @@ class ReactShallowRenderer {
     // because DOM refs are not available.
   }
 
-  _updateClassComponent(element, context) {
+  _updateClassComponent(element: ReactElement, context: null | Object) {
     const {props, type} = element;
 
     const oldState = this._instance.state || emptyObject;
@@ -589,7 +699,10 @@ class ReactShallowRenderer {
     // because DOM refs are not available.
   }
 
-  _updateStateFromStaticLifecycle(props) {
+  _updateStateFromStaticLifecycle(props: Object) {
+    if (this._element === null) {
+      return;
+    }
     const {type} = this._element;
 
     if (typeof type.getDerivedStateFromProps === 'function') {
@@ -652,7 +765,7 @@ function shouldConstruct(Component) {
 }
 
 function getMaskedContext(contextTypes, unmaskedContext) {
-  if (!contextTypes) {
+  if (!contextTypes || !unmaskedContext) {
     return emptyObject;
   }
   const context = {};
diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.internal.js b/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.internal.js
index 7f0dcc687af4e..2b1f3f56b8238 100644
--- a/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.internal.js
+++ b/packages/react-test-renderer/src/__tests__/ReactShallowRendererHooks-test.internal.js
@@ -270,4 +270,51 @@ describe('ReactShallowRenderer with hooks', () => {
       </div>,
     );
   });
+
+  it('should not leak state when component type changes', () => {
+    function SomeComponent({defaultName}) {
+      const [name] = React.useState(defaultName);
+
+      return (
+        <div>
+          <p>
+            Your name is: <span>{name}</span>
+          </p>
+        </div>
+      );
+    }
+
+    function SomeOtherComponent({defaultName}) {
+      const [name] = React.useState(defaultName);
+
+      return (
+        <div>
+          <p>
+            Your name is: <span>{name}</span>
+          </p>
+        </div>
+      );
+    }
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(
+      <SomeComponent defaultName={'Dominic'} />,
+    );
+    expect(result).toEqual(
+      <div>
+        <p>
+          Your name is: <span>Dominic</span>
+        </p>
+      </div>,
+    );
+
+    result = shallowRenderer.render(<SomeOtherComponent defaultName={'Dan'} />);
+    expect(result).toEqual(
+      <div>
+        <p>
+          Your name is: <span>Dan</span>
+        </p>
+      </div>,
+    );
+  });
 });

From c8134905093350001da36b143f68db23971f1e4f Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <dg@domgan.com>
Date: Thu, 17 Jan 2019 14:24:17 +0000
Subject: [PATCH 4/6] More follow ups

---
 .../react-test-renderer/src/ReactShallowRenderer.js   | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index fad1a87162693..8120db15f547c 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -298,27 +298,26 @@ class ReactShallowRenderer {
 
     const useMemo = <T>(
       nextCreate: () => T,
-      inputs: Array<mixed> | void | null,
+      deps: Array<mixed> | void | null,
     ): T => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
 
-      const nextInputs =
-        inputs !== undefined && inputs !== null ? inputs : [nextCreate];
+      const nextDeps = deps === undefined ? null : deps;
 
       if (
         this._workInProgressHook !== null &&
         this._workInProgressHook.memoizedState !== null
       ) {
         const prevState = this._workInProgressHook.memoizedState;
-        const prevInputs = prevState[1];
-        if (areHookInputsEqual(nextInputs, prevInputs)) {
+        const prevDeps = prevState[1];
+        if (areHookInputsEqual(nextDeps, prevDeps)) {
           return prevState[0];
         }
       }
 
       const nextValue = nextCreate();
-      (this._workInProgressHook: any).memoizedState = [nextValue, nextInputs];
+      (this._workInProgressHook: any).memoizedState = [nextValue, nextDeps];
       return nextValue;
     };
 

From 08bb591ec5ecd760dbbb2512cbdea8e4de20cca0 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <dg@domgan.com>
Date: Thu, 17 Jan 2019 14:59:57 +0000
Subject: [PATCH 5/6] Revert last change

---
 packages/react-test-renderer/src/ReactShallowRenderer.js | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index 8120db15f547c..2c864d532c6e5 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -298,12 +298,13 @@ class ReactShallowRenderer {
 
     const useMemo = <T>(
       nextCreate: () => T,
-      deps: Array<mixed> | void | null,
+      inputs: Array<mixed> | void | null,
     ): T => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
 
-      const nextDeps = deps === undefined ? null : deps;
+      const nextDeps =
+        inputs !== undefined && inputs !== null ? inputs : [nextCreate];
 
       if (
         this._workInProgressHook !== null &&

From 4bc5250f56e1737b787df7735c3f6caf480f0e81 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <dg@domgan.com>
Date: Thu, 17 Jan 2019 17:22:22 +0000
Subject: [PATCH 6/6] Fix issue based on feedback

---
 .../react-test-renderer/src/ReactShallowRenderer.js   | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index 2c864d532c6e5..890596721b8d2 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -298,13 +298,12 @@ class ReactShallowRenderer {
 
     const useMemo = <T>(
       nextCreate: () => T,
-      inputs: Array<mixed> | void | null,
+      deps: Array<mixed> | void | null,
     ): T => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
 
-      const nextDeps =
-        inputs !== undefined && inputs !== null ? inputs : [nextCreate];
+      const nextDeps = deps !== undefined && deps !== null ? deps : null;
 
       if (
         this._workInProgressHook !== null &&
@@ -312,8 +311,10 @@ class ReactShallowRenderer {
       ) {
         const prevState = this._workInProgressHook.memoizedState;
         const prevDeps = prevState[1];
-        if (areHookInputsEqual(nextDeps, prevDeps)) {
-          return prevState[0];
+        if (nextDeps !== null) {
+          if (areHookInputsEqual(nextDeps, prevDeps)) {
+            return prevState[0];
+          }
         }
       }