Hoist await and yield expressions

The helper wraps the original argument with a new function. This is a
problem if the argument contains yield or await expressions, since the
function they now run is is not a generator or marked as asynchronous.

Instead hoist these expressions so they're resolved outside of the
helper.

Fixes avajs/ava#987.

Author: novemberborn

index.js

@@ -31,6 +31,8 @@ var buildHelper = template([
 	'}'
 ].join('\n'));
 
+var hoistExpression = template('var ARG = EXP;');
+
 var assertionVisitor = {
 	CallExpression: function (path, state) {
 		if (isThrowsMember(path.get('callee'))) {
@@ -40,14 +42,24 @@ var assertionVisitor = {
 				return;
 			}
 
+			var hoisted = hoistAwaitAndYield(path, arg0);
+
 			path.node.arguments[0] = wrapWithHelper({
 				HELPER_ID: t.identifier(this.avaThrowHelper()),
-				EXP: arg0,
+				EXP: hoisted ? hoisted.arg0 : arg0,
 				LINE: t.numericLiteral(arg0.loc.start.line),
 				COLUMN: t.numericLiteral(arg0.loc.start.column),
 				SOURCE: t.stringLiteral(state.file.code.substring(arg0.start, arg0.end)),
 				FILE: t.stringLiteral(state.file.opts.filename)
 			}).expression;
+
+			if (hoisted) {
+				path.replaceWithMultiple(
+					hoisted.expressions.concat(
+						t.expressionStatement(path.node)
+					)
+				);
+			}
 		}
 	}
 };
@@ -83,3 +95,32 @@ function isThrowsMember(path) {
 			path.get('property').isIdentifier({name: 'notThrows'})
 		);
 }
+
+function hoistAwaitAndYield(path, arg0) {
+	if (t.isAwaitExpression(arg0) || t.isYieldExpression(arg0)) {
+		var arg = path.scope.generateUidIdentifier('arg');
+		return {
+			expressions: [hoistExpression({ARG: arg, EXP: arg0})],
+			arg0: arg
+		};
+	}
+
+	if (t.isCallExpression(arg0)) {
+		var expressions = [];
+		arg0.arguments = arg0.arguments.map(function (arg) {
+			var hoisted = hoistAwaitAndYield(path, arg);
+			if (!hoisted) {
+				return arg;
+			}
+
+			expressions = expressions.concat(hoisted.expressions);
+			return hoisted.arg0;
+		});
+		return {
+			expressions: expressions,
+			arg0: arg0
+		};
+	}
+
+	return null;
+}

test.js

@@ -51,24 +51,28 @@ const HELPER = `function _avaThrowsHelper(fn, data) {
   }
 }\n`;
 
-function wrapped(throws, expression, line, column) {
+function wrapped(throws, line, column, expression, source = expression) { // eslint-disable-line max-params
 	return `t.${throws}(_avaThrowsHelper(function () {
   return ${expression};
 }, {
   line: ${line},
   column: ${column},
-  source: "${expression}",
+  source: "${source}",
   filename: "some-file.js"
 }));`;
 }
 
+function indent(str) {
+	return str.replace(/\n/g, '\n  ').trimRight();
+}
+
 test('creates a helper', t => {
 	const input = 't.throws(foo())';
 	const {code} = transform(input);
 
 	const expected = [
 		HELPER,
-		wrapped('throws', 'foo()', 1, 9)
+		wrapped('throws', 1, 9, 'foo()')
 	].join('\n');
 
 	t.is(code, expected);
@@ -81,14 +85,62 @@ test('creates the helper only once', t => {
 
 	const expected = [
 		HELPER,
-		wrapped('throws', 'foo()', 1, 9),
-		wrapped('throws', 'bar()', 2, 9)
+		wrapped('throws', 1, 9, 'foo()'),
+		wrapped('throws', 2, 9, 'bar()')
 	].join('\n');
 
 	t.is(code, expected);
 	addExample(input, code);
 });
 
+test('hoists await expressions', t => {
+	const input = `async function test() {
+  t.throws(foo(await bar(), await baz(), qux));
+  t.throws(await quux);
+}`;
+	const {code} = transform(input);
+
+	const expected = `${HELPER}
+async function test() {
+  var _arg = await bar();
+
+  var _arg2 = await baz();
+
+  ${indent(wrapped('throws', 2, 11, 'foo(_arg, _arg2, qux)', 'foo(await bar(), await baz(), qux)'))}
+
+  var _arg3 = await quux;
+
+  ${indent(wrapped('throws', 3, 11, '_arg3', 'await quux'))}
+}`;
+
+	t.is(code, expected);
+	addExample(input, code);
+});
+
+test('hoists yield expressions', t => {
+	const input = `function* test() {
+  t.throws(foo(yield bar(), yield baz(), qux));
+  t.throws(yield quux);
+}`;
+	const {code} = transform(input);
+
+	const expected = `${HELPER}
+function* test() {
+  var _arg = yield bar();
+
+  var _arg2 = yield baz();
+
+  ${indent(wrapped('throws', 2, 11, 'foo(_arg, _arg2, qux)', 'foo(yield bar(), yield baz(), qux)'))}
+
+  var _arg3 = yield quux;
+
+  ${indent(wrapped('throws', 3, 11, '_arg3', 'yield quux'))}
+}`;
+
+	t.is(code, expected);
+	addExample(input, code);
+});
+
 test('does nothing if it does not match', t => {
 	const input = 't.is(foo());';
 	const {code} = transform(input);
@@ -103,7 +155,7 @@ test('helps notThrows', t => {
 
 	const expected = [
 		HELPER,
-		wrapped('notThrows', 'baz()', 1, 12)
+		wrapped('notThrows', 1, 12, 'baz()')
 	].join('\n');
 
 	t.is(code, expected);