perf: cache AST for evaluator on reearth/core (#473)

* perf: cache AST for evaluator on reearth/core * fix: a way to clear cache * test: fix
reearth · Feb 20, 2023 · da6bb3a · da6bb3a
1 parent ebb50de
commit da6bb3a
Show file tree

Hide file tree

Showing 9 changed files with 128 additions and 25 deletions.
diff --git a/src/core/Map/Layer/hooks.ts b/src/core/Map/Layer/hooks.ts
@@ -1,7 +1,14 @@
 import { useAtom } from "jotai";
-import { useCallback, useLayoutEffect, useMemo, useRef } from "react";
+import { useCallback, useEffect, useLayoutEffect, useMemo, useRef } from "react";
 
-import { computeAtom, DataType, type Atom, evalFeature, ComputedFeature } from "../../mantle";
+import {
+  clearAllExpressionCaches,
+  computeAtom,
+  DataType,
+  type Atom,
+  evalFeature,
+  ComputedFeature,
+} from "../../mantle";
 import type { DataRange, Feature, Layer } from "../../mantle";
 
 export type { Atom as Atoms } from "../../mantle";
@@ -74,6 +81,22 @@ export default function useHooks(
     };
   }, [layer, forceUpdateFeatures]);
 
+  // Clear expression cache if layer is unmounted
+  useEffect(
+    () => () => {
+      window.requestIdleCallback(() => {
+        // This is a little heavy task, and not critical for main functionality, so we can run this at idle time.
+        computedLayer?.originalFeatures.forEach(f => {
+          clearAllExpressionCaches(
+            computedLayer.layer.type === "simple" ? computedLayer.layer : undefined,
+            f,
+          );
+        });
+      });
+    },
+    [computedLayer],
+  );
+
   return {
     computedLayer,
     handleFeatureRequest: requestFetch,

diff --git a/src/core/engines/Cesium/Feature/utils.tsx b/src/core/engines/Cesium/Feature/utils.tsx
@@ -20,7 +20,7 @@ import {
 } from "react";
 import { type CesiumComponentRef, Entity } from "resium";
 
-import { Data, LayerSimple, TimeInterval } from "@reearth/core/mantle";
+import { Data, Layer, LayerSimple, TimeInterval } from "@reearth/core/mantle";
 
 import type { ComputedFeature, ComputedLayer, FeatureComponentProps, Geometry } from "../..";
 
@@ -141,11 +141,14 @@ const tagKeys = Object.keys(tagObj) as (keyof Tag)[];
 
 const tagKey = "__reearth_tag";
 
-export const extractSimpleLayer = (layer: ComputedLayer | undefined): LayerSimple | void => {
-  if (layer?.layer.type !== "simple") {
+export const extractSimpleLayer = (
+  layer: ComputedLayer | Layer | undefined,
+): LayerSimple | void => {
+  const l = layer && "layer" in layer ? layer.layer : layer;
+  if (l?.type !== "simple") {
     return;
   }
-  return layer.layer;
+  return l;
 };
 
 export const extractSimpleLayerData = (layer: ComputedLayer | undefined): Data | void => {

diff --git a/src/core/mantle/evaluator/index.ts b/src/core/mantle/evaluator/index.ts
@@ -10,6 +10,8 @@ import {
 
 import { evalSimpleLayerFeature, evalSimpleLayer } from "./simple";
 
+export { clearAllExpressionCaches } from "./simple";
+
 export type EvalContext = {
   getFeatures: (d: Data, r?: DataRange) => Promise<Feature[] | undefined>;
   getAllFeatures: (d: Data) => Promise<Feature[] | undefined>;

diff --git a/src/core/mantle/evaluator/simple/expression/expression.test.ts b/src/core/mantle/evaluator/simple/expression/expression.test.ts
@@ -1,6 +1,16 @@
-import { expect, test, describe } from "vitest";
+import jsep from "jsep";
+import { expect, test, describe, beforeEach } from "vitest";
 
-import { replaceDefines, removeBackslashes } from "./expression";
+import { Feature } from "../../../types";
+
+import {
+  replaceDefines,
+  removeBackslashes,
+  Expression,
+  EXPRESSION_CACHES,
+  clearExpressionCaches,
+} from "./expression";
+import { createRuntimeAst } from "./runtime";
 
 describe("replaceDefines", () => {
   test("should replace defined placeholders with the corresponding values in the expression string", () => {
@@ -28,3 +38,20 @@ describe("removeBackslashes", () => {
     expect(result).toBe("@#%@#%");
   });
 });
+
+describe("expression caches", () => {
+  beforeEach(() => {
+    EXPRESSION_CACHES.clear();
+  });
+
+  test("should remove caches", () => {
+    const key = "czm_HEIGHT > 2 ? color('red') : color('blue')";
+    const feature = { properties: { HEIGHT: 1 } } as Feature;
+    const expression = new Expression("${HEIGHT} > 2 ? color('red') : color('blue')", feature);
+    expect(EXPRESSION_CACHES.get(key)).toEqual(createRuntimeAst(expression, jsep(key)));
+
+    clearExpressionCaches(key, feature, undefined);
+
+    expect(EXPRESSION_CACHES.get(key)).toBeUndefined();
+  });
+});
diff --git a/src/core/mantle/evaluator/simple/expression/expression.ts b/src/core/mantle/evaluator/simple/expression/expression.ts
@@ -13,6 +13,8 @@ export type JPLiteral = {
   literalValue: any;
 };
 
+export const EXPRESSION_CACHES = new Map<string, Node | Error>();
+
 export class Expression {
   #expression: string;
   #runtimeAst: Node | Error;
@@ -28,24 +30,30 @@ export class Expression {
       this.#feature?.properties,
     );
 
-    if (literalJP.length !== 0) {
-      for (const elem of literalJP) {
-        jsep.addLiteral(elem.literalName, elem.literalValue);
+    const cachedAST = EXPRESSION_CACHES.get(expression);
+    if (cachedAST) {
+      this.#runtimeAst = cachedAST;
+    } else {
+      if (literalJP.length !== 0) {
+        for (const elem of literalJP) {
+          jsep.addLiteral(elem.literalName, elem.literalValue);
+        }
       }
-    }
 
-    // customize jsep operators
-    jsep.addBinaryOp("=~", 0);
-    jsep.addBinaryOp("!~", 0);
+      // customize jsep operators
+      jsep.addBinaryOp("=~", 0);
+      jsep.addBinaryOp("!~", 0);
 
-    let ast;
-    try {
-      ast = jsep(expression);
-    } catch (e) {
-      throw new Error(`failed to generate ast: ${e}`);
-    }
+      let ast;
+      try {
+        ast = jsep(expression);
+      } catch (e) {
+        throw new Error(`failed to generate ast: ${e}`);
+      }
 
-    this.#runtimeAst = createRuntimeAst(this, ast);
+      this.#runtimeAst = createRuntimeAst(this, ast);
+      EXPRESSION_CACHES.set(expression, this.#runtimeAst);
+    }
   }
 
   evaluate() {
@@ -71,3 +79,13 @@ export function replaceDefines(expression: string, defines: any): string {
 export function removeBackslashes(expression: string): string {
   return expression.replace(backslashRegex, backslashReplacement);
 }
+
+export function clearExpressionCaches(
+  expression: string,
+  feature: Feature | undefined,
+  defines: any | undefined,
+) {
+  expression = replaceDefines(expression, defines);
+  [expression] = replaceVariables(removeBackslashes(expression), feature?.properties);
+  EXPRESSION_CACHES.delete(expression);
+}
diff --git a/src/core/mantle/evaluator/simple/expression/index.ts b/src/core/mantle/evaluator/simple/expression/index.ts
@@ -1 +1 @@
-export { Expression } from "./expression";
+export { Expression, clearExpressionCaches } from "./expression";
diff --git a/src/core/mantle/evaluator/simple/index.ts b/src/core/mantle/evaluator/simple/index.ts
@@ -14,7 +14,7 @@ import {
 import { defined } from "../../utils";
 
 import { ConditionalExpression } from "./conditionalExpression";
-import { Expression } from "./expression";
+import { clearExpressionCaches, Expression } from "./expression";
 import { evalTimeInterval } from "./interval";
 
 export async function evalSimpleLayer(
@@ -72,6 +72,31 @@ export function evalLayerAppearances(
   );
 }
 
+export function clearAllExpressionCaches(
+  layer: LayerSimple | undefined,
+  feature: Feature | undefined,
+) {
+  const appearances: Partial<LayerAppearanceTypes> = pick(layer, appearanceKeys);
+
+  Object.entries(appearances).forEach(([, v]) => {
+    Object.entries(v).forEach(([, expressionContainer]) => {
+      if (hasExpression(expressionContainer)) {
+        const styleExpression = expressionContainer.expression;
+        if (typeof styleExpression === "object" && styleExpression.conditions) {
+          styleExpression.conditions.forEach(([expression1, expression2]) => {
+            clearExpressionCaches(expression1, feature, layer?.defines);
+            clearExpressionCaches(expression2, feature, layer?.defines);
+          });
+        } else if (typeof styleExpression === "boolean" || typeof styleExpression === "number") {
+          clearExpressionCaches(String(styleExpression), feature, layer?.defines);
+        } else if (typeof styleExpression === "string") {
+          clearExpressionCaches(styleExpression, feature, layer?.defines);
+        }
+      }
+    });
+  });
+}
+
 function hasExpression(e: any): e is ExpressionContainer {
   return typeof e === "object" && e && "expression" in e;
 }

diff --git a/src/core/mantle/index.ts b/src/core/mantle/index.ts
@@ -1,5 +1,5 @@
 export * from "./types";
 export * from "./compat";
-export { evalFeature } from "./evaluator";
+export { evalFeature, clearAllExpressionCaches } from "./evaluator";
 export { computeAtom } from "./atoms";
 export type { Atom } from "./atoms";
diff --git a/src/test/setup.ts b/src/test/setup.ts
@@ -33,4 +33,9 @@ Object.defineProperty(window, "matchMedia", {
   })),
 });
 
+Object.defineProperty(window, "requestIdleCallback", {
+  writable: true,
+  value: vi.fn(),
+});
+
 afterEach(cleanup);