From d41a91d2b66b27b954662d90f3d36ce5480e1710 Mon Sep 17 00:00:00 2001
From: Nathan Rajlich <nathan@tootallnate.net>
Date: Tue, 4 Feb 2020 18:34:48 -0800
Subject: [PATCH] Fix nested `yield` statements

---
 src/generator-to-promise.ts | 58 ++++++++++++++++++++++---------------
 src/index.ts                |  6 ++--
 test/test.ts                | 30 +++++++++++++++++--
 3 files changed, 66 insertions(+), 28 deletions(-)

diff --git a/src/generator-to-promise.ts b/src/generator-to-promise.ts
index cab1c2f..69053e5 100644
--- a/src/generator-to-promise.ts
+++ b/src/generator-to-promise.ts
@@ -4,7 +4,11 @@ interface Deferred<T> {
 	reject: (reason?: any) => void;
 }
 
-function isGen(fn: any): fn is GeneratorFunction {
+function isGenerator(fn: any): fn is Generator {
+	return fn && fn.next && fn.throw;
+}
+
+function isGeneratorFunction(fn: any): fn is GeneratorFunction {
 	return (
 		typeof fn == 'function' && fn.constructor.name == 'GeneratorFunction'
 	);
@@ -28,10 +32,10 @@ function createDeferred<T>(): Deferred<T> {
 	return { promise, resolve: r, reject: j };
 }
 
-export default function generatorToPromise<T>(
+export default function generatorFnToPromise<T>(
 	generatorFunction: any
 ): (...args: any[]) => Promise<T> {
-	if (!isGen(generatorFunction)) {
+	if (!isGeneratorFunction(generatorFunction)) {
 		if (typeof generatorFunction === 'function') {
 			return function(this: any, ...args: any[]) {
 				return Promise.resolve(true).then(() => {
@@ -43,29 +47,37 @@ export default function generatorToPromise<T>(
 	}
 
 	return function(this: any, ...args: any[]): Promise<T> {
-		const deferred = createDeferred<T>();
 		const generator = generatorFunction.apply(this, args);
-		(function next(err?: Error | null, value?: any) {
-			let genState = null;
-			try {
-				if (err) {
-					genState = generator.throw(err);
-				} else {
-					genState = generator.next(value);
-				}
-			} catch (e) {
-				genState = { value: Promise.reject(e), done: true };
-			}
+		return generatorToPromise(generator);
+	};
+}
 
-			if (genState.done) {
-				deferred.resolve(genState.value);
+function generatorToPromise<T>(this: any, generator: any): Promise<T> {
+	const deferred = createDeferred<T>();
+	(function next(err?: Error | null, value?: any) {
+		let genState = null;
+		try {
+			if (err) {
+				genState = generator.throw(err);
 			} else {
-				Promise.resolve(genState.value)
-					.then(promiseResult => next(null, promiseResult))
-					.catch(err => next(err));
+				genState = generator.next(value);
 			}
-		})();
+		} catch (e) {
+			genState = { value: Promise.reject(e), done: true };
+		}
 
-		return deferred.promise;
-	};
+		if (isGenerator(genState.value)) {
+			genState.value = generatorToPromise(genState.value);
+		}
+
+		if (genState.done) {
+			deferred.resolve(genState.value);
+		} else {
+			Promise.resolve(genState.value)
+				.then(promiseResult => next(null, promiseResult))
+				.catch(err => next(err));
+		}
+	})();
+
+	return deferred.promise;
 }
diff --git a/src/index.ts b/src/index.ts
index 7da3c48..4f3f68b 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -5,7 +5,7 @@ import { visit, namedTypes as n, builders as b } from 'ast-types';
 import { Context, RunningScriptOptions, runInNewContext } from 'vm';
 
 import _supportsAsync from './supports-async';
-import generatorToPromise from './generator-to-promise';
+import generatorToPromiseFn from './generator-to-promise';
 
 /**
  * Compiles sync JavaScript code into JavaScript with async Functions.
@@ -174,11 +174,11 @@ namespace degenerator {
 		if (isAsyncFunction(fn)) {
 			return fn;
 		} else {
-			const rtn = (generatorToPromise(fn) as unknown) as T;
+			const rtn = (generatorToPromiseFn(fn) as unknown) as T;
 			Object.defineProperty(rtn, 'toString', {
 				value: fn.toString.bind(fn),
 				enumerable: false
-			})
+			});
 			return rtn;
 		}
 	}
diff --git a/test/test.ts b/test/test.ts
index 009b0da..49a133a 100644
--- a/test/test.ts
+++ b/test/test.ts
@@ -134,9 +134,9 @@ describe('degenerator()', () => {
 				}
 			);
 			if (supportsAsync) {
-				assert(/await b\(\)/.test(fn+''));
+				assert(/await b\(\)/.test(fn + ''));
 			} else {
-				assert(/yield b\(\)/.test(fn+''));
+				assert(/yield b\(\)/.test(fn + ''));
 			}
 		});
 		it('should be able to await non-promises', () => {
@@ -188,5 +188,31 @@ describe('degenerator()', () => {
 				'Expected a "function" to be returned for `foo`, but got "number"'
 			);
 		});
+		it('should be compile if branches', () => {
+			function ifA(): string {
+				if (a()) {
+					return 'foo';
+				}
+				return 'bar';
+			}
+			function a() {
+				if (b()) {
+					return false;
+				}
+				return true;
+			}
+			function b() {
+				return false;
+			}
+			const fn = compile<() => Promise<string>>(
+				`${ifA};${a}`,
+				'ifA',
+				['b'],
+				{ sandbox: { b } }
+			);
+			return fn().then((val: string) => {
+				assert.equal(val, 'foo');
+			});
+		});
 	});
 });