Support yield in do expression (#10101)

Co-authored-by: Huáng Jùnliàng <jlhwung@gmail.com>

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/await-yield/exec.js

@@ -0,0 +1,15 @@
+async function* asyncGenerator(x) {
+  const y = do {
+    let z;
+    yield 3;
+    yield await x;
+  };
+
+  return y;
+}
+
+const promise = Promise.resolve(5);
+const gen = asyncGenerator(promise);
+expect(gen.next()).resolves.toMatchObject({ value: 3, done: false });
+expect(gen.next()).resolves.toMatchObject({ value: 5, done: false });
+expect(gen.next(10)).resolves.toMatchObject({ value: 10, done: true });

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/await-yield/input.js

@@ -0,0 +1,8 @@
+async function* asyncGenerator(x) {
+  const y = do {
+    let z;
+    yield await x;
+  };
+
+  return y;
+}

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/await-yield/options.json

@@ -0,0 +1,3 @@
+{
+  "minNodeVersion": "10.0.0"
+}

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/await-yield/output.js

@@ -0,0 +1,7 @@
+async function* asyncGenerator(x) {
+  const y = yield* async function* () {
+    let z;
+    return yield await x;
+  }();
+  return y;
+}

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/yield/exec.js

@@ -0,0 +1,14 @@
+function * generator() {
+  yield 1;
+  const y = do {
+    let z;
+    yield 2;
+  };
+
+  return y;
+}
+
+const gen = generator();
+expect(gen.next().value).toBe(1);
+expect(gen.next().value).toBe(2);
+expect(gen.next(3).value).toBe(3);

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/yield/input.js

@@ -0,0 +1,8 @@
+function * g() {
+  const y = do {
+    let z;
+    yield 1;
+  };
+
+  return y;
+}

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/yield/options.json

@@ -0,0 +1,3 @@
+{
+  "minNodeVersion": "8.0.0"
+}

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/yield/output.js

@@ -0,0 +1,7 @@
+function* g() {
+  const y = yield* function* () {
+    let z;
+    return yield 1;
+  }();
+  return y;
+}

packages/babel-traverse/src/path/replacement.ts

@@ -228,6 +228,7 @@ export function replaceExpressionWithStatements(
 
   const functionParent = this.getFunctionParent();
   const isParentAsync = functionParent?.is("async");
+  const isParentGenerator = functionParent?.is("generator");
 
   const container = t.arrowFunctionExpression([], t.blockStatement(nodes));
 
@@ -277,16 +278,30 @@ export function replaceExpressionWithStatements(
   callee.arrowFunctionToExpression();
 
   // (() => await xxx)() -> await (async () => await xxx)();
-  if (
+  const needToAwaitFunction =
     isParentAsync &&
     traverse.hasType(
       (this.get("callee.body") as NodePath<t.BlockStatement>).node,
       "AwaitExpression",
       t.FUNCTION_TYPES,
-    )
+    );
-  ) {
+  const needToYieldFunction =
+    isParentGenerator &&
+    traverse.hasType(
+      (this.get("callee.body") as NodePath<t.BlockStatement>).node,
+      "YieldExpression",
+      t.FUNCTION_TYPES,
+    );
+  if (needToAwaitFunction) {
     callee.set("async", true);
-    this.replaceWith(t.awaitExpression((this as ThisType).node));
+    // yield* will await the generator return result
+    if (!needToYieldFunction) {
+      this.replaceWith(t.awaitExpression((this as ThisType).node));
+    }
+  }
+  if (needToYieldFunction) {
+    callee.set("generator", true);
+    this.replaceWith(t.yieldExpression((this as ThisType).node, true));
   }
 
   return callee.get("body.body");