diff --git a/src/ifrefactor.js b/src/ifrefactor.js
index 5095c43..62ab3bb 100644
--- a/src/ifrefactor.js
+++ b/src/ifrefactor.js
@@ -84,25 +84,28 @@ export const FirstPassIfVisitor = {
 const containsReturnOrAwait = matcher(['ReturnStatement', 'AwaitExpression'], NoSubFunctionsVisitor);
 
 export const SecondPassIfVisitor = extend({
-  IfStatement(path) {
-    const alt = path.node.alternate;
-    if (!path.node.consequent.body.length && alt && alt.body.length) {
-      path.node.consequent = path.node.alternate;
-      path.node.alternate = null;
-      path.node.test = unaryExpression('!', path.node.test);
-    }
-    const ifContainsAwait = containsAwait(path.get('consequent'));
-    const elseContainsAwait = containsAwait(path.get('alternate'));
+  IfStatement: {
+    exit(path) {
+      const alt = path.node.alternate;
+      if (!path.node.consequent.body.length && alt && alt.body.length) {
+        path.node.consequent = path.node.alternate;
+        path.node.alternate = null;
+        path.node.test = unaryExpression('!', path.node.test);
+      }
+      const ifContainsAwait = containsAwait(path.get('consequent'));
+      const elseContainsAwait = containsAwait(path.get('alternate'));
 
-    const {node} = path;
-    if (ifContainsAwait) {
-      node.consequent = wrapIfBranch(node.consequent);
-    }
-    if (elseContainsAwait) {
-      node.alternate = wrapIfBranch(node.alternate);
-    }
-    if (ifContainsAwait || elseContainsAwait) {
-      path.replaceWith(awaitExpression(wrapFunction(blockStatement([node]))));
+      const {node} = path;
+      const parentDirtyAllowed = path.parentPath.parent.dirtyAllowed;
+      if (ifContainsAwait) {
+        node.consequent = wrapIfBranch(node.consequent, parentDirtyAllowed);
+      }
+      if (elseContainsAwait) {
+        node.alternate = wrapIfBranch(node.alternate, parentDirtyAllowed);
+      }
+      if (ifContainsAwait || elseContainsAwait) {
+        path.replaceWith(awaitExpression(wrapFunction(blockStatement([node]))));
+      }
     }
   },
   BlockStatement(path) {
@@ -146,12 +149,15 @@ export const SecondPassIfVisitor = extend({
         subNode.consequent.body.splice(-1);
       }
       extendElse(subNode, remainder);
+      path.parent.dirtyAllowed = true;
     }
   }
 }, NoSubFunctionsVisitor)
 
-const wrapIfBranch =
-  branch => blockStatement([returnStatement(wrapFunction(branch))]);
+function wrapIfBranch(branch, parentDirtyAllowed) {
+  const dirtyAllowed = !parentDirtyAllowed || isReturnStatement(branch.body[branch.body.length - 1]);
+  return blockStatement([returnStatement(wrapFunction(branch, dirtyAllowed))]);
+}
 
 function extendElse(ifStmt, extraBody) {
   const body = ((ifStmt.alternate || {}).body || []).concat(extraBody);
diff --git a/src/refactor.js b/src/refactor.js
index acfeeee..d8a9c85 100644
--- a/src/refactor.js
+++ b/src/refactor.js
@@ -65,8 +65,8 @@ export const RefactorVisitor = extend({
     // ->
     //
     // await Promise.all([
-    //   function () {return a();}(),
-    //   function () {return await b();}()
+    //   async function () {return a();}(),
+    //   async function () {return await b();}()
     // ])
     //
     // (which is optimized away to:)
diff --git a/src/utils.js b/src/utils.js
index 7b7773d..5a35656 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -36,9 +36,9 @@ export function matcher(types, base) {
   }
 }
 
-export function wrapFunction(body) {
+export function wrapFunction(body, dirtyAllowed) {
   const func = functionExpression(null, [], body, false, true);
-  func.dirtyAllowed = true;
+  func.dirtyAllowed = dirtyAllowed;
   return callExpression(func, []);
 }
 
diff --git a/test/fixtures/21-incorrect-return/.babelrc b/test/fixtures/21-incorrect-return/.babelrc
new file mode 100644
index 0000000..e624888
--- /dev/null
+++ b/test/fixtures/21-incorrect-return/.babelrc
@@ -0,0 +1,5 @@
+{
+  "plugins": [
+    ["../../../src"]
+  ]
+}
diff --git a/test/fixtures/21-incorrect-return/actual.js b/test/fixtures/21-incorrect-return/actual.js
new file mode 100644
index 0000000..1ab2fe8
--- /dev/null
+++ b/test/fixtures/21-incorrect-return/actual.js
@@ -0,0 +1,10 @@
+async function handler() {
+    var response = await fetch('http://address');
+    if (!response.ok) {
+        return null; // 1
+    }
+    var json = await response.json(); // 2
+    return {
+      a: 3
+    };
+}
diff --git a/test/fixtures/21-incorrect-return/expected.js b/test/fixtures/21-incorrect-return/expected.js
new file mode 100644
index 0000000..2a85daa
--- /dev/null
+++ b/test/fixtures/21-incorrect-return/expected.js
@@ -0,0 +1,22 @@
+function handler() {
+    var response, json;
+    return Promise.resolve().then(function () {
+        return fetch('http://address');
+    }).then(function (_resp) {
+        response = _resp;
+
+        if (!response.ok) {
+            return null; // 1
+        } else {
+            return Promise.resolve().then(function () {
+                return response.json();
+            }).then(function (_resp) {
+                json = _resp; // 2
+
+                return {
+                    a: 3
+                };
+            });
+        }
+    });
+}
diff --git a/test/fixtures/complex-loop/expected.js b/test/fixtures/complex-loop/expected.js
index 78f5aeb..600d739 100644
--- a/test/fixtures/complex-loop/expected.js
+++ b/test/fixtures/complex-loop/expected.js
@@ -18,18 +18,22 @@ function test() {
       if (_resp) {
         return _recursive();
       } else {
-        if (_test && i === 2) {
-          return _recursive;
-        } else {
-          if (_test) {
+        return Promise.resolve().then(function () {
+          if (_test && i === 2) {
+            return _recursive;
+          } else {
             return Promise.resolve().then(function () {
-              return db.destroy();
-            }).then(function (_resp) {
-              a = _resp;
-              return _recursive();
-            });
+              if (_test) {
+                return Promise.resolve().then(function () {
+                  return db.destroy();
+                }).then(function (_resp) {
+                  a = _resp;
+                  return _recursive();
+                });
+              }
+            }).then(function () {});
           }
-        }
+        }).then(function () {});
       }
     });
   }
diff --git a/test/fixtures/guard/expected.js b/test/fixtures/guard/expected.js
index 512dd2d..6334fca 100644
--- a/test/fixtures/guard/expected.js
+++ b/test/fixtures/guard/expected.js
@@ -1,7 +1,9 @@
 function test(a, b) {
   return Promise.resolve().then(function () {
     if (!(a === b)) {
-      return someOp((a + b) / 2);
+      return Promise.resolve().then(function () {
+        return someOp((a + b) / 2);
+      }).then(function () {});
     }
-  }).then(function () {});
+  });
 }
diff --git a/test/fixtures/if-flow/expected.js b/test/fixtures/if-flow/expected.js
index 4ef8c87..0d5b5ad 100644
--- a/test/fixtures/if-flow/expected.js
+++ b/test/fixtures/if-flow/expected.js
@@ -13,7 +13,9 @@ function test() {
         _test2 = _resp;
 
         if (_test2) {
-          return e();
+          return Promise.resolve().then(function () {
+            return e();
+          }).then(function () {});
         }
       }).then(function () {
         if (_test2 && f()) {
@@ -31,5 +33,5 @@ function test() {
         }
       });
     }
-  }).then(function () {});
+  });
 }