From fee31d26465b16ee60f822d9aa4e551b17dd3281 Mon Sep 17 00:00:00 2001
From: kschat <kyleschattler@gmail.com>
Date: Fri, 26 Apr 2019 14:21:51 -0400
Subject: [PATCH] feat: support async/Promise-returning functions

---
 src/thunk.ts            |  4 ++--
 src/trampoline.ts       | 51 +++++++++++++++++++++++++++++++----------
 test/thunk.spec.ts      | 16 ++++++-------
 test/trampoline.spec.ts | 15 +++++++++++-
 4 files changed, 63 insertions(+), 23 deletions(-)

diff --git a/src/thunk.ts b/src/thunk.ts
index b4bff56..d574632 100644
--- a/src/thunk.ts
+++ b/src/thunk.ts
@@ -7,8 +7,8 @@ export interface Thunk<T> extends Function {
 
 export type ThunkOrValue<T> = T | Thunk<T>;
 
-export type UnwrapThunk<T> = {
-  0: T extends Thunk<infer U> ? UnwrapThunk<U> : T;
+export type UnwrapThunkDeep<T> = {
+  0: T extends Thunk<infer U> ? UnwrapThunkDeep<U> : T;
 }[
   T extends ThunkOrValue<T> ? 0 : never
 ];
diff --git a/src/trampoline.ts b/src/trampoline.ts
index 714dcf8..d55ff18 100644
--- a/src/trampoline.ts
+++ b/src/trampoline.ts
@@ -1,31 +1,58 @@
 import { ArgumentTypes } from './types';
 import {
   Thunk,
-  UnwrapThunk,
+  UnwrapThunkDeep,
   isThunk,
   toThunk,
   ThunkOrValue,
 } from './thunk';
 
-export type Cont<A extends any[], R> = (...args: A) => Thunk<UnwrapThunk<R>>;
+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>): UnwrapThunk<ReturnType<F>>;
+  (...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 trampolineFunction = (...args: ArgumentTypes<F>): UnwrapThunk<ReturnType<F>> => {
-    let result: ThunkOrValue<ReturnType<F>> = fn(...args);
+  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 },
+  );
+};
 
-    while (isThunk<ReturnType<F>>(result)) {
-      result = result();
-    }
+export const trampolineAsync = <F extends ((...args: any[]) => any)>(fn: F): TrampolineAsync<F> => {
+  const cont = (...args: ArgumentTypes<F>) => toThunk(() => fn(...args));
 
-    return result;
-  };
+  return Object.assign(
+    async (...args: ArgumentTypes<F>): Promise<Unbox<ReturnType<F>>> => {
+      let result: ThunkOrValue<ReturnType<F>> = await fn(...args);
 
-  trampolineFunction.cont = (...args: ArgumentTypes<F>) => toThunk(() => fn(...args));
+      while (isThunk<ReturnType<F>>(result)) {
+        result = await result();
+      }
 
-  return trampolineFunction;
+      return result;
+    },
+    { cont },
+  );
 };
diff --git a/test/thunk.spec.ts b/test/thunk.spec.ts
index 98d8b31..f1dfc74 100644
--- a/test/thunk.spec.ts
+++ b/test/thunk.spec.ts
@@ -3,7 +3,7 @@ import {
   isThunk,
   toThunk,
   THUNK_SYMBOL,
-  UnwrapThunk,
+  UnwrapThunkDeep,
   Thunk,
   ThunkOrValue,
 } from '../src/thunk';
@@ -32,13 +32,13 @@ describe('thunk', () => {
 
   describe('UnwrapThunk<T>', () => {
     it('removes recursively removes all instances of Thunk<T> from T', () => {
-      typeAssert<IsExact<UnwrapThunk<1>, 1>>(true);
-      typeAssert<IsExact<UnwrapThunk<Thunk<1>>, 1>>(true);
-      typeAssert<IsExact<UnwrapThunk<ThunkOrValue<1>>, 1>>(true);
-      typeAssert<IsExact<UnwrapThunk<ThunkOrValue<Thunk<1>>>, 1>>(true);
-      typeAssert<IsExact<UnwrapThunk<ThunkOrValue<ThunkOrValue<1> | ThunkOrValue<2>>>, 1 | 2>>(true);
-      typeAssert<IsExact<UnwrapThunk<Thunk<1 | Thunk<2>>>, 1 | 2>>(true);
-      typeAssert<IsExact<UnwrapThunk<Thunk<1 | Thunk<2> | ThunkOrValue<3>>>, 1 | 2 | 3>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<1>, 1>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<Thunk<1>>, 1>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<ThunkOrValue<1>>, 1>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<ThunkOrValue<Thunk<1>>>, 1>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<ThunkOrValue<ThunkOrValue<1> | ThunkOrValue<2>>>, 1 | 2>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<Thunk<1 | Thunk<2>>>, 1 | 2>>(true);
+      typeAssert<IsExact<UnwrapThunkDeep<Thunk<1 | Thunk<2> | ThunkOrValue<3>>>, 1 | 2 | 3>>(true);
     });
   });
 });
diff --git a/test/trampoline.spec.ts b/test/trampoline.spec.ts
index 6fb54a4..00f596f 100644
--- a/test/trampoline.spec.ts
+++ b/test/trampoline.spec.ts
@@ -1,5 +1,5 @@
 import { assert as typeAssert, IsExact, Has } from 'conditional-type-checks';
-import { trampoline, isThunk, ThunkOrValue } from '../src';
+import { trampoline, isThunk, ThunkOrValue, trampolineAsync } from '../src';
 import { ArgumentTypes } from '../src/types';
 
 describe('trampoline', () => {
@@ -99,5 +99,18 @@ describe('trampoline', () => {
       expect(fn()).toBeInstanceOf(Function);
       expect(fn()()).toBe('hello');
     });
+
+    it('supports async functions', async () => {
+      const sleep = async (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms));
+
+      const factorial = trampolineAsync(async (n: number, acc: number = 1): Promise<ThunkOrValue<number>> => {
+        await sleep(10);
+        return n
+          ? factorial.cont(n - 1, acc * n)
+          : acc;
+      });
+
+      expect(await factorial(2)).toBe(2);
+    });
   });
 });