optimisticReducer.tsx

import { useCallback, Reducer, Dispatch, useState, useMemo, useEffect, useRef } from 'react';

type Batch<S, A extends object> = {
  state: S,
  actions: A[],
  next: Batch<S, A> | null
}

class StreamReducer<S, A extends object> {
  /**
   * The current reduction of all actions in history.
   */
  state: S;

  /**
   * Batches form a singly-linked list, but only retain a reference to
   * the latest batch. All others are retained via `WeakMap`
   */
  batch: Batch<S, A>;

  /**
   * Maps actions to the batches they belong to, this is used to 
   * recall a previous batch that we need to re-reduce.
   */
  uncommitted = new WeakMap<A, Batch<S, A>>();

  reducer;

  constructor(reducer: Reducer<S, A>, initialState: S) {
    this.reducer = reducer;
    this.state = initialState;
    this.batch = {
      state: initialState,
      actions: [],
      next: null
    }
  }

  do(action: A): S {
    const state = this.state = (0, this.reducer)(this.state, action);
    this.uncommitted.set(action, this.batch);

    // Push the action, and if the batch is full, start a new batch.
    if(this.batch.actions.push(action) >= 5) {
      const nextBatch = {
        state,
        actions: [],
        next: null,
      };

      this.batch.next = nextBatch;
      this.batch = nextBatch;
    }
    return this.state;
  }

  undo(action: A): S {
    const batch = this.uncommitted.get(action);
    if(!batch) throw new Error("unknown action");

    let cursor = { ...batch };

    // Reset history 
    // We _have_ to reuse the same batch object here, we can't create a new one
    // there is a potential previous batch which has a next referencing this batch
    // and we can't fix it up.
    this.batch = batch;
    this.state = batch.state;
    batch.actions = [];

    // Replay history creating new batches
    while(cursor !== null) {
      for(const redo of cursor.actions) {
        if(redo === action) continue;
        this.do(redo);
      }

      if(!cursor.next) break;
      cursor = cursor.next;
    }

    return this.state;
  }
}

function useOptimisticReducer<S, A extends object, I>(
  reducer: Reducer<S, A>,
  initializerArg: I,
  initializer: (arg: I) => S
): [S, Dispatch<A>, Dispatch<A>] {
    const [state, setState] = useState<S>(() => initializer(initializerArg));

    const streamRef = useRef<StreamReducer<S, A>>(null!);
    streamRef.current ??= new StreamReducer(reducer, state);
    const { current: stream } = streamRef;

    return [state,
      useCallback((action: A) => setState(stream.do(action)), [stream]),
      useCallback((action: A) => setState(stream.undo(action)), [stream]),
    ];
}

// Usage Example:
type State = { amount: number };
type CountAction = 
  | { type: "increment", amount: number }
  | { type: "decrement", amount: number }
  | { type: "set", amount: number }


const runningSumReducer = (state: State, action: CountAction): State => {
  switch (action.type) {
    case 'increment':
      return { ...state, amount: state.amount + action.amount };
    case 'decrement':
      return { ...state, amount: state.amount - action.amount };
    case 'set':
      return { amount: action.amount }
    default:
      return state;
  }
};

// Example Component
const CounterComponent = () => {
  const [state, dispatch, recall] = useOptimisticReducer(runningSumReducer, undefined, () => ({ amount: 0 }));

  /*
  Enable this to test memory leaks, observe that it does not leak
  useEffect(() => {
    const io = setInterval(() => dispatch({ type: 'increment', amount: 1 }), 0);
    return () => clearInterval(io);
  }, [dispatch])
  */

  return (
    <div>
      <p>Count: {state.amount}</p>
      <button onClick={() => {
        const action = { type: 'increment', amount: 1 } satisfies CountAction;
        dispatch(action);
        setTimeout(() => recall(action), 1000);
      }}>Increment</button>
      <button onClick={() => {
        const action = { type: 'decrement', amount: 1 } satisfies CountAction;
        dispatch(action);
        setTimeout(() => recall(action), 1000);
      }}>Decrement</button>
      <button onClick={() => dispatch({ type: 'set', amount: 0 })}>Reset</button>
    </div>
  );
};

export default CounterComponent;