functional-programming-snake/js/solutions/safeio.js

// m07b safeio: a stack-safe IO. Same idea as m07's IO, plus one more constructor, so that `flatMap` only builds
// data and no program shape can overflow the stack, however deeply its `flatMap`s are nested.
//   {tag:'Pure', value}       - done, with a value
//   {tag:'Effect', op, k}     - perform `op`, then continue with `k(answer)`
//   {tag:'FlatMap', io, f}    - run `io`, then continue with `f(result)`
// Ops are the same plain data as in m07: {tag:'Draw', shape} (answer undefined), {tag:'RandomSeed'} (answer a number);
// the pure interpreter throws for any other op.
/** @import { Op, SafeIO, Trace } from './types' */

/** @type {<A>(value: A) => SafeIO<A>} */
export const pure = (value) => ({ tag: 'Pure', value })

/** @type {<A>(op: Op) => SafeIO<A>} */
export const effect = (op) => ({ tag: 'Effect', op, k: pure })

// Only builds a node: never calls f, not even when io is Pure.
/** @type {<A, B>(io: SafeIO<A>, f: (a: A) => SafeIO<B>) => SafeIO<B>} */
export const flatMap = (io, f) => ({ tag: 'FlatMap', io, f })

/** @type {<A, B>(io: SafeIO<A>, f: (a: A) => B) => SafeIO<B>} */
export const map = (io, f) => flatMap(io, (a) => pure(f(a)))

// A left fold on purpose: with SafeIO the nesting no longer matters.
/** @type {(ios: SafeIO<any>[]) => SafeIO<undefined>} */
export const sequence = (ios) => ios.reduce((acc, io) => flatMap(acc, () => io), pure(undefined))

// One rewrite towards a program whose top is a Pure or an Effect. Constant work, no recursion.
//   FlatMap(Pure(a), f)          -> f(a)
//   FlatMap(Effect(op, k), f)    -> Effect(op, r => FlatMap(k(r), f))
//   FlatMap(FlatMap(m, g), f)    -> FlatMap(m, x => FlatMap(g(x), f))     (re-association)
// Pure and Effect are already in that shape and come back unchanged.
/** @type {<A>(io: SafeIO<A>) => SafeIO<A>} */
export const step = (io) => {
  if (io.tag !== 'FlatMap') return io
  const { io: inner, f } = io
  if (inner.tag === 'Pure') return f(inner.value)
  if (inner.tag === 'Effect') return { tag: 'Effect', op: inner.op, k: (r) => flatMap(inner.k(r), f) }
  return flatMap(inner.io, (x) => flatMap(inner.f(x), f))
}

// Given (not an exercise): a trampoline driver without loops or mutation.
// `next(s)` returns `[s2, outs]`. Applies `next` until `isDone(s)`, returns `[last, allOuts]` (outs in order).
// Running `n` moves is running `n/2` moves and then `n/2` more, so the recursion is only log2(n) deep; when the
// budget runs out before the end, the driver doubles it and goes on. The outputs are kept in a balanced tree
// and flattened once at the end, so no array is copied once per move.
const leaf = (items) => ({ tag: 'Leaf', items })
const node = (l, r) => ({ tag: 'Node', l, r })
const toArray = (t) => (t.tag === 'Leaf' ? t.items : toArray(t.l).concat(toArray(t.r)))

const runFor = (s, n, isDone, next) => {
  if (isDone(s)) return { state: s, outs: leaf([]) }
  if (n === 1) {
    const [s2, outs] = next(s)
    return { state: s2, outs: leaf(outs) }
  }
  const half = Math.floor(n / 2)
  const first = runFor(s, half, isDone, next)
  const second = runFor(first.state, n - half, isDone, next)
  return { state: second.state, outs: node(first.outs, second.outs) }
}

/** @type {<S, O>(start: S, isDone: (s: S) => boolean, next: (s: S) => [S, O[]]) => [S, O[]]} */
export const drive = (start, isDone, next) => {
  const go = (s, budget, acc) => {
    const r = runFor(s, budget, isDone, next)
    const outs = node(acc, r.outs)
    return isDone(r.state) ? [r.state, toArray(outs)] : go(r.state, budget * 2, outs)
  }
  return go(start, 1, leaf([]))
}

const answer = (op, seed) => {
  if (op.tag === 'RandomSeed') return seed
  if (op.tag === 'Draw') return undefined
  throw new Error("SafeIO.trace answers only Module 7's operations")
}

// One move of the pure interpreter: rewrite a FlatMap, or answer an Effect and record its op.
/** @type {<A>(io: SafeIO<A>, seed: number) => [SafeIO<A>, Op[]]} */
export const traceStep = (io, seed) => {
  if (io.tag === 'FlatMap') return [step(io), []]
  if (io.tag === 'Effect') return [io.k(answer(io.op, seed)), [io.op]]
  return [io, []]
}

/** @type {<A>(io: SafeIO<A>, seed: number) => Trace<A>} */
export const trace = (io, seed) => {
  const [last, ops] = drive(io, (p) => p.tag === 'Pure', (p) => traceStep(p, seed))
  return { result: /** @type {any} */ (last).value, ops } // drive stops only at a Pure
}