transform.cjs 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  1. 'use strict';
  2. const acorn = require('acorn');
  3. const MagicString = require('magic-string');
  4. const estreeWalker = require('estree-walker');
  5. function _interopDefaultCompat (e) { return e && typeof e === 'object' && 'default' in e ? e.default : e; }
  6. function _interopNamespaceCompat(e) {
  7. if (e && typeof e === 'object' && 'default' in e) return e;
  8. const n = Object.create(null);
  9. if (e) {
  10. for (const k in e) {
  11. n[k] = e[k];
  12. }
  13. }
  14. n.default = e;
  15. return n;
  16. }
  17. const acorn__namespace = /*#__PURE__*/_interopNamespaceCompat(acorn);
  18. const MagicString__default = /*#__PURE__*/_interopDefaultCompat(MagicString);
  19. const kInjected = "__unctx_injected__";
  20. function createTransformer(options = {}) {
  21. options = {
  22. asyncFunctions: ["withAsyncContext"],
  23. helperModule: "unctx",
  24. helperName: "executeAsync",
  25. objectDefinitions: {},
  26. ...options
  27. };
  28. const objectDefinitionFunctions = Object.keys(options.objectDefinitions);
  29. const matchRE = new RegExp(
  30. `\\b(${[...options.asyncFunctions, ...objectDefinitionFunctions].join(
  31. "|"
  32. )})\\(`
  33. );
  34. function shouldTransform(code) {
  35. return typeof code === "string" && matchRE.test(code);
  36. }
  37. function transform(code, options_ = {}) {
  38. if (!options_.force && !shouldTransform(code)) {
  39. return;
  40. }
  41. const ast = acorn__namespace.parse(code, {
  42. sourceType: "module",
  43. ecmaVersion: "latest",
  44. locations: true
  45. });
  46. const s = new MagicString__default(code);
  47. const lines = code.split("\n");
  48. let detected = false;
  49. estreeWalker.walk(ast, {
  50. enter(node) {
  51. if (node.type === "CallExpression") {
  52. const functionName = _getFunctionName(node.callee);
  53. if (options.asyncFunctions.includes(functionName)) {
  54. transformFunctionArguments(node);
  55. if (functionName !== "callAsync") {
  56. const lastArgument = node.arguments[node.arguments.length - 1];
  57. if (lastArgument && lastArgument.loc) {
  58. s.appendRight(toIndex(lastArgument.loc.end), ",1");
  59. }
  60. }
  61. }
  62. if (objectDefinitionFunctions.includes(functionName)) {
  63. for (const argument of node.arguments) {
  64. if (argument.type !== "ObjectExpression") {
  65. continue;
  66. }
  67. for (const property of argument.properties) {
  68. if (property.type !== "Property" || property.key.type !== "Identifier") {
  69. continue;
  70. }
  71. if (options.objectDefinitions[functionName]?.includes(
  72. property.key?.name
  73. )) {
  74. transformFunctionBody(property.value);
  75. }
  76. }
  77. }
  78. }
  79. }
  80. }
  81. });
  82. if (!detected) {
  83. return;
  84. }
  85. s.appendLeft(
  86. 0,
  87. `import { ${options.helperName} as __executeAsync } from "${options.helperModule}";`
  88. );
  89. return {
  90. code: s.toString(),
  91. magicString: s
  92. };
  93. function toIndex(pos) {
  94. return lines.slice(0, pos.line - 1).join("\n").length + pos.column + 1;
  95. }
  96. function transformFunctionBody(function_) {
  97. if (function_.type !== "ArrowFunctionExpression" && function_.type !== "FunctionExpression") {
  98. return;
  99. }
  100. if (!function_.async) {
  101. return;
  102. }
  103. const body = function_.body;
  104. let injectVariable = false;
  105. estreeWalker.walk(body, {
  106. enter(node, parent) {
  107. if (node.type === "AwaitExpression" && !node[kInjected]) {
  108. detected = true;
  109. injectVariable = true;
  110. injectForNode(node, parent);
  111. } else if (node.type === "IfStatement" && node.consequent.type === "ExpressionStatement" && node.consequent.expression.type === "AwaitExpression") {
  112. detected = true;
  113. injectVariable = true;
  114. node.consequent.expression[kInjected] = true;
  115. injectForNode(node.consequent.expression, node);
  116. }
  117. if (node.type === "ArrowFunctionExpression" || node.type === "FunctionExpression" || node.type === "FunctionDeclaration") {
  118. return this.skip();
  119. }
  120. }
  121. });
  122. if (injectVariable && body.loc) {
  123. s.appendLeft(toIndex(body.loc.start) + 1, "let __temp, __restore;");
  124. }
  125. }
  126. function transformFunctionArguments(node) {
  127. for (const function_ of node.arguments) {
  128. transformFunctionBody(function_);
  129. }
  130. }
  131. function injectForNode(node, parent) {
  132. const isStatement = parent?.type === "ExpressionStatement";
  133. if (!node.loc || !node.argument.loc) {
  134. return;
  135. }
  136. s.remove(toIndex(node.loc.start), toIndex(node.argument.loc.start));
  137. s.remove(toIndex(node.loc.end), toIndex(node.argument.loc.end));
  138. s.appendLeft(
  139. toIndex(node.argument.loc.start),
  140. isStatement ? `;(([__temp,__restore]=__executeAsync(()=>` : `(([__temp,__restore]=__executeAsync(()=>`
  141. );
  142. s.appendRight(
  143. toIndex(node.argument.loc.end),
  144. isStatement ? `)),await __temp,__restore());` : `)),__temp=await __temp,__restore(),__temp)`
  145. );
  146. }
  147. }
  148. return {
  149. transform,
  150. shouldTransform
  151. };
  152. }
  153. function _getFunctionName(node) {
  154. if (node.type === "Identifier") {
  155. return node.name;
  156. } else if (node.type === "MemberExpression") {
  157. return _getFunctionName(node.property);
  158. }
  159. return "";
  160. }
  161. exports.createTransformer = createTransformer;