From 1581424d18b5a7f3f151862ceeaab9d5e3fd63b0 Mon Sep 17 00:00:00 2001
From: Brian Ng <bng412@gmail.com>
Date: Thu, 9 Feb 2017 15:06:41 -0600
Subject: [PATCH] Prevent multiple return statements in a loop when replacing
 expressions (#5030)

---
 .../fixtures/do-expressions/while-if/actual.js     | 10 ++++++++++
 .../fixtures/do-expressions/while-if/expected.js   | 13 +++++++++++++
 packages/babel-traverse/src/path/replacement.js    | 14 ++++++++++----
 3 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/actual.js
 create mode 100644 packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/expected.js

diff --git a/packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/actual.js b/packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/actual.js
new file mode 100644
index 000000000000..d3706c14b3dc
--- /dev/null
+++ b/packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/actual.js
@@ -0,0 +1,10 @@
+let p
+let a = do {
+  while (p = p.parentPath) {
+    if (a) {
+      'a'
+    } else {
+      'b'
+    }
+  }
+};
diff --git a/packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/expected.js b/packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/expected.js
new file mode 100644
index 000000000000..5c0298f045b3
--- /dev/null
+++ b/packages/babel-plugin-transform-do-expressions/test/fixtures/do-expressions/while-if/expected.js
@@ -0,0 +1,13 @@
+let p;
+let a = function () {
+  var _ret;
+
+  while (p = p.parentPath) {
+    if (a) {
+      _ret = 'a';
+    } else {
+      _ret = 'b';
+    }
+  }
+  return _ret;
+}();
diff --git a/packages/babel-traverse/src/path/replacement.js b/packages/babel-traverse/src/path/replacement.js
index 1a8cdbef5495..47e706d0d2bb 100644
--- a/packages/babel-traverse/src/path/replacement.js
+++ b/packages/babel-traverse/src/path/replacement.js
@@ -219,10 +219,16 @@ export function replaceExpressionWithStatements(nodes: Array<Object>) {
 
       const loop = path.findParent((path) => path.isLoop());
       if (loop) {
-        const callee = this.get("callee");
-
-        const uid = callee.scope.generateDeclaredUidIdentifier("ret");
-        callee.get("body").pushContainer("body", t.returnStatement(uid));
+        let uid = loop.getData("expressionReplacementReturnUid");
+
+        if (!uid) {
+          const callee = this.get("callee");
+          uid = callee.scope.generateDeclaredUidIdentifier("ret");
+          callee.get("body").pushContainer("body", t.returnStatement(uid));
+          loop.setData("expressionReplacementReturnUid", uid);
+        } else {
+          uid = t.identifier(uid.name);
+        }
 
         path.get("expression").replaceWith(
           t.assignmentExpression("=", uid, path.node.expression)