From 4e8435cbb2f8269e3f7b1f04869faf01fba57684 Mon Sep 17 00:00:00 2001
From: Lukas Taegert <lukas.taegert@tngtech.com>
Date: Wed, 16 May 2018 15:14:40 +0200
Subject: [PATCH] Resolve #2191: Add spaces after return and yield if necessary

---
 src/ast/nodes/ReturnStatement.ts                  | 15 ++++++++++++++-
 src/ast/nodes/YieldExpression.ts                  | 15 ++++++++++++++-
 .../return-statement/missing-space/_config.js     |  3 +++
 .../return-statement/missing-space/_expected.js   |  5 +++++
 .../return-statement/missing-space/main.js        |  5 +++++
 .../yield-expression/missing-space/_config.js     |  3 +++
 .../yield-expression/missing-space/_expected.js   |  7 +++++++
 .../yield-expression/missing-space/main.js        |  7 +++++++
 8 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 test/form/samples/return-statement/missing-space/_config.js
 create mode 100644 test/form/samples/return-statement/missing-space/_expected.js
 create mode 100644 test/form/samples/return-statement/missing-space/main.js
 create mode 100644 test/form/samples/yield-expression/missing-space/_config.js
 create mode 100644 test/form/samples/yield-expression/missing-space/_expected.js
 create mode 100644 test/form/samples/yield-expression/missing-space/main.js

diff --git a/src/ast/nodes/ReturnStatement.ts b/src/ast/nodes/ReturnStatement.ts
index 4312ed0b540..bb00cc9622f 100644
--- a/src/ast/nodes/ReturnStatement.ts
+++ b/src/ast/nodes/ReturnStatement.ts
@@ -2,17 +2,30 @@ import { UNKNOWN_EXPRESSION } from '../values';
 import { ExecutionPathOptions } from '../ExecutionPathOptions';
 import * as NodeType from './NodeType';
 import { ExpressionNode, StatementBase } from './shared/Node';
+import { RenderOptions } from '../../utils/renderHelpers';
+import MagicString from 'magic-string';
 
 export default class ReturnStatement extends StatementBase {
 	type: NodeType.tReturnStatement;
 	argument: ExpressionNode | null;
 
 	hasEffects(options: ExecutionPathOptions) {
-		return super.hasEffects(options) || !options.ignoreReturnAwaitYield();
+		return (
+			!options.ignoreReturnAwaitYield() || (this.argument && this.argument.hasEffects(options))
+		);
 	}
 
 	initialise() {
 		this.included = false;
 		this.scope.addReturnExpression(this.argument || UNKNOWN_EXPRESSION);
 	}
+
+	render(code: MagicString, options: RenderOptions) {
+		if (this.argument) {
+			this.argument.render(code, options);
+			if (this.argument.start === this.start + 6 /* 'return'.length */) {
+				code.prependLeft(this.start + 6, ' ');
+			}
+		}
+	}
 }
diff --git a/src/ast/nodes/YieldExpression.ts b/src/ast/nodes/YieldExpression.ts
index a5ae93f174b..74b4ffb9f7c 100644
--- a/src/ast/nodes/YieldExpression.ts
+++ b/src/ast/nodes/YieldExpression.ts
@@ -1,6 +1,8 @@
 import { ExecutionPathOptions } from '../ExecutionPathOptions';
 import * as NodeType from './NodeType';
 import { ExpressionNode, NodeBase } from './shared/Node';
+import { RenderOptions } from '../../utils/renderHelpers';
+import MagicString from 'magic-string';
 
 export default class YieldExpression extends NodeBase {
 	type: NodeType.tYieldExpression;
@@ -8,6 +10,17 @@ export default class YieldExpression extends NodeBase {
 	delegate: boolean;
 
 	hasEffects(options: ExecutionPathOptions) {
-		return super.hasEffects(options) || !options.ignoreReturnAwaitYield();
+		return (
+			!options.ignoreReturnAwaitYield() || (this.argument && this.argument.hasEffects(options))
+		);
+	}
+
+	render(code: MagicString, options: RenderOptions) {
+		if (this.argument) {
+			this.argument.render(code, options);
+			if (this.argument.start === this.start + 5 /* 'yield'.length */) {
+				code.prependLeft(this.start + 5, ' ');
+			}
+		}
 	}
 }
diff --git a/test/form/samples/return-statement/missing-space/_config.js b/test/form/samples/return-statement/missing-space/_config.js
new file mode 100644
index 00000000000..ab1265aacfe
--- /dev/null
+++ b/test/form/samples/return-statement/missing-space/_config.js
@@ -0,0 +1,3 @@
+module.exports = {
+	description: 'Inserts space when simplifying return statement without space'
+};
diff --git a/test/form/samples/return-statement/missing-space/_expected.js b/test/form/samples/return-statement/missing-space/_expected.js
new file mode 100644
index 00000000000..e5a7078164c
--- /dev/null
+++ b/test/form/samples/return-statement/missing-space/_expected.js
@@ -0,0 +1,5 @@
+function test() {
+	return null;
+}
+
+console.log(test());
diff --git a/test/form/samples/return-statement/missing-space/main.js b/test/form/samples/return-statement/missing-space/main.js
new file mode 100644
index 00000000000..9183a564360
--- /dev/null
+++ b/test/form/samples/return-statement/missing-space/main.js
@@ -0,0 +1,5 @@
+function test() {
+	return!1||null;
+}
+
+console.log(test());
diff --git a/test/form/samples/yield-expression/missing-space/_config.js b/test/form/samples/yield-expression/missing-space/_config.js
new file mode 100644
index 00000000000..3dace68f84f
--- /dev/null
+++ b/test/form/samples/yield-expression/missing-space/_config.js
@@ -0,0 +1,3 @@
+module.exports = {
+	description: 'Inserts space when simplifying yield expression without space'
+};
diff --git a/test/form/samples/yield-expression/missing-space/_expected.js b/test/form/samples/yield-expression/missing-space/_expected.js
new file mode 100644
index 00000000000..5393bdee8a4
--- /dev/null
+++ b/test/form/samples/yield-expression/missing-space/_expected.js
@@ -0,0 +1,7 @@
+function* test() {
+	yield null;
+}
+
+for (const x of test()) {
+	console.log(x);
+}
diff --git a/test/form/samples/yield-expression/missing-space/main.js b/test/form/samples/yield-expression/missing-space/main.js
new file mode 100644
index 00000000000..556e6b2e008
--- /dev/null
+++ b/test/form/samples/yield-expression/missing-space/main.js
@@ -0,0 +1,7 @@
+function* test() {
+	yield!1||null;
+}
+
+for (const x of test()) {
+	console.log(x);
+}