Kliesli Arrow Optics

#note.md

This implementation can now be found in production here

arrow-optics.ts

// This is all based on the idea here: https://gist.github.com/serras/5152ec18ec5223b676cc67cac0e99b70
//
import type { $, Kind } from "../kind.ts";
import type { Traversable } from "../traversable.ts";
import type { Refinement } from "../refinement.ts";
import type { Predicate } from "../predicate.ts";
import type { Either } from "../either.ts";
import type { Option } from "../option.ts";
import type { Monad } from "../monad.ts";
import type { Fn } from "../fn.ts";

import * as I from "../identity.ts";
import * as O from "../option.ts";
import * as A from "../array.ts";
import * as R from "../record.ts";
import * as E from "../either.ts";
import { isNotNil } from "../nilable.ts";
import { apply, flow, identity, pipe, todo, unsafeCoerce } from "../fn.ts";

// Build up language of Kliesli Optics

/**
 * Following are the runtime tags associated
 * with various forms of Optics.
 */

export const LensTag = "Lens" as const;
export type LensTag = typeof LensTag;

export const PrismTag = "Prism" as const;
export type PrismTag = typeof PrismTag;

export const TraversalTag = "Traversal" as const;
export type TraversalTag = typeof TraversalTag;

export type OpticTag = LensTag | PrismTag | TraversalTag;

/**
 * Type level mapping from Tag to URI. Since an
 * Optic get function is a Kliesli Arrow a => mb, we
 * associate the Optic Tags as follows:
 *
 * LensTag => Identity
 * PrismTag => Option
 * TraversalTag => Array
 */
export type ToURI<T extends OpticTag> = T extends LensTag ? I.URI
  : T extends PrismTag ? O.URI
  : T extends TraversalTag ? A.URI
  : never;

/**
 * Type level mapping from URI to Tag
 */
export type ToTag<U extends Kind> = U extends A.URI ? TraversalTag
  : U extends O.URI ? PrismTag
  : U extends I.URI ? LensTag
  : never;

/**
 * Join will give us the "loosest" of two tags. This is used to
 * determine the abstraction level that an Optic operatates at. The
 * most contstrained is Identity while the least constrained is Array.
 * The typescript version of the source optics Getters are as follows:
 *
 * ```ts
 * type Lens<S, A>      = { get: (s: S) => Identity<A> };
 * type Prism<S, A>     = { get: (s: S) =>   Option<A> };
 * type Traversal<S, A> = { get: (s: S) =>    Array<A> };
 * ```
 *
 * Here we can see that Lens is constrained to get exactly one A,
 * Prism is constrained to get zero or one A, and Traversal is
 * constrained to get zero, one, or many As. Because of this,
 * Lens can always be lifted to a Prism and Prism can always be
 * lifted to Traversal. All Optics share the same modify function
 * over S and A.
 *
 * Thus Join is like GT where Array > Option > Identity.
 */
export type Join<U extends OpticTag, V extends OpticTag> = U extends
  TraversalTag ? TraversalTag
  : V extends TraversalTag ? TraversalTag
  : U extends PrismTag ? PrismTag
  : V extends PrismTag ? PrismTag
  : LensTag;

/**
 * The runtime level GTE for Join
 */
export function join<A extends OpticTag, B extends OpticTag>(
  a: A,
  b: B,
): Join<A, B> {
  if (a === TraversalTag || b === TraversalTag) {
    return TraversalTag as unknown as Join<A, B>;
  } else if (a === PrismTag || b === PrismTag) {
    return PrismTag as unknown as Join<A, B>;
  } else {
    return LensTag as unknown as Join<A, B>;
  }
}

/**
 * Our new Optic definition. Instead of get and set we use get and modify as
 * set can be derived from modify(() => value). This drastically simplifies
 * implementation.
 */
export type Optic<T extends OpticTag, S, A> = {
  readonly tag: T;
  readonly get: Fn<[S], $<ToURI<T>, [A, never, never], [never], [never]>>;
  readonly modify: Fn<[Fn<[A], A>], Fn<[S], S>>;
};

/**
 * We recover the Lens type from the generic Optic
 */
export type Lens<S, A> = Optic<LensTag, S, A>;

/**
 * We recover the Prism type from the generic Optic
 */
export type Prism<S, A> = Optic<PrismTag, S, A>;

/**
 * We recover the Traversal type from the generic Optic
 */
export type Traversal<S, A> = Optic<TraversalTag, S, A>;

export function optic<U extends OpticTag, S, A>(
  tag: U,
  get: Fn<[S], $<ToURI<U>, [A, never, never], [never], [never]>>,
  modify: Fn<[Fn<[A], A>], Fn<[S], S>>,
): Optic<U, S, A> {
  return { tag, get, modify };
}

/**
 * Construct a Lens from get and modify functions.
 */
export function lens<S, A>(
  get: (s: S) => A,
  modify: Fn<[Fn<[A], A>], Fn<[S], S>>,
): Lens<S, A> {
  return optic(LensTag, get, modify);
}

/**
 * Construct a Prism from get and modify functions.
 */
export function prism<S, A>(
  get: (s: S) => Option<A>,
  modify: Fn<[Fn<[A], A>], Fn<[S], S>>,
): Prism<S, A> {
  return optic(PrismTag, get, modify);
}

/**
 * Construct a Traversal from get and modify functions.
 */
export function traversal<S, A>(
  get: (s: S) => ReadonlyArray<A>,
  modify: Fn<[Fn<[A], A>], Fn<[S], S>>,
): Traversal<S, A> {
  return optic(TraversalTag, get, modify);
}

/**
 * A helper function mapping an Optic tag to it's associate ADT Monad
 * instance.
 */
export function getMonad<U extends OpticTag>(tag: U): Monad<ToURI<U>> {
  return (tag === TraversalTag
    ? A.MonadArray
    : tag === PrismTag
    ? O.MonadOption
    : I.MonadIdentity) as unknown as Monad<ToURI<U>>;
}

/**
 * Cast an Optic<U> as an Optic<V>. This non-exported function only
 * works for the following cases
 *
 * Lens => Lens
 * Lens => Prism
 * Lens => Traversal
 *
 * Prism => Prism
 * Prism => Traversal
 *
 * Traversal => Traversal
 *
 * As long as only Optics over Identity, Option, and Array are in
 * this file. There should be no way to break this casting.
 */
function cast<U extends OpticTag, V extends OpticTag, S, A>(
  G: Optic<U, S, A>,
  tag: V,
): Optic<V, S, A> {
  // Covers Lens => Lens, Prism => Prism, Traversal => Traversal
  if (G.tag === tag as OpticTag) {
    return unsafeCoerce(G);
    // Covers Lens => Traversal, Prism => Traversal
  } else if (tag == TraversalTag) {
    // deno-lint-ignore no-explicit-any
    return unsafeCoerce(traversal((s) => A.of(G.get(s)) as any, G.modify));
    // Covers Lens => Prism
  } else {
    // deno-lint-ignore no-explicit-any
    return unsafeCoerce(prism((s) => O.of(G.get(s)) as any, G.modify));
  }
}

/**
 * Compose two Optics by lifting them to matching ADTs, then chain
 * using the Monad for that ADT. Using a monad here was easier than
 * implementing Arrow all over the fun library
 */
export function compose<V extends OpticTag, A, I>(second: Optic<V, A, I>) {
  return <U extends OpticTag, S>(
    first: Optic<U, S, A>,
  ): Optic<Join<U, V>, S, I> => {
    const tag = join(first.tag, second.tag);
    const _chain = getMonad(tag).chain;
    const _first = cast(first, tag);
    const _second = cast(second, tag);
    return {
      tag,
      get: flow(_first.get, _chain(_second.get)),
      modify: flow(_second.modify, _first.modify),
    };
  };
}

/**
 * The starting place for most Optics. Create an Optic over the
 * identity function.
 */
export function id<A>(): Lens<A, A> {
  return lens(identity, identity);
}

/**
 * Given an Optic over a structure with a property P, construct a
 * new Optic at that property P.
 */
export function prop<A, P extends keyof A>(prop: P) {
  return <U extends OpticTag, S>(
    sa: Optic<U, S, A>,
  ): Optic<Join<U, LensTag>, S, A[P]> =>
    pipe(
      sa,
      compose(
        lens((s) => s[prop], (fii) => (a) => ({ ...a, [prop]: fii(a[prop]) })),
      ),
    );
}

/**
 * Given an Optic over a structure with properties P, construct a new
 * optic that only focuses on those properties
 */
export function props<A, P extends keyof A>(...props: [P, P, ...Array<P>]) {
  return <U extends OpticTag, S>(
    first: Optic<U, S, A>,
  ): Optic<Join<U, LensTag>, S, { [K in P]: A[K] }> =>
    pipe(
      first,
      compose(lens(
        R.pick(props),
        (faa) => (a) => ({ ...a, ...faa(pipe(a, R.pick(props))) }),
      )),
    );
}

/**
 * Given an optic over an array, focus on a value at an index in the
 * array.
 */
export function index(i: number) {
  return <U extends OpticTag, S, A>(
    first: Optic<U, S, ReadonlyArray<A>>,
  ): Optic<Join<U, PrismTag>, S, A> =>
    pipe(first, compose(prism(A.lookup(i), A.modifyAt(i))));
}

/**
 * Given an optic over a record, focus on a value at a key in that
 * record.
 */
export function key(key: string) {
  return <U extends OpticTag, S, A>(
    first: Optic<U, S, Readonly<Record<string, A>>>,
  ): Optic<Join<U, PrismTag>, S, A> =>
    pipe(first, compose(prism(R.lookup(key), R.modifyAt(key))));
}

/**
 * Given an optic over a record, focus on an Option(value) at
 * the given key, allowing one to delete the key by modifying
 * with a None value.
 *
 * TODO: Clean this implementation up.
 */
export function atKey(key: string) {
  return <U extends OpticTag, S, A>(
    first: Optic<U, S, Readonly<Record<string, A>>>,
  ): Optic<Join<U, LensTag>, S, Option<A>> =>
    pipe(
      first,
      compose(lens(R.lookupAt(key), (faa) => (s) =>
        pipe(
          s,
          R.lookup(key),
          faa,
          O.fold(() => R.deleteAt(key), R.insertAt(key)),
          apply(s),
        ))),
    );
}

/**
 * Given an Optic over a structure with a nilable property P, construct a
 * new Option Optic at that property P.
 */
export function nilableProp<A, P extends keyof A>(prop: P) {
  return <U extends OpticTag, S>(
    sa: Optic<U, S, A>,
  ): Optic<Join<U, PrismTag>, S, NonNullable<A[P]>> =>
    pipe(
      sa,
      compose(prism((s) =>
        O.fromNullable(s[prop]), (fii) => (a) =>
        isNotNil(a[prop])
          ? ({ ...a, [prop]: fii(a[prop] as NonNullable<A[P]>) })
          : a)),
    );
}

/**
 * Given an Optic focused on A, filter out or refine that A.
 */
export function filter<A>(
  r: Predicate<A>,
): <U extends OpticTag, S>(
  first: Optic<U, S, A>,
) => Optic<Join<U, PrismTag>, S, A>;
export function filter<A, B extends A>(
  r: Refinement<A, B>,
): <U extends OpticTag, S>(
  first: Optic<U, S, A>,
) => Optic<Join<U, PrismTag>, S, B>;
export function filter<A>(predicate: Predicate<A>) {
  return <U extends OpticTag, S>(
    first: Optic<U, S, A>,
  ): Optic<Join<U, PrismTag>, S, A> =>
    pipe(
      first,
      compose(prism(O.fromPredicate(predicate), (fii) => (a) =>
        predicate(a) ? fii(a) : a)),
    );
}

/**
 * Traverse a U using a Traversable for U. This
 */
export function traverse<T extends Kind>(T: Traversable<T>) {
  return <U extends OpticTag, S, A, B, C, D, E>(
    first: Optic<U, S, $<T, [A, B, C], [D], [E]>>,
  ): Optic<Join<U, TraversalTag>, S, A> =>
    pipe(
      first,
      compose(traversal(
        T.reduce((as, a) => {
          // Interior mutability is questionable here
          as.push(a);
          return as;
        }, [] as A[]),
        (faa) => T.map(faa),
      )),
    );
}

/**
 * Construct a set function for a given Optic
 */
export function set<U extends OpticTag, S, A>(
  optic: Optic<U, S, A>,
): (a: A) => (s: S) => S {
  return (a) => optic.modify(() => a);
}

/**
 * Given an optic focused on an Option<A>, construct
 * an Optic focused within the Option.
 */
export function some<U extends OpticTag, S, A>(
  optic: Optic<U, S, Option<A>>,
): Optic<Join<U, PrismTag>, S, A> {
  return pipe(optic, compose(prism(identity, O.map)));
}

/**
 * Given an optic focused on an Either<B, A>, construct
 * an Optic focused on the Right value of the Either.
 */
export function right<U extends OpticTag, S, B, A>(
  optic: Optic<U, S, Either<B, A>>,
): Optic<Join<U, PrismTag>, S, A> {
  return pipe(optic, compose(prism(E.getRight, E.map)));
}

/**
 * Given an optic focused on an Either<B, A>, construct
 * an Optic focused on the Left value of the Either.
 */
export function left<U extends OpticTag, S, B, A>(
  optic: Optic<U, S, Either<B, A>>,
): Optic<Join<U, PrismTag>, S, B> {
  return pipe(optic, compose(prism(E.getLeft, E.mapLeft)));
}

// -- Lets Test ---

// Nested structure
type Structure = {
  name: Option<string>;
  age?: Either<string, number>;
  first: {
    second: {
      third: number;
      fourth?: {
        fifth: {
          sixth: number;
          seventh: readonly {
            name: string;
          }[];
        };
      };
    };
  };
};

// Mock the structure
const makeStructure = (third: number): Structure => ({
  name: O.none,
  age: E.left("No age listed"),
  first: { second: { third } },
});

const makeDeepStructure = (third: number, names: string[]): Structure => ({
  name: O.some("Brandon"),
  age: E.right(37),
  first: {
    second: {
      third,
      fourth: {
        fifth: { sixth: 2 * third, seventh: names.map((name) => ({ name })) },
      },
    },
  },
});

// Get a nested value from a structure
const structure = makeStructure(5);
const deepStructure = makeDeepStructure(100, [
  "Brandon",
  "Tina",
  "Rufus",
  "Clementine",
]);

const ot1 = pipe(
  id<Structure>(),
  prop("name"),
  some,
);

console.log({
  test: "ot1",
  get: ot1.get(structure),
  getDeep: ot1.get(deepStructure),
  modify: pipe(structure, ot1.modify((a) => a.toUpperCase())),
  modifyDeep: pipe(deepStructure, ot1.modify((a) => a.toUpperCase())),
});

const et1 = pipe(
  id<Structure>(),
  nilableProp("age"),
  right,
);

console.log({
  test: "et1",
  get: et1.get(structure),
  getDeep: et1.get(deepStructure),
  modify: pipe(structure, et1.modify((a) => a + 1)),
  modifyDeep: pipe(deepStructure, et1.modify((a) => a + 1)),
});

const et2 = pipe(
  id<Structure>(),
  nilableProp("age"),
  left,
);

console.log({
  test: "et2",
  get: et2.get(structure),
  getDeep: et2.get(deepStructure),
  modify: pipe(structure, et2.modify((a) => a.toUpperCase())),
  modifyDeep: pipe(deepStructure, et2.modify((a) => a.toUpperCase())),
});

const t1 = pipe(
  id<Structure>(),
  prop("first"),
  prop("second"),
  prop("third"),
  filter((n) => n > 0),
);

const t2 = pipe(
  id<Structure>(),
  prop("first"),
  prop("second"),
  nilableProp("fourth"),
  prop("fifth"),
  prop("sixth"),
);

const _ = pipe(
  id<Structure>(),
  prop("first"),
  prop("second"),
  prop("fourth"), // Returns Optic<LensTag, Structure, { fifth: { sixth: number } } | undefined>
  // prop("fifth"), // Fails with string not assignable to never
  // prop("sixth"),
);

const t3 = pipe(
  id<Structure>(),
  prop("first"),
  prop("second"),
  nilableProp("fourth"),
  prop("fifth"),
  prop("seventh"),
  traverse(A.TraversableArray),
  prop("name"),
);

// Returns:
// {
//   structure: { first: { second: { third: 5 } } },
//   t1: { tag: "Some", value: 5 },
//   t2: { tag: "None" },
//   t3: [],
//   t3deep: [ "Brandon", "Tina", "Rufus", "Clementine" ]
// }
console.log({
  structure,
  t1: t1.get(structure),
  t2: t2.get(structure),
  t3: t3.get(structure),
  t3deep: t3.get(deepStructure),
});

const modifiedDeep = pipe(
  deepStructure,
  t3.modify((name) => name.toUpperCase()),
);

// {
//   "first": {
//     "second": {
//       "third": 100,
//       "fourth": {
//         "fifth": {
//           "sixth": 200,
//           "seventh": [
//             {
//               "name": "BRANDON"
//             },
//             {
//               "name": "TINA"
//             },
//             {
//               "name": "RUFUS"
//             },
//             {
//               "name": "CLEMENTINE"
//             }
//           ]
//         }
//       }
//     }
//   }
// }
console.log(JSON.stringify(modifiedDeep, null, 2));

const traverseFilter = pipe(t3, filter((name) => name.includes("a")));
const tfModify = pipe(
  deepStructure,
  traverseFilter.modify((name) => name.toUpperCase()),
);
console.log({
  deepStructure,
  traverseFilter: traverseFilter.get(deepStructure),
});
console.log(JSON.stringify(tfModify, null, 2));