Abusing TypeScript Generators
Tags: functional, typescript
Published: 2020-11-02 18:31:25

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/](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:

```ts
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
)
```

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:

```ts
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()
}
```

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 }
```

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:

```ts
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)
}
```

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
```

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:

```ts
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)
}
```

This will print out:

```ts
0
1
2
3
4
5
6
7
8
9
```

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](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`.

```ts
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
)
```

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>
```

So you can use it like:

```ts
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)
``` 

That will produce:

```
open
15
release
```

## 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](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:

```ts
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])
}
```
From [https://github.com/Matechs-Garage/matechs-effect/blob/master/packages/system/src/Stream/Stream/gen.ts](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:

```ts
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())
})
```

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:

```ts
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
)
```

will produce:

```
[0, 1, 2]
```

## 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](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.

```ts
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)
```

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

```ts
const { get, put } = yield* _(DbConnection)
```

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!