useReducerWithEmitEffect.js

const {useCallback, useEffect, useReducer, useRef} = require('react');

let effectCapture = null;

exports.useReducerWithEmitEffect = function(reducer, initialArg, init) {
  let updateCounter = useRef(0);
  let wrappedReducer = useCallback(function(oldWrappedState, action) {
    effectCapture = [];
    try {
      let newState = reducer(oldWrappedState.state, action.action);
      let lastAppliedContiguousUpdate = oldWrappedState.lastAppliedContiguousUpdate;
      let effects = oldWrappedState.effects || [];
      if (lastAppliedContiguousUpdate + 1 === action.updateCount) {
        lastAppliedContiguousUpdate++;
        effects.push(...effectCapture);
      }
      return {
        state: newState,
        lastAppliedContiguousUpdate,
        effects,
      };
    } finally {
      effectCapture = null;
    }
  }, [reducer]);
  let [wrappedState, rawDispatch] = useReducer(wrappedReducer, undefined, function() {
    let initialState;
    if (init !== undefined) {
      initialState = init(initialArg);
    } else {
      initialState = initialArg;
    }
    return {
      state: initialState,
      lastAppliedContiguousUpdate: 0,
      effects: null,
    };
  });
  let dispatch = useCallback(function(action) {
    updateCounter.current++;
    rawDispatch({updateCount: updateCounter.current, action});
  }, []);
  useEffect(function() {
    if (wrappedState.effects) {
      wrappedState.effects.forEach(function(eff) {
        eff();
      });
    }
    wrappedState.effects = null;
  });
  return [wrappedState.state, dispatch];
};

exports.emitEffect = function(fn) {
  if (!effectCapture) {
    throw new Error('emitEffect can only be called from a useReducerWithEmitEffect reducer');
  }
  effectCapture.push(fn);
};

useReducerWithEmitEffect.test.js

// For whatever reason, we need to mock this *and* use RTR._Scheduler below. Why? Who knows.
jest.mock('scheduler', () => require.requireActual('scheduler/unstable_mock'));

const React = require('react');
const ReactTestRenderer = require('react-test-renderer');

const {useReducerWithEmitEffect, emitEffect} = require('./useReducerWithEmitEffect.js');

let _state;
let _dispatch;
let _log;

beforeEach(() => {
  _state = _dispatch = undefined;
  _log = [];
});

function Foo() {
  React.useLayoutEffect(() => {
    _log.push('commit');
  });
  let [state, dispatch] = useReducerWithEmitEffect(function(state, action) {
    let calculation = `${state} + ${action} = ${state + action}`;
    _log.push(`reduce: ${calculation}`);
    emitEffect(() => {
      _log.push(`effect: ${calculation}`);
    });
    return state + action;
  }, 0);
  _state = state;
  _dispatch = dispatch;
  return state;
}

it('initializes', () => {
  const root = ReactTestRenderer.create(
    <Foo />,
    {unstable_isConcurrent: true},
  );
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(0);
  expect(_log).toEqual(['commit']);
  _log.length = 0;
});

it('dispatches', () => {
  const root = ReactTestRenderer.create(
    <Foo />,
    {unstable_isConcurrent: true},
  );
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(0);
  expect(_log).toEqual(['commit']);
  _log.length = 0;

  ReactTestRenderer.act(() => {
    _dispatch(1);
    // Initial effect run eagerly
    expect(_log).toEqual([
      'reduce: 0 + 1 = 1',
    ]);
    _log.length = 0;
  });
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(1);
  expect(_log).toEqual([
    'reduce: 0 + 1 = 1',
    'commit',
    'effect: 0 + 1 = 1',
  ]);
});

it('does two in series', () => {
  const root = ReactTestRenderer.create(
    <Foo />,
    {unstable_isConcurrent: true},
  );
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(0);
  expect(_log).toEqual(['commit']);
  _log.length = 0;

  ReactTestRenderer.act(() => {
    _dispatch(1);
    // Initial effect run eagerly
    expect(_log).toEqual([
      'reduce: 0 + 1 = 1',
    ]);
    _log.length = 0;
  });
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(1);
  expect(_log).toEqual([
    'reduce: 0 + 1 = 1',
    'commit',
    'effect: 0 + 1 = 1',
  ]);
  _log.length = 0;

  ReactTestRenderer.act(() => {
    _dispatch(2);
    // Why doesn't this one also run eagerly? I might've screwed up the
    // scheduler mock somehow.
    expect(_log).toEqual([
      // 'reduce: 1 + 2 = 3',
    ]);
    _log.length = 0;
  });
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(3);
  expect(_log).toEqual([
    'reduce: 1 + 2 = 3',
    'commit',
    'effect: 1 + 2 = 3',
  ]);
});

it('does two at once', () => {
  const root = ReactTestRenderer.create(
    <Foo />,
    {unstable_isConcurrent: true},
  );
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(0);
  expect(_log).toEqual(['commit']);
  _log.length = 0;

  ReactTestRenderer.act(() => {
    _dispatch(1);
    _dispatch(2);
    // Initial effect run eagerly
    expect(_log).toEqual([
      'reduce: 0 + 1 = 1',
    ]);
    _log.length = 0;
  });
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(3);
  expect(_log).toEqual([
    'reduce: 0 + 1 = 1',
    'reduce: 1 + 2 = 3',
    'commit',
    'effect: 0 + 1 = 1',
    'effect: 1 + 2 = 3',
  ]);
});

it('does low and hi pri', () => {
  const root = ReactTestRenderer.create(
    <Foo />,
    {unstable_isConcurrent: true},
  );
  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(0);
  expect(_log).toEqual(['commit']);
  _log.length = 0;

  ReactTestRenderer.act(() => {
    _dispatch(1);
    // Initial effect run eagerly
    expect(_log).toEqual([
      'reduce: 0 + 1 = 1',
    ]);
    _log.length = 0;
  });

  root.unstable_flushSync(() => {
    _dispatch(2);
  });
  // Only the hi-pri update runs, and no effects happen
  expect(_log).toEqual([
    'reduce: 0 + 2 = 2',
    'commit',
  ]);
  _log.length = 0;

  ReactTestRenderer._Scheduler.unstable_flushWithoutYielding();
  expect(_state).toBe(3);
  expect(_log).toEqual([
    'reduce: 0 + 1 = 1',
    'reduce: 1 + 2 = 3',
    'commit',
    'effect: 0 + 1 = 1',
    'effect: 1 + 2 = 3',
  ]);
});

@faceyspacey

What is the purpose of lastAppliedContiguousUpdate?

This allows the reducer to tell if the action being processed is really the 'next' action.

In the last test the high priority _dispatch(2) is going to be rendered without the _dispatch(1) being included but it will have an update count of 2. since lastAppliedContiguousUpdate will be 0 for this render (the first dispatch is no included during this render) the reducer knows not to emit effects because it is expecting a re-render to happen later that includes all the dispatches in sequence.

When the lower priority render happens (in the test it is just flushing everything synchronously so this is like rendering at the lowest priority I think) the reducer sees _dispatch(1) as having an updateCount of 1 which is "next" update. after that the reducer processes dispatch(2) which has an updateCount of 2 which again is "next". the effect list is built up and if this render is not interrupted again it will commit with the effects in place

It's a pretty ingenious system