feat: opt gc cost

packages/shader-lab/src/AstNodePool.ts

@@ -0,0 +1,54 @@
+import { IPoolElement } from "@galacean/engine";
+
+type Constructor<T> = new (...args: any[]) => T;
+
+export interface IInitializedPoolElement<T, C extends Constructor<T>> {
+  init(...args: ConstructorParameters<C>): void;
+}
+
+export const AstNodePoolSet: AstNodePool<any, any>[] = [];
+export function clearAllAstNodePool() {
+  for (let i = 0; i < AstNodePoolSet.length; i++) {
+    AstNodePoolSet[i].clear();
+  }
+}
+
+export class AstNodePool<C extends Constructor<T>, T extends IPoolElement & IInitializedPoolElement<T, C>> {
+  private _type: C;
+  private _elements: T[];
+  private _usedElementCount: number = -1;
+
+  constructor(type: C, initializeCount: number) {
+    this._type = type;
+    this._elements = new Array<T>(initializeCount);
+    for (let i = 0; i < initializeCount; i++) {
+      this._elements[i] = new type();
+    }
+    AstNodePoolSet.push(this);
+  }
+
+  garbageCollection(): void {
+    const elements = this._elements;
+    for (let i = elements.length - 1; i >= 0; i--) {
+      elements[i]?.dispose();
+    }
+    elements.length = 0;
+  }
+
+  get(...args: ConstructorParameters<C>): T {
+    this._usedElementCount++;
+    let element: T;
+    if (this._elements.length === this._usedElementCount) {
+      element = new this._type(...args);
+      this._elements.push(element);
+    } else {
+      element = this._elements[this._usedElementCount];
+    }
+    element.init(...args);
+    return element;
+  }
+
+  clear(): void {
+    this._usedElementCount = -1;
+  }
+}

packages/shader-lab/src/ShaderLab.ts

@@ -7,6 +7,7 @@ import { ShaderContentParser } from "./contentParser";
 // @ts-ignore
 import { ClearableObjectPool, Logger, ShaderLib, ShaderMacro, ShaderPlatformTarget } from "@galacean/engine";
 import { ShaderPosition, ShaderRange } from "./common";
+import { clearAllAstNodePool } from "./AstNodePool";
 
 export class ShaderLab implements IShaderLab {
   /**
@@ -76,6 +77,7 @@ export class ShaderLab implements IShaderLab {
 
     const lexer = new Lexer(ppdContent);
     const tokens = lexer.tokenize();
+    clearAllAstNodePool();
     const program = ShaderLab._parser.parse(tokens);
     const codeGen =
       backend === ShaderPlatformTarget.GLES100 ? GLES100Visitor.getVisitor() : GLES300Visitor.getVisitor();
@@ -108,7 +110,7 @@ export class ShaderLab implements IShaderLab {
     const passResult = [] as any;
     for (const subShader of structInfo.subShaders) {
       for (const pass of subShader.passes) {
-        if (!pass.isUsePass) continue;
+        if (pass.isUsePass) continue;
         const passInfo = this._parseShaderPass(
           pass.contents,
           pass.vertexEntry,

packages/shader-lab/src/codeGen/GLES100.ts

@@ -16,7 +16,7 @@ export class GLES100Visitor extends GLESVisitor {
 
   override getAttributeDeclare(): ICodeSegment[] {
     const ret: ICodeSegment[] = [];
-    for (const [_, item] of VisitorContext.context._referencedAttributeList) {
+    for (const item of Object.values(VisitorContext.context._referencedAttributeList)) {
       ret.push({
         text: `attribute ${item.typeInfo.typeLexeme} ${item.ident.lexeme};`,
         index: item.ident.location.start.index
@@ -27,7 +27,7 @@ export class GLES100Visitor extends GLESVisitor {
 
   override getVaryingDeclare(): ICodeSegment[] {
     const ret: ICodeSegment[] = [];
-    for (const [_, item] of VisitorContext.context._referencedVaryingList) {
+    for (const item of Object.values(VisitorContext.context._referencedVaryingList)) {
       ret.push({
         text: `varying ${item.typeInfo.typeLexeme} ${item.ident.lexeme};`,
         index: item.ident.location.start.index

packages/shader-lab/src/codeGen/GLES300.ts

@@ -23,7 +23,7 @@ export class GLES300Visitor extends GLESVisitor {
 
   override getAttributeDeclare(): ICodeSegment[] {
     const ret: ICodeSegment[] = [];
-    for (const [_, item] of VisitorContext.context._referencedAttributeList) {
+    for (const item of Object.values(VisitorContext.context._referencedAttributeList)) {
       ret.push({
         text: `in ${item.typeInfo.typeLexeme} ${item.ident.lexeme};`,
         index: item.ident.location.start.index
@@ -35,7 +35,9 @@ export class GLES300Visitor extends GLESVisitor {
   override getVaryingDeclare(): ICodeSegment[] {
     const ret: ICodeSegment[] = [];
     const qualifier = VisitorContext.context.stage === EShaderStage.FRAGMENT ? "in" : "out";
-    for (const [_, item] of VisitorContext.context._referencedVaryingList) {
+    const values = Object.values(VisitorContext.context._referencedVaryingList);
+    for (let i = 0; i < values.length; i++) {
+      const item = values[i];
       ret.push({
         text: `${item.qualifier ?? qualifier} ${item.typeInfo.typeLexeme} ${item.ident.lexeme};`,
         index: item.ident.location.start.index
@@ -80,13 +82,13 @@ export class GLES300Visitor extends GLESVisitor {
 
   override visitVariableIdentifier(node: ASTNode.VariableIdentifier): string {
     if (VisitorContext.context.stage === EShaderStage.FRAGMENT && node.lexeme === "gl_FragColor") {
-      if (!VisitorContext.context._referencedVaryingList.has(V3_GL_FragColor)) {
-        VisitorContext.context._referencedVaryingList.set(V3_GL_FragColor, {
+      if (!VisitorContext.context._referencedVaryingList[V3_GL_FragColor]) {
+        VisitorContext.context._referencedVaryingList[V3_GL_FragColor] = {
           ident: new Token(ETokenType.ID, V3_GL_FragColor, ShaderLab.createPosition(0, 0, 0)),
           typeInfo: new SymbolType(EKeyword.VEC4, "vec4"),
           qualifier: "out",
           astNode: node
-        });
+        };
       }
       return V3_GL_FragColor;
     }

packages/shader-lab/src/codeGen/GLESVisitor.ts

@@ -109,15 +109,17 @@ export abstract class GLESVisitor extends CodeGenVisitor {
   ): ICodeSegment[] {
     const { _referencedGlobals } = VisitorContext.context;
 
-    if (lastLength === _referencedGlobals.size) {
+    if (lastLength === Object.keys(_referencedGlobals).length) {
       for (const precision of data.globalPrecisions) {
         textList.push({ text: precision.codeGen(this), index: precision.location.start.index });
       }
       return textList;
     }
 
-    lastLength = _referencedGlobals.size;
-    for (const [ident, sm] of _referencedGlobals) {
+    lastLength = Object.keys(_referencedGlobals).length;
+    for (const ident in _referencedGlobals) {
+      const sm = _referencedGlobals[ident];
+
       if (_serialized.has(ident)) continue;
       _serialized.add(ident);
 

packages/shader-lab/src/codeGen/VisitorContext.ts

@@ -23,9 +23,9 @@ export class VisitorContext {
 
   stage: EShaderStage;
 
-  _referencedAttributeList: Map<string, IParamInfo & { qualifier?: string }> = new Map();
-  _referencedGlobals: Map<string, SymbolInfo | ASTNode.PrecisionSpecifier> = new Map();
-  _referencedVaryingList: Map<string, IParamInfo & { qualifier?: string }> = new Map();
+  _referencedAttributeList: Record<string, IParamInfo & { qualifier?: string }> = Object.create(null);
+  _referencedGlobals: Record<string, SymbolInfo | ASTNode.PrecisionSpecifier> = Object.create(null);
+  _referencedVaryingList: Record<string, IParamInfo & { qualifier?: string }> = Object.create(null);
 
   _curFn?: ASTNode.FunctionProtoType;
 
@@ -39,9 +39,9 @@ export class VisitorContext {
   reset() {
     this.attributeList.length = 0;
     this.attributeStructs.length = 0;
-    this._referencedAttributeList.clear();
-    this._referencedGlobals.clear();
-    this._referencedVaryingList.clear();
+    this._referencedAttributeList = Object.create(null);
+    this._referencedGlobals = Object.create(null);
+    this._referencedVaryingList = Object.create(null);
   }
 
   isAttributeStruct(type: string) {
@@ -53,41 +53,41 @@ export class VisitorContext {
   }
 
   referenceAttribute(ident: string) {
-    if (this._referencedAttributeList.has(ident)) return;
+    if (this._referencedAttributeList[ident]) return;
 
     const prop = this.attributeList.find((item) => item.ident.lexeme === ident);
     if (!prop) {
       Logger.error("referenced attribute not found:", ident);
       return;
     }
-    this._referencedAttributeList.set(ident, prop);
+    this._referencedAttributeList[ident] = prop;
   }
 
   referenceVarying(ident: string) {
-    if (this._referencedVaryingList.has(ident)) return;
+    if (this._referencedVaryingList[ident]) return;
 
     const prop = this.varyingStruct?.propList.find((item) => item.ident.lexeme === ident);
     if (!prop) {
       Logger.error("referenced varying not found:", ident);
       return;
     }
-    this._referencedVaryingList.set(ident, prop);
+    this._referencedVaryingList[ident] = prop;
   }
 
   referenceGlobal(ident: string, type: ESymbolType) {
-    if (this._referencedGlobals.has(ident)) return;
+    if (this._referencedGlobals[ident]) return;
 
     if (type === ESymbolType.FN) {
       const fnEntries = this._passSymbolTable.getAllFnSymbols(ident);
       for (let i = 0; i < fnEntries.length; i++) {
         const key = i === 0 ? ident : ident + i;
-        this._referencedGlobals.set(key, fnEntries[i]);
+        this._referencedGlobals[key] = fnEntries[i];
       }
       return;
     }
     const sm = this.passSymbolTable.lookup({ ident, symbolType: type });
     if (sm) {
-      this._referencedGlobals.set(ident, sm);
+      this._referencedGlobals[ident] = sm;
     }
   }
 }