Fix some shaderlab bug

packages/shader-lab/src/Preprocessor.ts

@@ -5,6 +5,7 @@ import { ShaderLib } from "@galacean/engine";
 export enum MacroValueType {
   Number, // 1, 1.1
   Symbol, // variable name
+  MemberAccess, // member access, e.g. input.v_uv, v.rgb
   FunctionCall, // function call, e.g. clamp(a, 0.0, 1.0)
   Other // shaderLab does not check this
 }
@@ -27,6 +28,7 @@ export class Preprocessor {
   private static readonly _macroRegex =
     /^\s*#define\s+(\w+)[ ]*(\(([^)]*)\))?[ ]+(\(?\w+\)?.*?)(?:\/\/.*|\/\*.*?\*\/)?\s*$/gm;
   private static readonly _symbolReg = /^[a-zA-Z_][a-zA-Z0-9_]*$/;
+  private static readonly _memberAccessReg = /^([a-zA-Z_][a-zA-Z0-9_]*)(\.[a-zA-Z_][a-zA-Z0-9_]*)+$/;
   private static readonly _funcCallReg = /^([a-zA-Z_][a-zA-Z0-9_]*)\s*\((.*)\)$/;
   private static readonly _macroDefineIncludeMap = new Map<string, MacroDefineList>();
 
@@ -61,6 +63,10 @@ export class Preprocessor {
         const referencedName = valueType === MacroValueType.FunctionCall ? info.functionCallName : info.value;
         if (info.params.indexOf(referencedName) !== -1) continue;
         if (out.indexOf(referencedName) === -1) out.push(referencedName);
+      } else if (valueType === MacroValueType.MemberAccess) {
+        // Extract root symbol: "input.v_uv" → "input"
+        const rootName = info.value.substring(0, info.value.indexOf("."));
+        if (out.indexOf(rootName) === -1) out.push(rootName);
       } else if (valueType === MacroValueType.Other) {
         // #if _VERBOSE
         Logger.warn(
@@ -110,6 +116,8 @@ export class Preprocessor {
         valueType = MacroValueType.Number;
       } else if (this._symbolReg.test(value)) {
         valueType = MacroValueType.Symbol;
+      } else if (this._memberAccessReg.test(value)) {
+        valueType = MacroValueType.MemberAccess;
       } else {
         const callMatch = this._funcCallReg.exec(value);
         if (callMatch) {

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

@@ -31,13 +31,19 @@ export abstract class CodeGenVisitor {
 
   protected static _tmpArrayPool = new ReturnableObjectPool(TempArray<string>, 10);
 
+  protected static readonly _memberAccessReg = /\b(\w+)\.(\w+)\b/g;
+
   defaultCodeGen(children: NodeChild[]) {
     const pool = CodeGenVisitor._tmpArrayPool;
     let ret = pool.get();
     ret.dispose();
     for (const child of children) {
       if (child instanceof BaseToken) {
-        ret.array.push(child.lexeme);
+        if (child.type === Keyword.MACRO_DEFINE_EXPRESSION) {
+          ret.array.push(this._transformMacroDefineValue(child.lexeme));
+        } else {
+          ret.array.push(child.lexeme);
+        }
       } else {
         ret.array.push(child.codeGen(this));
       }
@@ -46,6 +52,32 @@ export abstract class CodeGenVisitor {
     return ret.array.join(" ");
   }
 
+  protected _transformMacroDefineValue(
+    lexeme: string,
+    overrideMap?: Record<string, "varying" | "attribute" | "mrt">
+  ): string {
+    const context = VisitorContext.context;
+    const structVarMap = overrideMap ?? context._structVarMap;
+    if (!structVarMap) return lexeme;
+
+    const spaceIdx = lexeme.indexOf(" ");
+    if (spaceIdx === -1) return lexeme;
+
+    const macroName = lexeme.substring(0, spaceIdx);
+    let value = lexeme.substring(spaceIdx);
+
+    const reg = CodeGenVisitor._memberAccessReg;
+    reg.lastIndex = 0;
+    value = value.replace(reg, (match, varName, propName) => {
+      const role = structVarMap[varName];
+      if (!role) return match;
+      context.referenceStructPropByName(role, propName);
+      return propName;
+    });
+
+    return macroName + value;
+  }
+
   visitPostfixExpression(node: ASTNode.PostfixExpression): string {
     const children = node.children;
     const derivationLength = children.length;
@@ -351,6 +383,8 @@ export abstract class CodeGenVisitor {
     const fnName = fnNode.protoType.ident.lexeme;
     const context = VisitorContext.context;
 
+    this._collectStructVars(fnNode, context);
+
     if (fnName == context.stageEntry) {
       const statements = fnNode.statements.codeGen(this);
       return `void main() ${statements}`;
@@ -359,6 +393,74 @@ export abstract class CodeGenVisitor {
     }
   }
 
+  private _collectStructVars(fnNode: ASTNode.FunctionDefinition, context: VisitorContext): void {
+    const map = context._structVarMap;
+    // Clear previous function's mappings
+    for (const key in map) delete map[key];
+
+    // Collect from function parameters
+    const paramList = fnNode.protoType.parameterList;
+    if (paramList) {
+      for (const param of paramList) {
+        if (param.ident && param.typeInfo && typeof param.typeInfo.type === "string") {
+          const role = context.getStructRole(param.typeInfo.typeLexeme);
+          if (role) map[param.ident.lexeme] = role;
+        }
+      }
+    }
+
+    // Collect from local variable declarations in function body
+    this._collectStructVarsFromNode(fnNode.statements, context, map);
+  }
+
+  private _collectStructVarsFromNode(
+    node: TreeNode,
+    context: VisitorContext,
+    map: Record<string, "varying" | "attribute" | "mrt">
+  ): void {
+    const children = node.children;
+    for (let i = 0; i < children.length; i++) {
+      const child = children[i];
+      if (child instanceof ASTNode.InitDeclaratorList) {
+        const typeLexeme = child.typeInfo?.typeLexeme;
+        if (typeLexeme) {
+          const role = context.getStructRole(typeLexeme);
+          if (role) {
+            // Extract variable name from SingleDeclaration or comma-separated identifiers
+            this._extractVarNamesFromInitDeclaratorList(child, map, role);
+          }
+        }
+      } else if (child instanceof TreeNode) {
+        this._collectStructVarsFromNode(child, context, map);
+      }
+    }
+  }
+
+  protected _extractVarNamesFromInitDeclaratorList(
+    node: ASTNode.InitDeclaratorList,
+    map: Record<string, "varying" | "attribute" | "mrt">,
+    role: "varying" | "attribute" | "mrt"
+  ): void {
+    const children = node.children;
+    if (children.length === 1) {
+      // SingleDeclaration: type ident
+      const singleDecl = children[0] as ASTNode.SingleDeclaration;
+      const identChildren = singleDecl.children;
+      if (identChildren.length >= 2 && identChildren[1] instanceof BaseToken) {
+        map[identChildren[1].lexeme] = role;
+      }
+    } else if (children.length >= 3) {
+      // InitDeclaratorList , ident ...
+      const initDeclList = children[0];
+      if (initDeclList instanceof ASTNode.InitDeclaratorList) {
+        this._extractVarNamesFromInitDeclaratorList(initDeclList, map, role);
+      }
+      if (children[2] instanceof BaseToken) {
+        map[children[2].lexeme] = role;
+      }
+    }
+  }
+
   protected _reportError(loc: ShaderRange | ShaderPosition, message: string): void {
     // #if _VERBOSE
     this.errors.push(new GSError(GSErrorName.CompilationError, message, loc, ShaderLab._processingPassText));

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

@@ -5,6 +5,7 @@ import { Keyword } from "../common/enums/Keyword";
 import { ASTNode, TreeNode } from "../parser/AST";
 import { ShaderData } from "../parser/ShaderInfo";
 import { ESymbolType, FnSymbol, StructSymbol, SymbolInfo } from "../parser/symbolTable";
+import { NodeChild } from "../parser/types";
 import { CodeGenVisitor } from "./CodeGenVisitor";
 import { ICodeSegment } from "./types";
 import { VisitorContext } from "./VisitorContext";
@@ -37,10 +38,15 @@ export abstract class GLESVisitor extends CodeGenVisitor {
     this.reset();
 
     const shaderData = node.shaderData;
-    VisitorContext.context._passSymbolTable = shaderData.symbolTable;
+    const context = VisitorContext.context;
+    context._passSymbolTable = shaderData.symbolTable;
 
     const outerGlobalMacroDeclarations = shaderData.getOuterGlobalMacroDeclarations();
 
+    // Build combined _globalStructVarMap from both entry functions before per-stage processing.
+    // This must happen here because vertex runs first and doesn't yet know fragment's variables.
+    this._buildGlobalStructVarMap(vertexEntry, fragmentEntry, shaderData, context);
+
     return {
       vertex: this._vertexMain(vertexEntry, shaderData, outerGlobalMacroDeclarations),
       fragment: this._fragmentMain(fragmentEntry, shaderData, outerGlobalMacroDeclarations)
@@ -109,6 +115,9 @@ export abstract class GLESVisitor extends CodeGenVisitor {
       }
     });
 
+    // Pre-register global #define member access references for this stage
+    this._registerGlobalMacroReferences(outerGlobalMacroDeclarations, context);
+
     const globalCodeArray = this._globalCodeArray;
     VisitorContext.context.referenceGlobal(entry, ESymbolType.FN);
 
@@ -173,6 +182,9 @@ export abstract class GLESVisitor extends CodeGenVisitor {
       }
     });
 
+    // Pre-register global #define member access references for this stage
+    this._registerGlobalMacroReferences(outerGlobalMacroStatements, context);
+
     const globalCodeArray = this._globalCodeArray;
     VisitorContext.context.referenceGlobal(entry, ESymbolType.FN);
 
@@ -193,6 +205,127 @@ export abstract class GLESVisitor extends CodeGenVisitor {
     return globalCode;
   }
 
+  /**
+   * Build _globalStructVarMap from both entry functions before per-stage processing.
+   * Classifies struct types by their position in function signatures:
+   * - vertex param[0] → attribute, vertex return type → varying
+   * - fragment param[0] → varying, fragment return type → mrt
+   */
+  private _buildGlobalStructVarMap(
+    vertexEntry: string,
+    fragmentEntry: string,
+    data: ShaderData,
+    context: VisitorContext
+  ): void {
+    const map = context._globalStructVarMap;
+    const lookupSymbol = GLESVisitor._lookupSymbol;
+    const { symbolTable } = data;
+
+    // Map struct type names to roles based on function signature positions
+    const structTypeRoles: Record<string, "varying" | "attribute" | "mrt"> = Object.create(null);
+
+    // Vertex entry: param[0] type → attribute, return type → varying
+    lookupSymbol.set(vertexEntry, ESymbolType.FN);
+    const vertexFns = <FnSymbol[]>symbolTable.getSymbols(lookupSymbol, true, []);
+    for (const fn of vertexFns) {
+      const proto = fn.astNode.protoType;
+      const param0 = proto.parameterList?.[0];
+      if (param0 && typeof param0.typeInfo.type === "string") {
+        structTypeRoles[param0.typeInfo.typeLexeme] = "attribute";
+      }
+      if (typeof proto.returnType.type === "string") {
+        structTypeRoles[<string>proto.returnType.type] = "varying";
+      }
+    }
+
+    // Fragment entry: param[0] type → varying, return type → mrt
+    lookupSymbol.set(fragmentEntry, ESymbolType.FN);
+    const fragmentFns = <FnSymbol[]>symbolTable.getSymbols(lookupSymbol, true, []);
+    for (const fn of fragmentFns) {
+      const proto = fn.astNode.protoType;
+      const param0 = proto.parameterList?.[0];
+      if (param0 && typeof param0.typeInfo.type === "string") {
+        structTypeRoles[param0.typeInfo.typeLexeme] = "varying";
+      }
+      if (typeof proto.returnType.type === "string") {
+        structTypeRoles[<string>proto.returnType.type] = "mrt";
+      }
+    }
+
+    // Scan all entry functions' params and local vars, classify by structTypeRoles
+    for (const fn of [...vertexFns, ...fragmentFns]) {
+      const fnNode = fn.astNode;
+      const paramList = fnNode.protoType.parameterList;
+      if (paramList) {
+        for (const param of paramList) {
+          if (param.ident && param.typeInfo && typeof param.typeInfo.type === "string") {
+            const role = structTypeRoles[param.typeInfo.typeLexeme];
+            if (role) map[param.ident.lexeme] = role;
+          }
+        }
+      }
+      this._collectStructVarsFromBody(fnNode.statements, structTypeRoles, map);
+    }
+  }
+
+  private _collectStructVarsFromBody(
+    node: TreeNode,
+    structTypeRoles: Record<string, "varying" | "attribute" | "mrt">,
+    map: Record<string, "varying" | "attribute" | "mrt">
+  ): void {
+    const children = node.children;
+    for (let i = 0; i < children.length; i++) {
+      const child = children[i];
+      if (child instanceof ASTNode.InitDeclaratorList) {
+        const typeLexeme = child.typeInfo?.typeLexeme;
+        if (typeLexeme) {
+          const role = structTypeRoles[typeLexeme];
+          if (role) {
+            this._extractVarNamesFromInitDeclaratorList(child, map, role);
+          }
+        }
+      } else if (child instanceof TreeNode) {
+        this._collectStructVarsFromBody(child, structTypeRoles, map);
+      }
+    }
+  }
+
+  /**
+   * Pre-register attribute/varying/mrt references from global #define member access patterns,
+   * so that declarations are emitted by _getCustomStruct for the current stage.
+   */
+  private _registerGlobalMacroReferences(globalMacros: ASTNode.GlobalDeclaration[], context: VisitorContext): void {
+    const map = context._globalStructVarMap;
+    if (!Object.keys(map).length) return;
+    const reg = CodeGenVisitor._memberAccessReg;
+    for (const macro of globalMacros) {
+      this._scanMacroMemberAccess(macro.children, reg, map, context);
+    }
+  }
+
+  private _scanMacroMemberAccess(
+    children: NodeChild[],
+    reg: RegExp,
+    map: Record<string, "varying" | "attribute" | "mrt">,
+    context: VisitorContext
+  ): void {
+    for (const child of children) {
+      if (child instanceof BaseToken && child.type === Keyword.MACRO_DEFINE_EXPRESSION) {
+        const spaceIdx = child.lexeme.indexOf(" ");
+        if (spaceIdx === -1) continue;
+        const value = child.lexeme.substring(spaceIdx);
+        reg.lastIndex = 0;
+        let match: RegExpExecArray | null;
+        while ((match = reg.exec(value)) !== null) {
+          const role = map[match[1]];
+          if (role) context.referenceStructPropByName(role, match[2]);
+        }
+      } else if (child instanceof TreeNode) {
+        this._scanMacroMemberAccess(child.children, reg, map, context);
+      }
+    }
+  }
+
   private _getGlobalSymbol(out: ICodeSegment[]): void {
     const { _referencedGlobals } = VisitorContext.context;
 
@@ -233,6 +366,7 @@ export abstract class GLESVisitor extends CodeGenVisitor {
 
   private _getGlobalMacroDeclarations(macros: ASTNode.GlobalDeclaration[], out: ICodeSegment[]): void {
     const context = VisitorContext.context;
+    const globalMap = context._globalStructVarMap;
     const referencedGlobals = context._referencedGlobals;
     const referencedGlobalMacroASTs = context._referencedGlobalMacroASTs;
     referencedGlobalMacroASTs.length = 0;
@@ -253,7 +387,12 @@ export abstract class GLESVisitor extends CodeGenVisitor {
         let result: ICodeSegment[] = [];
         result.push(
           ...macro.macroExpressions.map((item) => ({
-            text: item instanceof BaseToken ? item.lexeme : item.codeGen(this),
+            text:
+              item instanceof BaseToken
+                ? item.type === Keyword.MACRO_DEFINE_EXPRESSION
+                  ? this._transformMacroDefineValue(item.lexeme, globalMap)
+                  : item.lexeme
+                : item.codeGen(this),
             index: item.location.start.index
           }))
         );
@@ -264,6 +403,8 @@ export abstract class GLESVisitor extends CodeGenVisitor {
           .sort((a, b) => a.index - b.index)
           .map((item) => item.text)
           .join("\n");
+      } else if (child instanceof BaseToken && child.type === Keyword.MACRO_DEFINE_EXPRESSION) {
+        text = this._transformMacroDefineValue(child.lexeme, globalMap);
       } else {
         text = macro.codeGen(this);
       }