Update a nested structure in an immutable fashion (like a lens)

updater.ts

import { isEqual } from 'lodash';

/*
 * Usage: imagine you have part of a reducer like this:
 *
 * switch(action.type) {
 *   case 'CAMPAIGN_INITIALISED':
 *     campaign = state[action.campaignId];
 *     return {
 *         ...state,
 *       [action.campaignId]: { ...campaign, initialised: true },
 *     };
 *
 * This gets unwieldy as you have more nested structures
 * now you can write it instead as:
 *
 * case 'CAMPAIGN_INITIALISED':
 *   return update(state)
 *     .get(action.campaignId)
 *     .get('initialised')
 *     .set((init) => true);
 *
 * you can probably make this into a more powerful abstraction if you dont evaluate them immediately
 * and make them composable. The implementation would also probably be easier to understand as a class
 * but hey, I was in functional mode
 */

interface Updater<S, C> {
  get: Getter<S, C>;
  set: Setter<S, C>;
}

type Mut<C> = (cur: C) => C;
type Getter<S, C> = <K extends keyof C>(key: K) => Updater<S, C[K]>
type Setter<S, C> = (mut: Mut<C>) => S;
type AwareSetter<S, C> = (mut: Mut<C>, top: boolean) => S;

export function update<S extends Object>(state: S): Updater<S, S> {

  function getter<C>(current: C, setter: AwareSetter<S, C>) {
    function get<K extends keyof C>(key: K): Updater<S, C[K]> {
      let newCurrent: C[K] = current[key];

      function newSetter(mut: Mut<C[K]>, top = true): S {
        return setter((cur: C) => {
          let newVal = mut(newCurrent);
          if (newVal === newCurrent) {
            return cur;
          }
          // If we're the first setter to be called do a more expensive
          // equality check to see if we can avoid updating the state
          if (top && isEqual(newVal, newCurrent)) {
            return cur;
          }
          return {
            ...(cur as any),
            [key as any]: newVal,
          };
        }, false);// pass false to nested setter
      }
      return {
        get: getter(newCurrent, newSetter),
        set: newSetter,
      };
    }
    return get;
  }
  let initialSetter = (mutator: Mut<S>, top = true) => mutator(state);

  return {
    get: getter(state, initialSetter),
    set: initialSetter,
  };
}