Immutably set a nested key on an object.

makeSetValue.ts

type PickByNullableType<D, T> = {
  [K in keyof D]: NonNullable<D[K]> extends T ? K : never;
};

type PickKeysWithNullableValueTypes<D, T> = keyof D &
  PickByNullableType<D, T>[keyof D];

function setLevel1Value<D extends object, K1 extends keyof D>(
  data: D,
  key1: K1,
  value: D[K1],
): D {
  return {
    ...data,
    [key1]: value,
  };
}

function makeSetLevel1Value<D extends object, K1 extends keyof D>(
  setter: React.Dispatch<React.SetStateAction<D>>,
  key1: K1,
) {
  return (value: D[K1]) => {
    setter((data) => setLevel1Value(data, key1, value));
  };
}

function setLevel2Value<
  D extends object,
  K1 extends PickKeysWithNullableValueTypes<D, object>,
  K2 extends keyof NonNullable<D[K1]>,
>(data: D, key1: K1, key2: K2, value: NonNullable<D[K1]>[K2]): D {
  return {
    ...data,
    [key1]: {
      ...data[key1],
      [key2]: value,
    },
  };
}

function makeSetLevel2Value<
  D extends object,
  K1 extends PickKeysWithNullableValueTypes<D, object>,
  K2 extends keyof NonNullable<D[K1]>,
>(setter: React.Dispatch<React.SetStateAction<D>>, key1: K1, key2: K2) {
  return (value: NonNullable<D[K1]>[K2]) => {
    setter((data) => setLevel2Value(data, key1, key2, value));
  };
}

export function makeSetValue<D extends object, K extends keyof D>(
  ...args: Parameters<typeof makeSetLevel1Value<D, K>>
): ReturnType<typeof makeSetLevel1Value<D, K>>;
export function makeSetValue<
  D extends object,
  K1 extends PickKeysWithNullableValueTypes<D, object>,
  K2 extends keyof NonNullable<D[K1]>,
>(
  ...args: Parameters<typeof makeSetLevel2Value<D, K1, K2>>
): ReturnType<typeof makeSetLevel2Value<D, K1, K2>>;
export function makeSetValue<
  D extends object,
  K1 extends PickKeysWithNullableValueTypes<D, object>,
  K2 extends keyof NonNullable<D[K1]>,
  S extends React.Dispatch<React.SetStateAction<D>>,
>(...args: [S, K1] | [S, K1, K2]) {
  return args.length === 2
    ? makeSetLevel1Value<D, K1>(...args)
    : makeSetLevel2Value<D, K1, K2>(...args);
}