Abusing TypeScript Generators

Michael Arnaldi - Nov 2 '20 - - Dev Community

In the previous chapters of the series we explored the @effect-ts/core ecosystem and some of its data-types like Effect, Managed, Layer and more.

We have explored, at length, the advantages of using statically typed functional programming in typescript to gain modularity, testability and safety of our code.

One thing is missing though, the API we have been using feels a bit awkward for people that come from non-functional languages, for the ones that come from functional languages there is a big missing, namely in scala for, in haskell do.

For the newcomers, let's first explore the reason behind this apis.

We have seen how Monads can represent the concept of sequential operations and we have been using chain to express this all over the place.

This style of programming highlight the function composition aspect, we would like to have a way to also express sequential steps in an imperative way.

For the people coming from javascript or typescript we would like to have something similar to what async/await does but generalised to every Monad.

Early approaches

Historically there have been a significant amount of attempts at approximating Do in the context of functional programming in typescript.

Namely we had 2 approaches:

  • the first one by Paul Grey detailed explained in his article at: https://paulgray.net/do-syntax-in-typescript/ using a fluent based api

  • the second one by Giulio Canti using a pipeable based api (that we also have as an alternative)

Let's look at the second one:

import * as T from "@effect-ts/core/Effect"
import { pipe } from "@effect-ts/core/Function"

const program = pipe(
  T.do,
  T.bind("a", () => T.succeed(1)),
  T.bind("b", () => T.succeed(2)),
  T.bind("c", ({ a, b }) => T.succeed(a + b)),
  T.map(({ c }) => c)
)

pipe(
  program,
  T.chain((c) =>
    T.effectTotal(() => {
      console.log(c)
    })
  ),
  T.runMain
)
Enter fullscreen mode Exit fullscreen mode

Basically what happens here is T.do initializes an empty state {} to hold the scope of the computation and bind uses it to progressively fill up the variables step by step.

This way each step has access to the whole set of variables defined up to that point.

Both the fluent & the pipeable APIs are really nice to use but still doesn't feel native typescript and there is a good degree of repetition in accessing the scope explicitly at every bind operation.

Enter the world of generators

Before anything in order to use this API you will need to enable "downlevelIteration": true in your tsconfig.json.

Let's start to explore what are generators:

function* countTo(n: number) {
  let i = 0
  while (i < n) {
    yield i++
  }
}

const iterator = countTo(10)

let current = iterator.next()

while (!current.done) {
  console.log(current)

  current = iterator.next()
}
Enter fullscreen mode Exit fullscreen mode

This will print out:

{ value: 0, done: false }
{ value: 1, done: false }
{ value: 2, done: false }
{ value: 3, done: false }
{ value: 4, done: false }
{ value: 5, done: false }
{ value: 6, done: false }
{ value: 7, done: false }
{ value: 8, done: false }
{ value: 9, done: false }
Enter fullscreen mode Exit fullscreen mode

So basically every yield give us a value and execution is controlled by the caller by calling consuming the generator.

Let's now proceed with the second thing to know:

function* countTo(n: number) {
  let i = 0
  while (i < n) {
    const a = yield i++

    console.log(a)
  }
}

const iterator = countTo(10)

let input = "a"
let current = iterator.next(input)

while (!current.done) {
  console.log(current)

  input += "a"
  current = iterator.next(input)
}
Enter fullscreen mode Exit fullscreen mode

This will print out:

{ value: 0, done: false }
aa
{ value: 1, done: false }
aaa
{ value: 2, done: false }
aaaa
{ value: 3, done: false }
aaaaa
{ value: 4, done: false }
aaaaaa
{ value: 5, done: false }
aaaaaaa
{ value: 6, done: false }
aaaaaaaa
{ value: 7, done: false }
aaaaaaaaa
{ value: 8, done: false }
aaaaaaaaaa
{ value: 9, done: false }
aaaaaaaaaaa
Enter fullscreen mode Exit fullscreen mode

So basically with a generator we can push the return values of yields from the consumer to populate the variable scope.

The types are just horrible, const a in the generator body is any and there is really no inference working.

Another feature that we are going to exploit is yield* to yield a generator into a generator, surprisingly this works well type wise
:)

Let's see:

function* constant<A>(a: A): Generator<A, A, A> {
  return yield a
}

function* countTo(n: number) {
  let i = 0
  while (i < n) {
    const a = yield* constant(i++)

    console.log(a)
  }
}

const iterator = countTo(10)

let current = iterator.next()

while (!current.done) {
  current = iterator.next(current.value)
}
Enter fullscreen mode Exit fullscreen mode

This will print out:

0
1
2
3
4
5
6
7
8
9
Enter fullscreen mode Exit fullscreen mode

and all the types work well, in the signature of constant we read Generator<A, A, A>, the first A is the type of the yield, the second is the type of the yielded value and the third is the type required by the next step.

The idea here is to use a generator to declare the body of a computation in a monadic context and then, at runtime, starve the generator building up a set of chained computations (initial attempts by https://github.com/nythrox and some others, forgive me for not remembering all).

Generator based do for Effect

Let's put together what we seen so far, we would like to avoid modifying the types of Effect and in general any type so we won't directly add a generator inside them. That would also break variance of the type.

We are going to instead use a function to lift an effect into a generator of effect.

import * as T from "@effect-ts/core/Effect"
import { pipe } from "@effect-ts/core/Function"
import type { _E, _R } from "@effect-ts/core/Utils"

export class GenEffect<K, A> {
  constructor(readonly op: K) {}

  *[Symbol.iterator](): Generator<GenEffect<K, A>, A, any> {
    return yield this
  }
}

const adapter = (_: any) => {
  return new GenEffect(_)
}

export function gen<Eff extends GenEffect<any, any>, AEff>(
  f: (i: {
    <R, E, A>(_: T.Effect<R, E, A>): GenEffect<T.Effect<R, E, A>, A>
  }) => Generator<Eff, AEff, any>
): T.Effect<_R<Eff["op"]>, _E<Eff["op"]>, AEff> {
  return T.suspend(() => {
    const iterator = f(adapter as any)
    const state = iterator.next()

    function run(
      state: IteratorYieldResult<Eff> | IteratorReturnResult<AEff>
    ): T.Effect<any, any, AEff> {
      if (state.done) {
        return T.succeed(state.value)
      }
      return T.chain_(state.value["op"], (val) => {
        const next = iterator.next(val)
        return run(next)
      })
    }

    return run(state)
  })
}

const program = gen(function* (_) {
  const a = yield* _(T.succeed(1))
  const b = yield* _(T.succeed(2))

  return a + b
})

pipe(
  program,
  T.chain((n) =>
    T.effectTotal(() => {
      console.log(n)
    })
  ),
  T.runMain
)
Enter fullscreen mode Exit fullscreen mode

And there we go, we have our nice imperative style of composing effects!

One thing I learned is: "if you have restrictions at the api level, find opportunities to exploit them"

In this case we have an unwanted _ function that we use to convert an effect to a generator of effects we can exploit that to lift a generic natural transformation from any type we want to effect, in the @effect-ts/core/Effect package the gen function can directly deal with:

f: (i: {
    <A>(_: Tag<A>): GenEffect<Has<A>, never, A>
    <E, A>(_: Option<A>, onNone: () => E): GenEffect<unknown, E, A>
    <A>(_: Option<A>): GenEffect<unknown, NoSuchElementException, A>
    <E, A>(_: Either<E, A>): GenEffect<unknown, E, A>
    <R, E, A>(_: Effect<R, E, A>): GenEffect<R, E, A>
    <R, E, A>(_: Managed<R, E, A>): GenEffect<R, E, A>
  }) => Generator<Eff, AEff, any>
Enter fullscreen mode Exit fullscreen mode

So you can use it like:

import * as E from "@effect-ts/core/Classic/Either"
import * as O from "@effect-ts/core/Classic/Option"
import * as T from "@effect-ts/core/Effect"
import * as M from "@effect-ts/core/Effect/Managed"
import { pipe } from "@effect-ts/core/Function"

const result = T.gen(function* (_) {
  const a = yield* _(O.some(1))
  const b = yield* _(O.some(2))
  const c = yield* _(E.right(3))
  const d = yield* _(T.access((_: { n: number }) => _.n))
  const e = yield* _(
    M.makeExit_(
      T.effectTotal(() => {
        console.log("open")

        return 5
      }),
      () =>
        T.effectTotal(() => {
          console.log("release")
        })
    )
  )

  yield* _(
    T.effectTotal(() => {
      console.log(a + b + c + d + e)
    })
  )
})

pipe(result, T.provideAll({ n: 4 }), T.runMain)
Enter fullscreen mode Exit fullscreen mode

That will produce:

open
15
release
Enter fullscreen mode Exit fullscreen mode

Multi-Shot Monads

The trick explored so far works well for effects that produce a single value, like Effect, Managed and any io-like effect type.

There is a problem with multi-shot effects like Stream or Array, for those the approach taken doesn't work, as we work through the computation above we notice how we progress the iterator at every step of the execution, in a stream or in an array we would have to replay the same for many elements but we have only one iterator.

Iterators are mutable so there is no way we can "clone" the iterator or do any form of defensive copy.

That is not the end though, supposing the body of the generator is pure, as in none of the lines of code outside of the ones inside the yields can modify any external variable or perform any side-effect, we can solve the issue by keeping the stack of computations around and by replaying the iterator every time (idea of https://github.com/mattiamanzati).

This technique has a performance complexity of O(n^2) where n is the number of yields (not the number of elements produced) so this must be taken into account when building large generators.

Let's look at the code of the Stream generator:

import type { Effect } from "../../Effect"
import { fromEither, service } from "../../Effect"
import { die } from "../../Effect/die"
import type { Either } from "../../Either"
import { NoSuchElementException, PrematureGeneratorExit } from "../../GlobalExceptions"
import type { Has, Tag } from "../../Has"
import * as L from "../../List"
import type { Option } from "../../Option"
import type { _E, _R } from "../../Utils"
import { isEither, isOption, isTag } from "../../Utils"
import { chain_ } from "./chain"
import { Stream } from "./definitions"
import { fail } from "./fail"
import { fromEffect } from "./fromEffect"
import { succeed } from "./succeed"
import { suspend } from "./suspend"

export class GenStream<R, E, A> {
  readonly _R!: (_R: R) => void
  readonly _E!: () => E
  readonly _A!: () => A
  constructor(readonly effect: Stream<R, E, A>) {}
  *[Symbol.iterator](): Generator<GenStream<R, E, A>, A, any> {
    return yield this
  }
}

const adapter = (_: any, __?: any) => {
  if (isOption(_)) {
    return new GenStream(
      _._tag === "None"
        ? fail(__ ? __() : new NoSuchElementException())
        : succeed(_.value)
    )
  } else if (isEither(_)) {
    return new GenStream(fromEffect(fromEither(() => _)))
  } else if (_ instanceof Stream) {
    return new GenStream(_)
  } else if (isTag(_)) {
    return new GenStream(fromEffect(service(_)))
  }
  return new GenStream(fromEffect(_))
}

export function gen<RBase, EBase, AEff>(): <Eff extends GenStream<RBase, EBase, any>>(
  f: (i: {
    <A>(_: Tag<A>): GenStream<Has<A>, never, A>
    <E, A>(_: Option<A>, onNone: () => E): GenStream<unknown, E, A>
    <A>(_: Option<A>): GenStream<unknown, NoSuchElementException, A>
    <E, A>(_: Either<E, A>): GenStream<unknown, E, A>
    <R, E, A>(_: Effect<R, E, A>): GenStream<R, E, A>
    <R, E, A>(_: Stream<R, E, A>): GenStream<R, E, A>
  }) => Generator<Eff, AEff, any>
) => Stream<_R<Eff>, _E<Eff>, AEff>
export function gen<EBase, AEff>(): <Eff extends GenStream<any, EBase, any>>(
  f: (i: {
    <A>(_: Tag<A>): GenStream<Has<A>, never, A>
    <E, A>(_: Option<A>, onNone: () => E): GenStream<unknown, E, A>
    <A>(_: Option<A>): GenStream<unknown, NoSuchElementException, A>
    <E, A>(_: Either<E, A>): GenStream<unknown, E, A>
    <R, E, A>(_: Effect<R, E, A>): GenStream<R, E, A>
    <R, E, A>(_: Stream<R, E, A>): GenStream<R, E, A>
  }) => Generator<Eff, AEff, any>
) => Stream<_R<Eff>, _E<Eff>, AEff>
export function gen<AEff>(): <Eff extends GenStream<any, any, any>>(
  f: (i: {
    <A>(_: Tag<A>): GenStream<Has<A>, never, A>
    <E, A>(_: Option<A>, onNone: () => E): GenStream<unknown, E, A>
    <A>(_: Option<A>): GenStream<unknown, NoSuchElementException, A>
    <E, A>(_: Either<E, A>): GenStream<unknown, E, A>
    <R, E, A>(_: Effect<R, E, A>): GenStream<R, E, A>
    <R, E, A>(_: Stream<R, E, A>): GenStream<R, E, A>
  }) => Generator<Eff, AEff, any>
) => Stream<_R<Eff>, _E<Eff>, AEff>
export function gen<Eff extends GenStream<any, any, any>, AEff>(
  f: (i: {
    <A>(_: Tag<A>): GenStream<Has<A>, never, A>
    <E, A>(_: Option<A>, onNone: () => E): GenStream<unknown, E, A>
    <A>(_: Option<A>): GenStream<unknown, NoSuchElementException, A>
    <E, A>(_: Either<E, A>): GenStream<unknown, E, A>
    <R, E, A>(_: Effect<R, E, A>): GenStream<R, E, A>
    <R, E, A>(_: Stream<R, E, A>): GenStream<R, E, A>
  }) => Generator<Eff, AEff, any>
): Stream<_R<Eff>, _E<Eff>, AEff>
export function gen(...args: any[]): any {
  function gen_<Eff extends GenStream<any, any, any>, AEff>(
    f: (i: any) => Generator<Eff, AEff, any>
  ): Stream<_R<Eff>, _E<Eff>, AEff> {
    return suspend(() => {
      function run(replayStack: L.List<any>): Stream<any, any, AEff> {
        const iterator = f(adapter as any)
        let state = iterator.next()
        for (const a of replayStack) {
          if (state.done) {
            return fromEffect(die(new PrematureGeneratorExit()))
          }
          state = iterator.next(a)
        }
        if (state.done) {
          return succeed(state.value)
        }
        return chain_(state.value["effect"], (val) => {
          return run(L.append_(replayStack, val))
        })
      }
      return run(L.empty())
    })
  }

  if (args.length === 0) {
    return (f: any) => gen_(f)
  }
  return gen_(args[0])
}
Enter fullscreen mode Exit fullscreen mode

From https://github.com/Matechs-Garage/matechs-effect/blob/master/packages/system/src/Stream/Stream/gen.ts

Apart from all the overloadings and the adapter code to support different monadic values the key part is:

suspend(() => {
  function run(replayStack: L.List<any>): Stream<any, any, AEff> {
    const iterator = f(adapter as any)
    let state = iterator.next()
    for (const a of replayStack) {
      if (state.done) {
        return fromEffect(die(new PrematureGeneratorExit()))
      }
      state = iterator.next(a)
    }
    if (state.done) {
      return succeed(state.value)
    }
    return chain_(state.value["effect"], (val) => {
      return run(L.append_(replayStack, val))
    })
  }
  return run(L.empty())
})
Enter fullscreen mode Exit fullscreen mode

As we can see we carry around a stack of result and we reconstruct the local scope at each call.

This can be used just like the previous one:

import * as E from "@effect-ts/core/Classic/Either"
import * as O from "@effect-ts/core/Classic/Option"
import * as T from "@effect-ts/core/Effect"
import * as S from "@effect-ts/core/Effect/Stream"
import { pipe } from "@effect-ts/core/Function"

const result = S.gen(function* (_) {
  const a = yield* _(O.some(0))
  const b = yield* _(E.right(1))
  const c = yield* _(T.succeed(2))
  const d = yield* _(S.fromArray([a, b, c]))

  return d
})

pipe(
  result,
  S.runCollect,
  T.chain((res) =>
    T.effectTotal(() => {
      console.log(res)
    })
  ),
  T.runMain
)
Enter fullscreen mode Exit fullscreen mode

will produce:

[0, 1, 2]
Enter fullscreen mode Exit fullscreen mode

Generalisation

Thanks to the HKT structure and the prelude described in the previous chapters we have been able to generalise this approach to work with any monad, the generic code comes in 2 functions available at: https://github.com/Matechs-Garage/matechs-effect/blob/master/packages/core/src/Prelude/DSL/gen.ts namely genF (to be used in one-shot cases, very efficient) and genWithHistoryF (to be used in multi-shot cases, n^2 complexity for the n-th yield).

Bonus:

Let's use all we've seen so far and build up a simulation of a simple program using 2 services, we'll simulate a message broker that flushes the messages on completion and a database that keeps state.

import "@effect-ts/core/Operators"

import * as Array from "@effect-ts/core/Classic/Array"
import * as Map from "@effect-ts/core/Classic/Map"
import * as T from "@effect-ts/core/Effect"
import * as L from "@effect-ts/core/Effect/Layer"
import * as M from "@effect-ts/core/Effect/Managed"
import * as Ref from "@effect-ts/core/Effect/Ref"
import type { _A } from "@effect-ts/core/Utils"
import { tag } from "@effect-ts/system/Has"

// make Database Live
export const makeDbLive = M.gen(function* (_) {
  const ref = yield* _(
    Ref.makeRef<Map.Map<string, string>>(Map.empty)["|>"](
      M.make((ref) => ref.set(Map.empty))
    )
  )

  return {
    get: (k: string) => ref.get["|>"](T.map(Map.lookup(k)))["|>"](T.chain(T.getOrFail)),
    put: (k: string, v: string) => ref["|>"](Ref.update(Map.insert(k, v)))
  }
})

// simulate a database connection to a key-value store
export interface DbConnection extends _A<typeof makeDbLive> {}

// Tag<DbConnection>
export const DbConnection = tag<DbConnection>()

// Database Live Layer
export const DbLive = L.fromManaged(DbConnection)(makeDbLive)

// make Broker Live
export const makeBrokerLive = M.gen(function* (_) {
  const ref = yield* _(
    Ref.makeRef<Array.Array<string>>(Array.empty)["|>"](
      M.make((ref) =>
        ref.get["|>"](
          T.chain((messages) =>
            T.effectTotal(() => {
              console.log(`Flush:`)
              messages.forEach((message) => {
                console.log("- " + message)
              })
            })
          )
        )
      )
    )
  )

  return {
    send: (message: string) => ref["|>"](Ref.update(Array.snoc(message)))
  }
})

// simulate a connection to a message broker
export interface BrokerConnection extends _A<typeof makeBrokerLive> {}

// Tag<BrokerConnection>
export const BrokerConnection = tag<BrokerConnection>()

// Broker Live Layer
export const BrokerLive = L.fromManaged(BrokerConnection)(makeBrokerLive)

// Main Live Layer
export const ProgramLive = L.all(DbLive, BrokerLive)

// Program Entry
export const main = T.gen(function* (_) {
  const { get, put } = yield* _(DbConnection)
  const { send } = yield* _(BrokerConnection)

  yield* _(put("ka", "a"))
  yield* _(put("kb", "b"))
  yield* _(put("kc", "c"))

  const a = yield* _(get("ka"))
  const b = yield* _(get("kb"))
  const c = yield* _(get("kc"))

  const s = `${a}-${b}-${c}`

  yield* _(send(s))

  return s
})

// run the program and print the output
main["|>"](T.provideSomeLayer(ProgramLive))["|>"](T.runMain)
Enter fullscreen mode Exit fullscreen mode

We can see how easy it is to access services using this generator based approach:

const { get, put } = yield* _(DbConnection)
Enter fullscreen mode Exit fullscreen mode

Directly yielding the Tag of a service will give you access to it's content.

Also if we highlight through the code we can see that the full types are correctly inferred and almost never explicitly specified.

Stay tuned

In the next article of the series, in 2 weeks time, we will continue to explore the data types available in the @effect-ts/core ecosystem!

. . . . . . . . . . . . . . . . .