React.memo but it ignores functions changes

memoIgnoringFunctions.js

import React, { memo, useMemo, useLayoutEffect, useRef } from 'react';
import objectHash from 'object-hash';
import _partition from 'lodash/partition';
import _fromPairs from 'lodash/fromPairs';

/**
 * an HOC that uses `React.memo` to memoize component expect it ignores changes
 * to functions while also keeping those functions up-to-date somehow
 */
function memoIgnoringFunctions(Component, propsAreEqual) {
  // call React.memo first
  const MemoedComponent = memo(Component, propsAreEqual);

  // the resulting component this HOC returns
  function MemoIgnoringFunctions(props) {
    const propEntries = Object.entries(props);

    const [functionPropEntries, nonFunctionPropEntries] = _partition(
      propEntries,
      entry => typeof entry[1] === 'function',
    );

    const nonFunctionProps = _fromPairs(nonFunctionPropEntries);
    const functionPropKeys = functionPropEntries.map(([key]) => key);

    // this ref is used to hold the most current values of the functions.
    // the layout effect runs before any other effects and is ideal for updating
    // `currentFunctions` ref with new function references
    const currentFunctionsRef = useRef({});
    useLayoutEffect(() => {
      const currentFunctions = currentFunctionsRef.current;

      for (const key of functionPropKeys) {
        currentFunctions[key] = props[key];
      }
    });

    // the `preservedFunctions` is a memoed value calculated from a hash of the
    // `functionPropKeys`. this `useMemo` will return a new set of
    // `preservedFunctions` when those function key changes
    // eslint-disable-next-line react-hooks/exhaustive-deps
    const preservedFunctions = useMemo(() => {
      return functionPropKeys.reduce((functionCache, functionKey) => {
        functionCache[functionKey] = (...args) => {
          // this is the real trick:
          //
          // we can get away with wrapping the functions because the functions
          // can simply lazily pull the latest function value from our ref.
          // the layout effect keeps the values in sync
          const currentFunction = currentFunctionsRef.current[functionKey];
          return currentFunction(...args);
        };

        return functionCache;
      }, {});
    }, [objectHash(functionPropKeys)]);

    return <MemoedComponent {...nonFunctionProps} {...preservedFunctions} />;
  }

  return MemoIgnoringFunctions;
}

export default memoIgnoringFunctions;

memoIgnoringFunctions.test.js

import React, { useState, useEffect } from 'react';
import { act, create } from 'react-test-renderer';
import memoIgnoringFunctions from './memoIgnoringFunctions';

it("memoizes components but doesn't consider different function references", async () => {
  const childEffectHandler = jest.fn();
  const parentEffectHandler = jest.fn();

  const done = new DeferredPromise();

  function Child({ onStuff, foo }) {
    useEffect(() => {
      childEffectHandler({ onStuff, foo });
    });

    return null;
  }

  const Memoed = memoIgnoringFunctions(Child);

  function Parent() {
    const [reRender, setReRender] = useState(false);
    const [foo, setFoo] = useState('foo');

    useEffect(() => {
      parentEffectHandler();
    });

    useEffect(() => {
      setReRender(true);
    }, []);

    useEffect(() => {
      if (reRender) {
        setFoo('bar');
      }
    }, [reRender]);

    useEffect(() => {
      if (foo === 'bar') {
        done.resolve();
      }
    }, [foo]);

    const handleStuff = () => {
      return foo;
    };

    return <Memoed onStuff={handleStuff} foo={foo} />;
  }

  await act(async () => {
    create(<Parent />);
    await done;
  });

  expect(parentEffectHandler).toHaveBeenCalledTimes(3);
  expect(childEffectHandler).toHaveBeenCalledTimes(2);
  expect(childEffectHandler.mock.calls.map(args => args[0])).toMatchInlineSnapshot(`
    Array [
      Object {
        "foo": "foo",
        "onStuff": [Function],
      },
      Object {
        "foo": "bar",
        "onStuff": [Function],
      },
    ]
  `);

  const first = childEffectHandler.mock.calls[0][0].onStuff;
  const second = childEffectHandler.mock.calls[1][0].onStuff;

  // these will both be the _latest_ value due to the layout effect updating them
  expect(first()).toBe('bar');
  expect(second()).toBe('bar');

  expect(first).toBe(second);
});

// this a promise that you can `.resolve` somewhere else
class DeferredPromise {
  constructor() {
    this.state = 'pending';
    this._promise = new Promise((resolve, reject) => {
      this.resolve = value => {
        this.state = 'fulfilled';
        resolve(value);
      };
      this.reject = reason => {
        this.state = 'rejected';
        reject(reason);
      };
    });

    this.then = this._promise.then.bind(this._promise);
    this.catch = this._promise.catch.bind(this._promise);
    this.finally = this._promise.finally.bind(this._promise);
  }

  [Symbol.toStringTag] = 'Promise';
}
re

The real trick here is the layout effect + ref + some laziness:

by nature of functions, we can defer getting the latest function value by pulling it at the time of invocation in the function wrappers and then use a layout effect with a ref to keep those values updated with every render.