Ensure that state changes are applied immediately on mount

.eslintrc.json

@@ -25,7 +25,12 @@
     "@typescript-eslint/no-explicit-any": "off",
     "@typescript-eslint/array-type": "off",
     "react-hooks/rules-of-hooks": "error",
-    "react-hooks/exhaustive-deps": "warn",
+    "react-hooks/exhaustive-deps": [
+      "warn",
+      {
+        "additionalHooks": "useImmediateEffect"
+      }
+    ],
     "react/prop-types": "off",
     "react/no-children-prop": "off",
     "sort-keys": "off",

src/hooks/useImmediateEffect.test.ts

@@ -1,5 +1,5 @@
 import React from 'react';
-import { renderHook } from 'react-hooks-testing-library';
+import { renderHook, act } from 'react-hooks-testing-library';
 import { useImmediateEffect } from './useImmediateEffect';
 
 it('calls effects immediately on mount', () => {
@@ -8,14 +8,37 @@ it('calls effects immediately on mount', () => {
   const effect = jest.fn();
 
   spy.mockImplementation(useEffect);
-  renderHook(() => useImmediateEffect(effect, [effect]));
 
-  expect(effect).toHaveBeenCalledTimes(1);
-  expect(effect).toHaveBeenCalledTimes(1);
+  renderHook(() => {
+    useImmediateEffect(effect, [effect]);
+    expect(effect).toHaveBeenCalledTimes(1);
+  });
 
+  expect(effect).toHaveBeenCalledTimes(1);
   expect(useEffect).toHaveBeenCalledWith(expect.any(Function), [
     expect.any(Function),
   ]);
 
-  useEffect.mockRestore();
+  spy.mockRestore();
+});
+
+it('behaves like useEffect otherwise', () => {
+  const spy = jest.spyOn(React, 'useEffect');
+  const effect = jest.fn();
+
+  const { result } = renderHook(() => {
+    const [ref, setState] = React.useState({});
+    useImmediateEffect(effect, [ref]);
+
+    expect(effect).toHaveBeenCalledTimes(1);
+    expect(spy).toHaveBeenCalled();
+
+    const forceUpdate = () => setState({});
+    return forceUpdate;
+  });
+
+  act(() => result.current()); // forceUpdate
+  expect(effect).toHaveBeenCalledTimes(2);
+
+  spy.mockRestore();
 });

src/hooks/useImmediateEffect.ts

@@ -1,4 +1,6 @@
-import { useRef, useEffect, useCallback } from 'react';
+/* eslint-disable react-hooks/exhaustive-deps */
+
+import { useRef, useEffect } from 'react';
 import { noop } from '../utils';
 
 enum LifecycleState {
@@ -9,29 +11,28 @@ enum LifecycleState {
 
 type Effect = () => () => void;
 
-/** This executes an effect immediately on initial render and then treats it as a normal effect */
+/** This is a drop-in replacement for useEffect that will execute the first effect that happens during initial mount synchronously */
 export const useImmediateEffect = (
   effect: Effect,
   changes: ReadonlyArray<any>
 ) => {
   const teardown = useRef(noop);
   const state = useRef(LifecycleState.WillMount);
-  const execute = useCallback(effect, changes);
 
   // On initial render we just execute the effect
   if (state.current === LifecycleState.WillMount) {
     state.current = LifecycleState.DidMount;
-    teardown.current = execute();
+    teardown.current = effect();
   }
 
   useEffect(() => {
     // Initially we skip executing the effect since we've already done so on
     // initial render, then we execute it as usual
     if (state.current === LifecycleState.Update) {
-      return (teardown.current = execute());
+      return (teardown.current = effect());
     } else {
       state.current = LifecycleState.Update;
       return teardown.current;
     }
-  }, [execute]);
+  }, changes);
 };

src/hooks/useImmediateState.test.ts

@@ -0,0 +1,46 @@
+import React from 'react';
+import { renderHook } from 'react-hooks-testing-library';
+import { useImmediateState } from './useImmediateState';
+
+it('updates state immediately during mount', () => {
+  let initialState;
+  let update = 0;
+
+  const setState = jest.fn();
+
+  const spy = jest.spyOn(React, 'useState').mockImplementation(state => {
+    expect(state).toEqual({ x: 'x', test: false });
+    initialState = state;
+    return [state, setState];
+  });
+
+  const { result } = renderHook(() => {
+    const [state, setState] = useImmediateState({ x: 'x', test: false });
+    if (update === 0) setState(s => ({ ...s, test: true }));
+    update++;
+    return state;
+  });
+
+  expect(setState).not.toHaveBeenCalled();
+  expect(result.current).toEqual({ x: 'x', test: true });
+  expect(result.current).toBe(initialState);
+  expect(update).toBe(1);
+
+  spy.mockRestore();
+});
+
+it('behaves like useState otherwise', () => {
+  const setState = jest.fn();
+  const spy = jest
+    .spyOn(React, 'useState')
+    .mockImplementation(state => [state, setState]);
+
+  renderHook(() => {
+    const [state, setState] = useImmediateState({ x: 'x' });
+    React.useEffect(() => setState({ x: 'y' }), [setState]);
+    return state;
+  });
+
+  expect(setState).toHaveBeenCalledTimes(1);
+  spy.mockRestore();
+});

src/hooks/useImmediateState.ts

@@ -0,0 +1,33 @@
+import { useRef, useEffect, useState, useCallback } from 'react';
+
+type SetStateAction<S> = S | ((prevState: S) => S);
+type SetState<S> = (action: SetStateAction<S>) => void;
+
+/** This is a drop-in replacement for useState, limited to object-based state. During initial mount it will mutably update the state, instead of scheduling a React update using setState */
+export const useImmediateState = <S extends {}>(init: S): [S, SetState<S>] => {
+  const isMounted = useRef(false);
+  const initialState = useRef<S>({ ...init });
+  const [state, setState] = useState<S>(initialState.current);
+
+  // This wraps setState and updates the state mutably on initial mount
+  // It also prevents setting the state when the component is unmounted
+  const updateState: SetState<S> = useCallback((action: SetStateAction<S>) => {
+    if (isMounted.current) {
+      setState(action);
+    } else if (typeof action === 'function') {
+      const update = (action as any)(initialState.current);
+      Object.assign(initialState.current, update);
+    } else {
+      Object.assign(initialState.current, action as any);
+    }
+  }, []);
+
+  useEffect(() => {
+    isMounted.current = true;
+    return () => {
+      isMounted.current = false;
+    };
+  }, []);
+
+  return [state, updateState];
+};

src/hooks/useMutation.ts

@@ -1,9 +1,10 @@
 import { DocumentNode } from 'graphql';
-import { useContext, useEffect, useRef, useState } from 'react';
+import { useContext, useCallback } from 'react';
 import { pipe, toPromise } from 'wonka';
 import { Context } from '../context';
 import { OperationResult } from '../types';
 import { CombinedError, createRequest } from '../utils';
+import { useImmediateState } from './useImmediateState';
 
 export interface UseMutationState<T> {
   fetching: boolean;
@@ -19,39 +20,30 @@ export type UseMutationResponse<T, V> = [
 export const useMutation = <T = any, V = object>(
   query: DocumentNode | string
 ): UseMutationResponse<T, V> => {
-  const isMounted = useRef(true);
   const client = useContext(Context);
-  const [state, setState] = useState<UseMutationState<T>>({
+  const [state, setState] = useImmediateState<UseMutationState<T>>({
     fetching: false,
     error: undefined,
     data: undefined,
   });
 
-  useEffect(
-    () => () => {
-      isMounted.current = false;
-    },
-    []
-  );
-
-  const executeMutation = (variables?: V) => {
-    setState({ fetching: true, error: undefined, data: undefined });
-
-    const request = createRequest(query, variables as any);
+  const executeMutation = useCallback(
+    (variables?: V) => {
+      setState({ fetching: true, error: undefined, data: undefined });
 
-    return pipe(
-      client.executeMutation(request),
-      toPromise
-    ).then(result => {
-      const { data, error } = result;
+      const request = createRequest(query, variables as any);
 
-      if (isMounted.current) {
+      return pipe(
+        client.executeMutation(request),
+        toPromise
+      ).then(result => {
+        const { data, error } = result;
         setState({ fetching: false, data, error });
-      }
-
-      return result;
-    });
-  };
+        return result;
+      });
+    },
+    [client, query, setState]
+  );
 
   return [state, executeMutation];
 };

src/hooks/useQuery.ts

@@ -1,11 +1,12 @@
 import { DocumentNode } from 'graphql';
-import { useCallback, useContext, useRef, useState } from 'react';
+import { useCallback, useContext, useRef } from 'react';
 import { pipe, subscribe } from 'wonka';
 import { Context } from '../context';
 import { OperationContext, RequestPolicy } from '../types';
 import { CombinedError, noop } from '../utils';
 import { useRequest } from './useRequest';
 import { useImmediateEffect } from './useImmediateEffect';
+import { useImmediateState } from './useImmediateState';
 
 export interface UseQueryArgs<V> {
   query: string | DocumentNode;
@@ -28,61 +29,54 @@ export type UseQueryResponse<T> = [
 export const useQuery = <T = any, V = object>(
   args: UseQueryArgs<V>
 ): UseQueryResponse<T> => {
-  const isMounted = useRef(true);
   const unsubscribe = useRef(noop);
   const client = useContext(Context);
 
-  // This creates a request which will keep a stable reference
-  // if request.key doesn't change
-  const request = useRequest(args.query, args.variables);
-
-  const [state, setState] = useState<UseQueryState<T>>({
+  // This is like useState but updates the state object
+  // immediately, when we're still before the initial mount
+  const [state, setState] = useImmediateState<UseQueryState<T>>({
     fetching: false,
-    error: undefined,
     data: undefined,
+    error: undefined,
   });
 
+  // This creates a request which will keep a stable reference
+  // if request.key doesn't change
+  const request = useRequest(args.query, args.variables);
+
   const executeQuery = useCallback(
     (opts?: Partial<OperationContext>) => {
       unsubscribe.current();
 
+      // If useQuery is currently paused, set fetching to
+      // false and abort; otherwise start the query
       if (args.pause) {
         setState(s => ({ ...s, fetching: false }));
         unsubscribe.current = noop;
         return;
+      } else {
+        setState(s => ({ ...s, fetching: true }));
       }
 
-      setState(s => ({ ...s, fetching: true }));
-
-      const [teardown] = pipe(
+      [unsubscribe.current] = pipe(
         client.executeQuery(request, {
           requestPolicy: args.requestPolicy,
           ...opts,
         }),
-        subscribe(
-          ({ data, error }) => {
-            if (isMounted.current) {
-              setState({ fetching: false, data, error });
-            }
-          }
-        )
+        subscribe(({ data, error }) => {
+          setState({ fetching: false, data, error });
+        })
       );
-
-      unsubscribe.current = teardown;
     },
-    [request, client, args.pause, args.requestPolicy]
+    [args.pause, args.requestPolicy, client, request, setState]
   );
 
-  // Calls executeQuery on initial render immediately, then
-  // treats it as a normal effect
+  // This calls executeQuery immediately during the initial mount and
+  // otherwise behaves like a normal useEffect; We call executeQuery
+  // everytime it, i.e. its input like request, changes
   useImmediateEffect(() => {
-    isMounted.current = true;
     executeQuery();
-
-    return () => {
-      isMounted.current = false;
-      unsubscribe.current();
-    };
+    return () => unsubscribe.current();
   }, [executeQuery]);
 
   return [state, executeQuery];