From eba0d4b635f68652de51c3221908a39d6eff9dd1 Mon Sep 17 00:00:00 2001
From: sanmaopep <maoyiwei.ebay@bytedance.com>
Date: Tue, 25 Feb 2025 19:16:11 +0800
Subject: [PATCH] feat: add isMatchAST API

---
 .../ast/variable-declaration.test.ts          | 21 +++++++++++++------
 .../variable-core/src/ast/index.ts            |  1 +
 .../variable-core/src/ast/utils/helpers.ts    | 11 ++++++++--
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/packages/variable-engine/variable-core/__tests__/ast/variable-declaration.test.ts b/packages/variable-engine/variable-core/__tests__/ast/variable-declaration.test.ts
index 632b5840..5be2837f 100644
--- a/packages/variable-engine/variable-core/__tests__/ast/variable-declaration.test.ts
+++ b/packages/variable-engine/variable-core/__tests__/ast/variable-declaration.test.ts
@@ -1,6 +1,13 @@
 import { vi, describe, test, expect } from 'vitest';
 
-import { ASTKind, ObjectType, VariableEngine, VariableDeclaration } from '../../src';
+import {
+  ASTKind,
+  ObjectType,
+  NumberType,
+  VariableEngine,
+  VariableDeclaration,
+  isMatchAST,
+} from '../../src';
 import { simpleVariableList } from '../../__mocks__/variables';
 import { getContainer } from '../../__mocks__/container';
 
@@ -26,7 +33,7 @@ describe('test Basic Variable Declaration', () => {
   });
 
   test('test globalVariableTable variables', () => {
-    expect(globalVariableTable.variables.map(_v => _v.key)).toMatchSnapshot();
+    expect(globalVariableTable.variables.map((_v) => _v.key)).toMatchSnapshot();
   });
 
   test('test get variable by key path', () => {
@@ -47,7 +54,7 @@ describe('test Basic Variable Declaration', () => {
       expect(globalVariableTable.getByKeyPath(_case)?.type.kind).toEqual(_resType);
     });
 
-    undefinedCases.forEach(_case => {
+    undefinedCases.forEach((_case) => {
       expect(globalVariableTable.getByKeyPath(_case)).toBeUndefined();
     });
   });
@@ -95,9 +102,9 @@ describe('test Basic Variable Declaration', () => {
     expect(previousVariable3.disposed).toBe(false);
     expect(previousVariable3.version).toBe(1); // 更新次数为一次
     expect(testScope.ast.version).toBe(2); // 调用了两次 fromJSON，因此更新了两次
-    expect(globalVariableTable.variables.map(_v => _v.key)).toMatchSnapshot();
+    expect(globalVariableTable.variables.map((_v) => _v.key)).toMatchSnapshot();
     expect(globalVariableTable.version).toBe(2); // 调用了两次 fromJSON，因此 version 是 2
-    expect(testScope.available.variables.map(_v => _v.key)).toMatchSnapshot();
+    expect(testScope.available.variables.map((_v) => _v.key)).toMatchSnapshot();
   });
 
   test('remove variables', () => {
@@ -141,7 +148,7 @@ describe('test Basic Variable Declaration', () => {
     declaration.subscribe(() => {
       declarationChangeTimes++;
     });
-    declaration.onTypeChange(type => {
+    declaration.onTypeChange((type) => {
       expect(type?.toJSON()).toMatchSnapshot();
       typeChangeTimes++;
     });
@@ -160,6 +167,7 @@ describe('test Basic Variable Declaration', () => {
       type: ASTKind.Number,
       meta: { label: 'test Label' },
     });
+    expect(isMatchAST(declaration.type, NumberType)).toBeTruthy();
     expect(declarationChangeTimes).toBe(2);
     expect(anyVariableChangeTimes).toBe(2);
     expect(typeChangeTimes).toBe(1);
@@ -178,6 +186,7 @@ describe('test Basic Variable Declaration', () => {
       },
       meta: { label: 'test Label' },
     });
+    expect(isMatchAST(declaration.type, ObjectType)).toBeTruthy();
     expect(declarationChangeTimes).toBe(3);
     expect(anyVariableChangeTimes).toBe(3);
     expect(typeChangeTimes).toBe(2);
diff --git a/packages/variable-engine/variable-core/src/ast/index.ts b/packages/variable-engine/variable-core/src/ast/index.ts
index e291ff6e..b8885980 100644
--- a/packages/variable-engine/variable-core/src/ast/index.ts
+++ b/packages/variable-engine/variable-core/src/ast/index.ts
@@ -17,3 +17,4 @@ export * from './expression';
 
 export { ASTFactory } from './factory';
 export { injectToAST, postConstructAST } from './utils/inversify';
+export { isMatchAST } from './utils/helpers';
diff --git a/packages/variable-engine/variable-core/src/ast/utils/helpers.ts b/packages/variable-engine/variable-core/src/ast/utils/helpers.ts
index 1831ecb5..4e75cfc9 100644
--- a/packages/variable-engine/variable-core/src/ast/utils/helpers.ts
+++ b/packages/variable-engine/variable-core/src/ast/utils/helpers.ts
@@ -13,7 +13,7 @@ export function updateChildNodeHelper(
     updateChildNode: (nextNode: ASTNode) => void;
     removeChildNode: () => void;
     nextJSON?: ASTNodeJSON;
-  },
+  }
 ): ASTNode | undefined {
   const currNode: ASTNode | undefined = getChildNode();
 
@@ -50,5 +50,12 @@ export function parseTypeJsonOrKind(typeJSONOrKind?: ASTNodeJSONOrKind): ASTNode
 
 // 获取所有的 children
 export function getAllChildren(ast: ASTNode): ASTNode[] {
-  return [...ast.children, ...ast.children.map(_child => getAllChildren(_child)).flat()];
+  return [...ast.children, ...ast.children.map((_child) => getAllChildren(_child)).flat()];
+}
+
+export function isMatchAST<TargetASTNode extends ASTNode>(
+  node?: ASTNode,
+  targetType?: { kind: string; new (...args: any[]): TargetASTNode }
+): node is TargetASTNode {
+  return node?.kind === targetType?.kind;
 }