trampoline-ts/src/trampoline.ts

A type-safe way to emulate tail-call optimization with trampolines

import { ArgumentTypes } from './types';
import {
  Thunk,
  UnwrapThunkDeep,
  isThunk,
  toThunk,
  ThunkOrValue,
} from './thunk';

export type UnwrapPromise<T> = T extends Promise<infer U> ? Exclude<U, Promise<T>> : T;

export type Unbox<T> = UnwrapThunkDeep<UnwrapPromise<T>>;

export type Cont<A extends any[], R> = (...args: A) => Thunk<Unbox<R>>;

export interface Trampoline<F extends ((...args: any[]) => any)> {
  (...args: ArgumentTypes<F>): Unbox<ReturnType<F>>;
  cont: Cont<ArgumentTypes<F>, ReturnType<F>>;
}

export interface TrampolineAsync<F extends ((...args: any[]) => any)> {
  (...args: ArgumentTypes<F>): Promise<Unbox<ReturnType<F>>>;
  cont: Cont<ArgumentTypes<F>, ReturnType<F>>;
}

export const trampoline = <F extends ((...args: any[]) => any)>(fn: F): Trampoline<F> => {
  const cont = (...args: ArgumentTypes<F>) => toThunk(() => fn(...args));

  return Object.assign(
    (...args: ArgumentTypes<F>): Unbox<ReturnType<F>> => {
      let result: ThunkOrValue<ReturnType<F>> = fn(...args);

      while (isThunk<ReturnType<F>>(result)) {
        result = result();
      }

      return result;
    },
    { cont },
  );
};

export const trampolineAsync = <F extends ((...args: any[]) => any)>(fn: F): TrampolineAsync<F> => {
  const cont = (...args: ArgumentTypes<F>) => toThunk(() => fn(...args));

  return Object.assign(
    async (...args: ArgumentTypes<F>): Promise<Unbox<ReturnType<F>>> => {
      let result: ThunkOrValue<ReturnType<F>> = await fn(...args);

      while (isThunk<ReturnType<F>>(result)) {
        result = await result();
      }

      return result;
    },
    { cont },
  );
};