lenses.ts

export type GenericLens<T, R> = {
  get(input: T): R;
  set(input: T, value: R): T;
  over(input: T, transform: (input: R) => R): T;
};

export type LensInput = string | number | GenericLens<any, any>

export type LensType<
  T,
  S extends LensInput[]
> = S extends []
  ? T
  : S extends [infer Head, ...infer Tail]
  ? Head extends keyof T
    ? Tail extends Array<string | number>
      ? LensType<T[Head], Tail>
      : never
    : Head extends { get: (whole: T) => infer R }
    ? R
    : never
  : never;

function merge<T extends object>(whole: T, part: Partial<T>): T {
  return Object.assign(
    Array.isArray(whole) ? [] : Object.create(null),
    whole,
    part
  );
}

const idlens = {
  get(whole: any) {
    return whole;
  },
  set(_whole: any, part: any) {
    return part;
  },
  over(whole: any, transform: (input: any) => any) {
    return transform(whole);
  },
};

export interface Lens<
  S extends LensInput[]
> {
  get<T>(whole: T): LensType<T, S>;
  set<T>(whole: T, part: LensType<T, S>): T;
  over<T>(
    whole: T,
    transform: (input: LensType<T, S>) => LensType<T, S>
  ): T;
}

export function lens<S extends LensInput[]>(
  ...selectors: S
): Lens<S> {
  return selectors.reduce(
    (prev: any, attr) =>
      typeof attr === "object"
        ? {
            get(whole: any) {
              return attr.get(prev.get(whole));
            },
            set(whole: any, part: any) {
              return prev.over(whole, (input: any) => attr.set(input, part));
            },
            over(whole: any, transform: (input: any) => any) {
              return prev.over(whole, (input: any) =>
                attr.over(input, transform)
              );
            },
          }
        : {
            get(whole: any) {
              return prev.get(whole)[attr];
            },
            set(whole: any, part: any) {
              return prev.over(whole, (input: any) =>
                merge(input, { [attr]: part })
              );
            },
            over(whole: any, transform: (input: any) => any) {
              return prev.over(whole, (input: any) =>
                merge(input, { [attr]: transform(input[attr]) })
              );
            },
          },
    idlens
  ) as any;
}

react-lenses.tsx

import {
  Dispatch,
  SetStateAction,
  useCallback,
  useEffect,
  useMemo,
  useState,
} from "react";
import { Lens, lens, LensInput, LensType } from "./lenses";

const SUBSCRIBER = Symbol("subscriber");
const VALUE = Symbol("value");
const PARENT = Symbol("parent");
const PATH = Symbol("path");

export interface LensContext<T> {
  [VALUE]: T;
  [SUBSCRIBER]: Set<() => void>;
}

class PrimaryLensContext<T extends object> implements LensContext<T> {
  [VALUE]: T;
  readonly [SUBSCRIBER] = new Set<() => void>();
  constructor(value: T) {
    this[VALUE] = value;
  }
}

class DeriveLensContext<T extends object, S extends LensInput[]>
  implements LensContext<LensType<T, S>>
{
  readonly [PARENT]: LensContext<T>;
  readonly [PATH]: Lens<S>;
  readonly [SUBSCRIBER]: Set<() => void>;
  constructor(parent: LensContext<T>, ...path: S) {
    this[PARENT] = parent;
    this[PATH] = lens(...path);
    this[SUBSCRIBER] = parent[SUBSCRIBER];
  }

  get [VALUE]() {
    return this[PATH].get(this[PARENT][VALUE]);
  }

  set [VALUE](value: LensType<T, S>) {
    this[PARENT][VALUE] = this[PATH].set(this[PARENT][VALUE], value);
  }
}

export function createLens<T extends object>(defValue: T): LensContext<T> {
  return new PrimaryLensContext(defValue);
}

export function useDeriveLens<T extends object, S extends LensInput[]>(
  ctx: LensContext<T>,
  ...selectors: S
): LensContext<LensType<T, S>> {
  return useMemo(
    () => new DeriveLensContext(ctx, ...selectors),
    [ctx, selectors]
  );
}

export function useLens<T, S extends LensInput[]>(
  ctx: LensContext<T>,
  ...selectors: S
) {
  const lenses = useMemo(() => lens(...selectors), [selectors]);
  const [value, setValue] = useState(() => lenses.get(ctx[VALUE]));
  const update = useCallback(
    () => setValue(() => lenses.get(ctx[VALUE])),
    [ctx, lenses]
  );
  useEffect(() => update(), [ctx, lenses]);
  useEffect(() => {
    const subs = ctx[SUBSCRIBER];
    subs.add(update);
    return () => void subs.delete(update);
  }, [ctx]);
  const set = useCallback(
    (value: SetStateAction<LensType<T, S>>) => {
      ctx[VALUE] =
        typeof value === "function"
          ? lenses.over(ctx[VALUE], value as any)
          : lenses.set(ctx[VALUE], value);
      ctx[SUBSCRIBER].forEach((f) => f());
    },
    [ctx, lenses]
  );
  return [value, set] as [
    LensType<T, S>,
    Dispatch<SetStateAction<LensType<T, S>>>
  ];
}

export function useLensTransformer<T, S extends LensInput[]>(
  ctx: LensContext<T>,
  ...selectors: S
) {
  const lenses = useMemo(() => lens(...selectors), [selectors]);
  return useCallback(
    (value: SetStateAction<LensType<T, S>>) => {
      ctx[VALUE] =
        typeof value === "function"
          ? lenses.over(ctx[VALUE], value as any)
          : lenses.set(ctx[VALUE], value);
      ctx[SUBSCRIBER].forEach((f) => f());
    },
    [ctx, lenses]
  );
}