Merge pull request #4359 from easyops-cn/steve/v3-refine-traverse

Steve/v3-refine-traverse

Author: qiaofengxi

packages/cook/src/cook.ts

@@ -301,6 +301,10 @@ export function cook(
               `Unsupported unicode flag in regular expression: ${node.raw}`
             );
           }
+          // Always create a new RegExp, because the AST will be reused.
+          return NormalCompletion(
+            new RegExp(node.regex.pattern, node.regex.flags)
+          );
         }
         return NormalCompletion(node.value);
       }

packages/cook/src/interfaces.ts

@@ -50,6 +50,7 @@ export interface EstreeLiteral {
   value: unknown;
   raw: string;
   regex?: {
+    pattern: string;
     flags: string;
   };
   loc: SourceLocation;

packages/cook/src/precookFunction.spec.ts

@@ -1,11 +1,16 @@
 import { describe, it, jest, expect } from "@jest/globals";
-import { precookFunction } from "./precookFunction.js";
+import { clearFunctionASTCache, precookFunction } from "./precookFunction.js";
+import { cook } from "./cook.js";
 
 const consoleWarn = jest
   .spyOn(console, "warn")
   .mockImplementation(() => void 0);
 
 describe("precookFunction", () => {
+  beforeEach(() => {
+    clearFunctionASTCache();
+  });
+
   it.each<[string, string, string[]]>([
     [
       "lexical variables in block statement",
@@ -206,7 +211,7 @@ describe("precookFunction", () => {
       `,
       ["c", "b", "d"],
     ],
-  ])("%s", (desc, source, result) => {
+  ])("%s", (_desc, source, result) => {
     expect(Array.from(precookFunction(source).attemptToVisitGlobals)).toEqual(
       result
     );
@@ -263,6 +268,26 @@ describe("precookFunction", () => {
     }
   );
 
+  it("should isolate regexp", () => {
+    const source = `function test() {
+      const r = /\\w/g;
+      r.exec("abc");
+      return r.lastIndex;
+    }`;
+    const fn = {
+      name: "test",
+      source,
+    };
+    const attempt1 = precookFunction(source, { cacheKey: fn });
+    const fn1 = cook(attempt1.function, source) as Function;
+    expect(fn1()).toBe(1);
+
+    const attempt2 = precookFunction(source, { cacheKey: fn });
+    const fn2 = cook(attempt2.function, source) as Function;
+    // The second RegExp is using the cached one
+    expect(fn2()).toBe(1);
+  });
+
   it("should warn unsupported type", () => {
     const { attemptToVisitGlobals } = precookFunction(
       "function test() { this }"

packages/cook/src/precookFunction.ts

@@ -1,8 +1,11 @@
-import { FunctionDeclaration } from "@babel/types";
+import type { FunctionDeclaration } from "@babel/types";
 import { parseAsEstree } from "./parse.js";
 import { precook, PrecookOptions } from "./precook.js";
 
+let ASTCache = new WeakMap<object, FunctionDeclaration>();
+
 export interface PrecookFunctionOptions extends PrecookOptions {
+  cacheKey?: object;
   typescript?: boolean;
 }
 
@@ -13,12 +16,22 @@ export interface PrecookFunctionResult {
 
 export function precookFunction(
   source: string,
-  { typescript, ...restOptions }: PrecookFunctionOptions = {}
+  { typescript, cacheKey, ...restOptions }: PrecookFunctionOptions = {}
 ): PrecookFunctionResult {
-  const func = parseAsEstree(source, { typescript });
+  let func = cacheKey ? ASTCache.get(cacheKey) : undefined;
+  if (!func) {
+    func = parseAsEstree(source, { typescript });
+    if (cacheKey) {
+      ASTCache.set(cacheKey, func);
+    }
+  }
   const attemptToVisitGlobals = precook(func, restOptions);
   return {
     function: func,
     attemptToVisitGlobals,
   };
 }
+
+export function clearFunctionASTCache(): void {
+  ASTCache = new WeakMap();
+}

packages/cook/src/preevaluate.spec.ts

@@ -1,8 +1,10 @@
+import { cook } from "./cook.js";
 import {
   isEvaluable,
   preevaluate,
   shouldAllowRecursiveEvaluations,
   isTrackAll,
+  clearExpressionASTCache,
 } from "./preevaluate.js";
 
 describe("isEvaluable", () => {
@@ -28,6 +30,10 @@ describe("isEvaluable", () => {
 });
 
 describe("preevaluate", () => {
+  beforeEach(() => {
+    clearExpressionASTCache();
+  });
+
   it("should work", () => {
     const { attemptToVisitGlobals, ...restResult } = preevaluate(
       "<% DATA, EVENT.detail %>"
@@ -52,6 +58,22 @@ describe("preevaluate", () => {
     });
   });
 
+  it("should isolate regexp", () => {
+    const precooked1 = preevaluate(
+      "<% ((r) => (r.exec('abc'), r.lastIndex))(/\\w/g) %>",
+      { cache: true }
+    );
+    const result1 = cook(precooked1.expression, precooked1.source);
+    expect(result1).toBe(1);
+
+    const precooked2 = preevaluate(
+      "<%= ((r) => (r.exec('abc'), r.lastIndex))(/\\w/g) %>",
+      { cache: true }
+    );
+    const result2 = cook(precooked2.expression, precooked2.source);
+    expect(result2).toBe(1);
+  });
+
   it("should throw SyntaxError", () => {
     expect(() => {
       preevaluate("<% DATA : EVENT %>");

packages/cook/src/preevaluate.ts

@@ -1,8 +1,12 @@
-import { Expression } from "@babel/types";
+import type { Expression } from "@babel/types";
 import { parseAsEstreeExpression } from "./parse.js";
 import { precook, PrecookOptions } from "./precook.js";
 
-export type PreevaluateOptions = Omit<PrecookOptions, "expressionOnly">;
+const ASTCache = new Map<string, Expression>();
+
+export type PreevaluateOptions = Omit<PrecookOptions, "expressionOnly"> & {
+  cache?: boolean;
+};
 
 export interface PreevaluateResult {
   expression: Expression;
@@ -15,16 +19,22 @@ export interface PreevaluateResult {
 // `raw` should always be asserted by `isEvaluable`.
 export function preevaluate(
   raw: string,
-  options?: PreevaluateOptions
+  { cache, ...restOptions }: PreevaluateOptions = {}
 ): PreevaluateResult {
   const fixes: string[] = [];
   const source = raw.replace(/^\s*<%[~=]?\s|\s%>\s*$/g, (m) => {
     fixes.push(m);
     return "";
   });
-  const expression = parseAsEstreeExpression(source);
+  let expression = cache ? ASTCache.get(source) : undefined;
+  if (!expression) {
+    expression = parseAsEstreeExpression(source);
+    if (cache) {
+      ASTCache.set(source, expression);
+    }
+  }
   const attemptToVisitGlobals = precook(expression, {
-    ...options,
+    ...restOptions,
     expressionOnly: true,
   });
   return {
@@ -47,3 +57,7 @@ export function shouldAllowRecursiveEvaluations(raw: string): boolean {
 export function isTrackAll(raw: string): boolean {
   return /^\s*<%=\s/.test(raw) && /\s%>\s*$/.test(raw);
 }
+
+export function clearExpressionASTCache(): void {
+  ASTCache.clear();
+}

packages/runtime/src/StoryboardFunctionRegistry.ts

@@ -148,6 +148,7 @@ export function StoryboardFunctionRegistryFactory({
       hooks: collector && {
         beforeVisit: collector.beforeVisit,
       },
+      cacheKey: fn,
     });
     const globalVariables = supply(
       precooked.attemptToVisitGlobals,

packages/runtime/src/internal/Router.ts

@@ -6,6 +6,10 @@ import type {
   Storyboard,
 } from "@next-core/types";
 import { HttpAbortError } from "@next-core/http";
+import {
+  clearExpressionASTCache,
+  clearFunctionASTCache,
+} from "@next-core/cook";
 import { uniqueId } from "lodash";
 import { NextHistoryState, NextLocation, getHistory } from "../history.js";
 import {
@@ -290,6 +294,11 @@ export class Router {
         ? previousApp.id !== currentApp.id
         : previousApp !== currentApp;
 
+    clearExpressionASTCache();
+    if (appChanged) {
+      clearFunctionASTCache();
+    }
+
     // TODO: handle favicon
 
     // Set `Router::#currentApp` before calling `getFeatureFlags()`

packages/runtime/src/internal/compute/evaluate.ts

@@ -151,6 +151,7 @@ function lowLevelEvaluate(
   let precooked: PreevaluateResult;
   try {
     precooked = preevaluate(raw, {
+      cache: true,
       withParent: true,
       hooks: {
         beforeVisitGlobal(node, parent) {

packages/utils/src/storyboard/expressions/track.ts

@@ -23,6 +23,7 @@ export function track(
       hasNonStaticUsage: false,
     };
     const { expression } = preevaluate(raw, {
+      cache: true,
       withParent: true,
       hooks: {
         beforeVisitGlobal: beforeVisitGlobalMember(usage, variableName),
@@ -58,6 +59,7 @@ export function trackAll(raw: string): trackAllResult | false {
       hasNonStaticUsage: false,
     };
     preevaluate(raw, {
+      cache: true,
       withParent: true,
       hooks: {
         beforeVisitGlobal: beforeVisitGlobalMember(usage, TRACK_NAMES, 1, true),