Merge pull request #1 from lightmare/generator-literals

flattening tail spread recursion

packages/babel-plugin-proposal-generator-literals/src/index.js

@@ -3,7 +3,53 @@ import syntaxGeneratorLiterals from "@babel/plugin-syntax-generator-literals";
 import { types as t } from "@babel/core";
 import template from "@babel/template";
 
-const literal = template.expression("{[Symbol.iterator]: function*()%%body%%}");
+const generatorEmpty = template.expression(`
+{
+  [Symbol.iterator]: function* () {}
+}
+`);
+
+const generatorNoSpread = template.expression(`
+(kList =>
+  ({
+    [kList]: %%LIST%%,
+    [Symbol.iterator]: function* () {
+      let result;
+      for (const getter of this[kList]) {
+        result = yield getter();
+      }
+      return result;
+    }
+  })
+)(Symbol.for('##generatorExpressionList'))
+`);
+
+const generatorWithTail = template.expression(`
+((kHead, kTail) =>
+  ({
+    [kHead]: %%HEAD_ARRAY%%,
+    [kTail]: %%TAIL_FUNC%%,
+    [Symbol.iterator]: function* () {
+      let iter = this;
+      do {
+        for (const getter of iter[kHead]) {
+          yield* getter();
+        }
+        iter = iter[kTail]();
+      } while (kHead in iter);
+      return yield* iter;
+    }
+  })
+)(Symbol.for('##generatorExpressionHead'), Symbol.for('##generatorExpressionTail'))
+`);
+
+function makeYieldGetter(expr) {
+  return t.arrowFunctionExpression([], expr);
+}
+
+function makeYieldIterable(expr) {
+  return t.isSpreadElement(expr) ? expr.argument : t.arrayExpression([expr]);
+}
 
 export default declare(api => {
   api.assertVersion(7);
@@ -13,18 +59,24 @@ export default declare(api => {
     inherits: syntaxGeneratorLiterals,
     visitor: {
       GeneratorExpression(path) {
-        path.replaceWith(
-          literal({
-            body: t.blockStatement(
-              path.node.elements.map(node => {
-                const isSpread = t.isSpreadElement(node);
-                return t.expressionStatement(
-                  t.yieldExpression(isSpread ? node.argument : node, isSpread),
-                );
-              }),
-            ),
-          }),
-        );
+        const { elements } = path.node;
+        if (elements.length === 0) {
+          path.replaceWith(generatorEmpty());
+        } else if (elements.some(expr => t.isSpreadElement(expr))) {
+          const getters = elements.map(makeYieldIterable).map(makeYieldGetter);
+          const tail = getters.pop();
+          path.replaceWith(
+            generatorWithTail({
+              HEAD_ARRAY: t.arrayExpression(getters),
+              TAIL_FUNC: tail,
+            }),
+          );
+        } else {
+          const getters = elements.map(makeYieldGetter);
+          path.replaceWith(
+            generatorNoSpread({ LIST: t.arrayExpression(getters) }),
+          );
+        }
       },
     },
   };

packages/babel-plugin-proposal-generator-literals/test/fixtures/infinite-multi-spread/exec.js

@@ -0,0 +1,17 @@
+const fiveZeros = *[0, 0, 0, 0, 0];
+const fiveNonZeros = *[1, 2, 3, 4, 5];
+
+let inf0, inf15;
+inf0 = *[...fiveZeros, ...inf15];
+inf15 = *[...fiveNonZeros, ...inf0];
+
+let count = 666002;
+let target = 999003;
+let sum = 0;
+
+for (const x of inf15) {
+  if (--count < 0) break;
+  sum += x;
+}
+
+expect(sum).toBe(target);

packages/babel-plugin-proposal-generator-literals/test/fixtures/infinite-sum/exec.js

@@ -0,0 +1,14 @@
+let zeroOne, oneZero;
+zeroOne = *[0, ...oneZero];
+oneZero = *[1, ...zeroOne];
+
+let count = 128000;
+let target = 64000;
+let sum = 0;
+
+for (const x of oneZero) {
+  if (--count < 0) break;
+  sum += x;
+}
+
+expect(sum).toBe(target);

packages/babel-plugin-proposal-generator-literals/test/fixtures/lexical-this/exec.js

@@ -0,0 +1,9 @@
+let ob = {
+  bar: 'A',
+  foo() {
+    return *['B', this.bar, 'R'];
+  }
+}
+
+expect(Array.from(ob.foo())).toStrictEqual(['B', 'A', 'R']);
+

packages/babel-plugin-proposal-generator-literals/test/fixtures/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": ["proposal-generator-literals"]
+}