transform.mjs 4.8 KB

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