TypeScript implementation of Haskell's do-notation

syntax.test.ts

import { identity, Identity } from 'fp-ts/lib/Identity';
import { io, IO } from 'fp-ts/lib/IO';

import { doM } from './syntax';

describe('Monad do-syntax', () => {
  beforeAll(() => {
    jest.spyOn(Identity.prototype, 'chain');
    jest.spyOn(IO.prototype, 'chain');
  });

  afterAll(() => {
    jest.resetAllMocks();
  });

  it('should be able to execute computations in context of arbitrary monad', () => {
    const id: IO<number> = doM(io)<number>(function *() {
      const begin = yield io.of(21);
      const add = yield io.of(begin).ap(io.of((x: number) => x + 2));
      return io.of(add);
    });

    const res = id.run();
    expect(res).toEqual(23);
    expect((IO.prototype.chain as jest.Mock).mock.calls.length).toEqual(2);
  });

  it('should accept generator and call .chain() sequentally', () => {
    const id: Identity<number> = doM(identity)<number>(function *() {
      const begin = yield identity.of(41);
      const add = yield identity.of(begin).ap(identity.of((x: number) => x + 1));
      return identity.of(add);
    });

    expect(id.extract()).toEqual(42);
    expect((Identity.prototype.chain as jest.Mock).mock.calls.length).toEqual(2);
  });
});

syntax.ts

import { Lazy } from 'fp-ts/lib/function';
import { URIS, HKT, Type } from 'fp-ts/lib/HKT';
import { Chain, Chain1 } from 'fp-ts/lib/Chain';

export function doM<F extends URIS>(F: Chain1<F>): <A>(gen: Lazy<IterableIterator<Type<F, A>>>) => Type<F, A>;
export function doM<F>(F: Chain<F>): <A>(gen: Lazy<IterableIterator<HKT<F, A>>>) => HKT<F, A> {
  return <A>(gen: Lazy<IterableIterator<HKT<F, A>>>) => {
    let g = gen();
    const step = (value?: A): HKT<F, A> => {
      const result = g.next(value);
      if (result.done) {
        g = gen();
        return result.value;
      } else {
        return F.chain(result.value, step);
      }
    };
    return step();
  };
}