improve

babel · Nov 19, 2023 · bf5b052 · bf5b052
1 parent e557b4d
commit bf5b052
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 86 deletions.
diff --git a/packages/babel-helper-remap-async-to-generator/src/index.ts b/packages/babel-helper-remap-async-to-generator/src/index.ts
@@ -107,4 +107,4 @@ export default function (
  }
 }
 
-export { onCallExpressionExit } from "@babel/helper-wrap-function";
+export { buildOnCallExpression } from "@babel/helper-wrap-function";
diff --git a/packages/babel-helper-wrap-function/src/index.ts b/packages/babel-helper-wrap-function/src/index.ts
@@ -17,6 +17,7 @@ import {
  logicalExpression,
 } from "@babel/types";
 import type * as t from "@babel/types";
+import type { PluginPass } from "@babel/core";
 
 type ExpressionWrapperBuilder<ExtraBody extends t.Node[]> = (
  replacements?: Parameters<ReturnType<typeof template.expression>>[0],
@@ -316,36 +317,45 @@ export default function wrapFunction(
  }
 }
 
-export function onCallExpressionExit(path: NodePath<t.CallExpression>) {
- if (path.parentPath.isCallExpression()) {
- const wrappedFn = path.parentPath.getData(
- "babel-helper-wrap-function_wrapped_function",
- );
-
- if (!wrappedFn || wrappedFn === path.node.callee) return;
-
- const fnPath = path.findParent(p => p.isFunction());
-
- if (
- fnPath
- .findParent(p => p.isLoop() || p.isFunction() || p.isClass())
- ?.isLoop() !== true
- ) {
- const ref = path.scope.generateUidIdentifier("ref");
- fnPath.parentPath.scope.push({
- id: ref,
- });
- const oldNode = path.node;
- const comments = path.node.leadingComments;
- if (comments) path.node.leadingComments = undefined;
- path.replaceWith(
- assignmentExpression(
- "=",
- cloneNode(ref),
- logicalExpression("||", cloneNode(ref), oldNode),
- ),
- );
- if (comments) oldNode.leadingComments = comments;
- }
- }
+export function buildOnCallExpression(helperName: string) {
+ return {
+ CallExpression: {
+ exit(path: NodePath<t.CallExpression>, state: PluginPass) {
+ if (!state.availableHelper(helperName)) {
+ return;
+ }
+ if (path.parentPath.isCallExpression()) {
+ const wrappedFn = path.parentPath.getData(
+ "babel-helper-wrap-function_wrapped_function",
+ );
+
+ if (!wrappedFn || wrappedFn === path.node.callee) return;
+
+ const fnPath = path.findParent(p => p.isFunction());
+
+ if (
+ fnPath
+ .findParent(p => p.isLoop() || p.isFunction() || p.isClass())
+ ?.isLoop() !== true
+ ) {
+ const ref = path.scope.generateUidIdentifier("ref");
+ fnPath.parentPath.scope.push({
+ id: ref,
+ });
+ const oldNode = path.node;
+ const comments = path.node.leadingComments;
+ if (comments) path.node.leadingComments = undefined;
+ path.replaceWith(
+ assignmentExpression(
+ "=",
+ cloneNode(ref),
+ logicalExpression("||", cloneNode(ref), oldNode),
+ ),
+ );
+ if (comments) oldNode.leadingComments = comments;
+ }
+ }
+ },
+ },
+ };
 }
diff --git a/packages/babel-plugin-proposal-function-sent/src/index.ts b/packages/babel-plugin-proposal-function-sent/src/index.ts
@@ -47,11 +47,7 @@ export default declare(api => {
  inherits: syntaxFunctionSent,
 
  visitor: {
- CallExpression(path, state) {
- if (state.availableHelper("callSkipFirstGeneratorNext")) {
- wrapFunction.onCallExpressionExit(path);
- }
- },
+ ...wrapFunction.buildOnCallExpression("callSkipFirstGeneratorNext"),
  MetaProperty(path, state) {
  if (!isFunctionSent(path.node)) return;
 

diff --git a/packages/babel-plugin-transform-async-generator-functions/src/index.ts b/packages/babel-plugin-transform-async-generator-functions/src/index.ts
@@ -89,18 +89,19 @@ export default declare(api => {
 
  // We don't need to pass the noNewArrows assumption, since
  // async generators are never arrow functions.
- if (state.availableHelper("callAsyncGenerator")) {
- remapAsyncToGenerator.default(path, {
- wrapAsync: "wrapAsyncGenerator",
- wrapAwait: "awaitAsyncGenerator",
- callAsync: "callAsyncGenerator",
- });
- } else {
- remapAsyncToGenerator.default(path, {
- wrapAsync: state.addHelper("wrapAsyncGenerator"),
- wrapAwait: state.addHelper("awaitAsyncGenerator"),
- });
- }
+ remapAsyncToGenerator.default(
+ path,
+ state.availableHelper("callAsyncGenerator")
+ ? {
+ wrapAsync: "wrapAsyncGenerator",
+ wrapAwait: "awaitAsyncGenerator",
+ callAsync: "callAsyncGenerator",
+ }
+ : {
+ wrapAsync: state.addHelper("wrapAsyncGenerator"),
+ wrapAwait: state.addHelper("awaitAsyncGenerator"),
+ },
+ );
  },
  };
 
@@ -114,13 +115,7 @@ export default declare(api => {
  require("@babel/plugin-syntax-async-generators").default,
 
  visitor: {
- CallExpression: {
- exit(path, state) {
- if (state.availableHelper("callAsyncGenerator")) {
- remapAsyncToGenerator.onCallExpressionExit(path);
- }
- },
- },
+ ...remapAsyncToGenerator.buildOnCallExpression("callAsyncGenerator"),
  Program(path, state) {
  // We need to traverse the ast here (instead of just vising Function
  // in the top level visitor) because for-await needs to run before the

diff --git a/packages/babel-plugin-transform-async-to-generator/src/index.ts b/packages/babel-plugin-transform-async-to-generator/src/index.ts
@@ -54,39 +54,24 @@ export default declare<State>((api, options: Options) => {
  name: "transform-async-to-generator",
 
  visitor: {
- CallExpression: {
- exit(path, state) {
- if (state.availableHelper("callAsync")) {
- remapAsyncToGenerator.onCallExpressionExit(path);
- }
- },
- },
+ ...remapAsyncToGenerator.buildOnCallExpression("callAsync"),
  Function(path, state) {
  if (!path.node.async || path.node.generator) return;
 
- if (
+ remapAsyncToGenerator.default(
+ path,
  state.availableHelper("callAsync") &&
- state.availableHelper("asyncToGenerator2")
- ) {
- remapAsyncToGenerator.default(
- path,
- {
- wrapAsync: "asyncToGenerator2",
- callAsync: "callAsync",
- },
- noNewArrows,
- ignoreFunctionLength,
- );
- } else {
- remapAsyncToGenerator.default(
- path,
- {
- wrapAsync: state.addHelper("asyncToGenerator"),
- },
- noNewArrows,
- ignoreFunctionLength,
- );
- }
+ state.availableHelper("asyncToGenerator2")
+ ? {
+ wrapAsync: "asyncToGenerator2",
+ callAsync: "callAsync",
+ }
+ : {
+ wrapAsync: state.addHelper("asyncToGenerator"),
+ },
+ noNewArrows,
+ ignoreFunctionLength,
+ );
  },
  },
  };