fix: prevent shadowing and bloat in generator expressions

* prevent shadowing of innocuous identifiers kList, kHead, kTail * reduce bloat by reusing iterator functions defined once * rewrite head expansion to avoid yield* singleElementArray
MaxMotovilov · Jul 12, 2021 · 954af12 · 954af12
1 parent 2977196
commit 954af12
Showing 1 changed file with 104 additions and 39 deletions.
diff --git a/packages/babel-plugin-proposal-generator-literals/src/index.js b/packages/babel-plugin-proposal-generator-literals/src/index.js
@@ -3,52 +3,103 @@ import syntaxGeneratorLiterals from "@babel/plugin-syntax-generator-literals";
 import { types as t } from "@babel/core";
 import template from "@babel/template";
 
-const generatorEmpty = template.expression(`
+const symbolFor = template.expression(`Symbol.for(%%STRING%%)`);
+
+const generatorObjectEmpty = 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 generatorObjectNoSpread = template.expression(`
+{
+  [%%ID_LIST%%]: %%LIST_ARRAY%%,
+  [Symbol.iterator]: %%ID_ITERATOR%%,
+}
 `);
 
-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;
+const iteratorFunctionNoSpread = template.expression(`
+function* () {
+  let result;
+  for (const getter of this[%%ID_LIST%%]) {
+    result = yield getter();
+  }
+  return result;
+}
+`);
+
+const generatorObjectWithSpread = template.expression(`
+{
+  [%%ID_HEAD%%]: %%HEAD_ARRAY%%,
+  [%%ID_TAIL%%]: %%TAIL_FUNC%%,
+  [Symbol.iterator]: %%ID_ITERATOR%%,
+}
+`);
+
+const iteratorFunctionWithSpread = template.expression(`
+function* () {
+  let iter = this;
+  do {
+    const head = iter[%%ID_HEAD%%];
+    for (let i = 0, len = head.length; i < len; ) {
+      const single = head[i++];
+      if (single !== null) {
+        yield single();
+      } else {
+        const spread = head[i++];
+        yield* spread();
+      }
     }
-  })
-)(Symbol.for('##generatorExpressionHead'), Symbol.for('##generatorExpressionTail'))
+    iter = iter[%%ID_TAIL%%]();
+  } while (%%ID_HEAD%% in iter);
+  return yield* iter;
+}
 `);
 
-function makeYieldGetter(expr) {
-  return t.arrowFunctionExpression([], expr);
+function getHelper(scope, what, builder) {
+  const key = "genexpr:" + what;
+  let id = scope.getData(key);
+  if (!id) {
+    id = scope.generateUidIdentifier(key);
+    scope.push({
+      id,
+      kind: "const",
+      init: builder(),
+    });
+    scope.setData(key, id);
+  }
+  return t.cloneNode(id);
+}
+
+function getSymbolFor(scope, what) {
+  return getHelper(scope, what, () =>
+    symbolFor({
+      STRING: t.stringLiteral("##generatorExpression" + what),
+    }),
+  );
 }
 
-function makeYieldIterable(expr) {
-  return t.isSpreadElement(expr) ? expr.argument : t.arrayExpression([expr]);
+function getIteratorNoSpread(scope) {
+  return getHelper(scope, "IteratorNoSpread", () =>
+    iteratorFunctionNoSpread({
+      ID_LIST: getSymbolFor(scope, "List"),
+    }),
+  );
+}
+
+function getIteratorWithSpread(scope) {
+  return getHelper(scope, "IteratorWithSpread", () =>
+    iteratorFunctionWithSpread({
+      ID_HEAD: getSymbolFor(scope, "Head"),
+      ID_TAIL: getSymbolFor(scope, "Tail"),
+    }),
+  );
+}
+
+function makeYieldGetter(expr) {
+  return t.isSpreadElement(expr)
+    ? [t.nullLiteral(), t.arrowFunctionExpression([], expr.argument)]
+    : t.arrowFunctionExpression([], expr);
 }
 
 export default declare(api => {
@@ -58,23 +109,37 @@ export default declare(api => {
     name: "proposal-generator-literals",
     inherits: syntaxGeneratorLiterals,
     visitor: {
-      GeneratorExpression(path) {
+      GeneratorExpression(path, state) {
         const { elements } = path.node;
         if (elements.length === 0) {
-          path.replaceWith(generatorEmpty());
+          path.replaceWith(generatorObjectEmpty());
         } else if (elements.some(expr => t.isSpreadElement(expr))) {
-          const getters = elements.map(makeYieldIterable).map(makeYieldGetter);
+          const getters = elements.flatMap(makeYieldGetter);
           const tail = getters.pop();
+          const last = elements[elements.length - 1];
+          if (t.isSpreadElement(last)) {
+            t.assertNullLiteral(getters.pop());
+          } else {
+            t.assertArrowFunctionExpression(tail);
+            tail.body = t.arrayExpression([tail.body]);
+          }
           path.replaceWith(
-            generatorWithTail({
+            generatorObjectWithSpread({
+              ID_HEAD: getSymbolFor(state.file.scope, "Head"),
+              ID_TAIL: getSymbolFor(state.file.scope, "Tail"),
+              ID_ITERATOR: getIteratorWithSpread(state.file.scope),
               HEAD_ARRAY: t.arrayExpression(getters),
               TAIL_FUNC: tail,
             }),
           );
         } else {
           const getters = elements.map(makeYieldGetter);
           path.replaceWith(
-            generatorNoSpread({ LIST: t.arrayExpression(getters) }),
+            generatorObjectNoSpread({
+              ID_LIST: getSymbolFor(state.file.scope, "List"),
+              ID_ITERATOR: getIteratorNoSpread(state.file.scope),
+              LIST_ARRAY: t.arrayExpression(getters),
+            }),
           );
         }
       },