Stack-safe traversable in fp-ts

safetraverse.ts

import * as r from "fp-ts/Reader"
import * as array from "fp-ts/ReadonlyArray"
import * as option from "fp-ts/Option"
import * as either from "fp-ts/Either"
import * as fn from "fp-ts/function"
import type * as cr from "fp-ts/ChainRec"
import type * as applicative from "fp-ts/Applicative"
import type * as hkt from "fp-ts/HKT"
import type * as functor from "fp-ts/lib/Functor"
import type * as foldable from "fp-ts/lib/Foldable"
import assert from "assert"

const DATA = array.unfold(0, n =>
  n < 100_000 ? option.of([n, n + 1]) : option.none
)

const work: (n: number) => r.Reader<number, number> = n =>
  fn.pipe(r.asks((multiplier: number) => multiplier * n))

const unsafeTraverse = array.traverse(r.Applicative)

const readerCR: cr.ChainRec2<r.URI> = {
  ...r.Chain,
  chainRec: (a, f) => env => {
    let result = f(a)(env)
    while (either.isLeft(result)) result = f(result.left)(env)
    return result.right
  },
}

const arrayST: SafeTraversable1<array.URI> = {
  ...array.Functor,
  ...array.Foldable,

  /**
   * Traverse an array using the `chainRec` implementation of the
   * provided functor.
   *
   * NB: This _does_ mutate the arrays it manages as it traverse,
   * but this mutation is entirely internal to this function and therefore
   * okay. Mutation is used to avoid overflowing the heap.
   */
  traverse:
    <F>(F: applicative.Applicative<F> & cr.ChainRec<F>) =>
    <A, B>(f: (a: A) => hkt.HKT<F, B>) =>
    (as: ReadonlyArray<A>): hkt.HKT<F, ReadonlyArray<B>> => {
      type state = [Array<A>, Array<B>]
      return F.chainRec([[...as], []] as state, ([remaining, building]) => {
        if (remaining.length === 0)
          return F.of(either.of(building as ReadonlyArray<B>))
        const a = remaining.shift() as A
        return F.map(f(a), b => {
          building.push(b)
          return either.left([remaining, building] as state)
        })
      })
    },
}

const safeTraverse = arrayST.traverse({
  ...r.Applicative,
  ...readerCR,
})

assert.throws(() => {
  console.log("running unsafe traversal")
  fn.pipe(DATA, unsafeTraverse(work), fn.apply(10))
  console.log("unsafe traversal done")
})

assert.doesNotThrow(() => {
  console.log("running safe traversal")
  fn.pipe(DATA, safeTraverse(work), fn.apply(10))
  console.log("safe traversal done")
})

interface SafeTraversable1<T extends hkt.URIS>
  extends functor.Functor1<T>,
    foldable.Foldable1<T> {
  readonly traverse: SafeTraverse1<T>
}

interface SafeTraverse1<T extends hkt.URIS> {
  <F extends hkt.URIS2>(F: applicative.Applicative2<F> & cr.ChainRec2<F>): <
    A,
    E,
    B
  >(
    f: (a: A) => hkt.Kind2<F, E, B>
  ) => (traversableOfA: hkt.Kind<T, A>) => hkt.Kind2<F, E, hkt.Kind<T, B>>
  <F>(F: applicative.Applicative<F> & cr.ChainRec<F>): <A, B>(
    f: (a: A) => hkt.HKT<F, B>
  ) => (traversableOfA: hkt.Kind<T, A>) => hkt.HKT<F, hkt.Kind<T, B>>
}