import flatten from 'lodash/flatten' ; import isFinite from 'lodash/isFinite' ; import { GLSLContext } from '.' ; import { AST_NODE_TYPES as SHADER_AST_NODE_TYPES , AST_TOKEN_TYPES as SHADER_AST_TOKEN_TYPES , STORAGE_CLASS , } from './ast/glsl-ast-node-types' ; import { ArrayExpression , BaseNode as ShaderBaseNode , BlockStatement , CallExpression as ShaderCallExpression , DataType , DoWhileStatement as ShaderDoWhileStatement , Expression as ShaderExpression , ExpressionStatement , ForStatement as ShaderForStatement , FunctionDeclaration as ShaderFunctionDeclaration , Identifier as ShaderIdentifier , IfStatement as ShaderIfStatement , Program as ShaderProgram , ReturnStatement as ShaderReturnStatement , Scope , Statement as ShaderStatement , UpdateExpression as ShaderUpdateExpression , VariableDeclaration as ShaderVariableDeclaration , VariableDeclarator as ShaderVariableDeclarator , WhileStatement as ShaderWhileStatement , } from './ast/glsl-tree' ; import { AST_NODE_TYPES } from './ast/ts-ast-node-types' ; import { ClassDeclaration , ClassProperty , Decorator , DoWhileStatement , Expression , ForStatement , FunctionDeclaration , Identifier , IfStatement , MethodDefinition , Program , Statement , VariableDeclaration , WhileStatement , } from './ast/ts-estree' ; import { Target } from './backends/ICodeGenerator' ; import { DataTypeInference } from './DataTypeInference' ; import { exportFunctions } from './utils/builtin-functions' ; import { getComponentDataTypeFromVector } from './utils/data-type' ; import { addIdentifierToScope , getIdentifierFromScope , getScopeByIdentifierName , } from './utils/node' ; enum PropertyDecorator { In = 'in' , Out = 'out' , Shared = 'shared' } interface IDecorator { name : string ; params : number [ ] ; } export const dataTypeMap = { void : SHADER_AST_TOKEN_TYPES . Void , bool : SHADER_AST_TOKEN_TYPES . Boolean , int : SHADER_AST_TOKEN_TYPES . Int32 , uint : SHADER_AST_TOKEN_TYPES . Uint32 , float : SHADER_AST_TOKEN_TYPES . Float , vec2 : SHADER_AST_TOKEN_TYPES . Vector2Float , vec3 : SHADER_AST_TOKEN_TYPES . Vector3Float , vec4 : SHADER_AST_TOKEN_TYPES . Vector4Float , ivec2 : SHADER_AST_TOKEN_TYPES . Vector2Int , ivec3 : SHADER_AST_TOKEN_TYPES . Vector3Int , ivec4 : SHADER_AST_TOKEN_TYPES . Vector4Int , uvec2 : SHADER_AST_TOKEN_TYPES . Vector2Uint , uvec3 : SHADER_AST_TOKEN_TYPES . Vector3Uint , uvec4 : SHADER_AST_TOKEN_TYPES . Vector4Uint , bvec2 : SHADER_AST_TOKEN_TYPES . Vector2Boolean , bvec3 : SHADER_AST_TOKEN_TYPES . Vector3Boolean , bvec4 : SHADER_AST_TOKEN_TYPES . Vector4Boolean , mat3 : SHADER_AST_TOKEN_TYPES . Matrix3x3Float , mat4 : SHADER_AST_TOKEN_TYPES . Matrix4x4Float , 'float[]' : SHADER_AST_TOKEN_TYPES . FloatArray , 'vec4[]' : SHADER_AST_TOKEN_TYPES . Vector4FloatArray , } ; const emptyShaderStatement = { type : SHADER_AST_NODE_TYPES . ReturnStatement , argument : null , } ; /**
 * 将 ESTree 转换成 ShaderTree，同时完成 Scope 创建
 */ export class Transformer { public target : Target = Target . GLSL100 ; private inference : DataTypeInference = new DataTypeInference ( ) ; private context : GLSLContext ; public transform ( program , context ? ) { if ( context ) { this . context = context ; } const programNode = { type : SHADER_AST_NODE_TYPES . Program , body : [ ] , scope : [ ] , } ; this . addBuiltinVariablesToRootScope ( programNode ) ; programNode . body = flatten ( program . body . map ( ( stmt ) => this . transformStatement ( stmt , programNode , programNode ) ) ) . filter ( ( stmt ) => stmt ) ; return programNode ; } private addBuiltinVariablesToRootScope ( programNode ) { addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'globalInvocationID' , dataType : SHADER_AST_TOKEN_TYPES . Vector3Int , } ) ; addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'workGroupSize' , dataType : SHADER_AST_TOKEN_TYPES . Vector3Int , } ) ; addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'workGroupID' , dataType : SHADER_AST_TOKEN_TYPES . Vector3Int , } ) ; addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'localInvocationID' , dataType : SHADER_AST_TOKEN_TYPES . Vector3Int , } ) ; addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'localInvocationIndex' , dataType : SHADER_AST_TOKEN_TYPES . Int32 , } ) ; addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'numWorkGroups' , dataType : SHADER_AST_TOKEN_TYPES . Vector3Int , } ) ; addIdentifierToScope ( programNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : 'imageLoad' , dataType : SHADER_AST_TOKEN_TYPES . Vector4Float , } , true ) ; } private transformStatement ( stmt , scopeNode , parentNode ) { let shaderStatement = null ; if ( stmt . type === AST_NODE_TYPES . VariableDeclaration ) { shaderStatement = this . transformVariableDeclaration ( stmt , scopeNode , parentNode ) ; } else if ( stmt . type === AST_NODE_TYPES . FunctionDeclaration ) { shaderStatement = this . transformFunctionDeclaration ( stmt , parentNode ) ; } else if ( stmt . type === AST_NODE_TYPES . DoWhileStatement ) { shaderStatement = this . transformDoWhileStatement ( stmt , parentNode ) ; } else if ( stmt . type === AST_NODE_TYPES . WhileStatement ) { shaderStatement = this . transformWhileStatement ( stmt , parentNode ) ; } else if ( stmt . type === AST_NODE_TYPES . IfStatement ) { shaderStatement = this . transformIfStatement ( stmt , parentNode ) ; } else if ( stmt . type === AST_NODE_TYPES . ForStatement ) { shaderStatement = this . transformForStatement ( stmt , parentNode ) ; } else if ( stmt . type === AST_NODE_TYPES . BreakStatement ) { shaderStatement = { type : SHADER_AST_NODE_TYPES . BreakStatement , } ; } else if ( stmt . type === AST_NODE_TYPES . ContinueStatement ) { shaderStatement = { type : SHADER_AST_NODE_TYPES . ContinueStatement , } ; } else if ( stmt . type === AST_NODE_TYPES . BlockStatement ) { shaderStatement = { type : SHADER_AST_NODE_TYPES . BlockStatement , parent : parentNode , scope : [ ] , body : [ ] , } ; shaderStatement . body = flatten ( stmt . body . map ( ( s ) => this . transformStatement ( s , shaderStatement as BlockStatement , shaderStatement as BlockStatement ) ) ) ; } else if ( stmt . type === AST_NODE_TYPES . ReturnStatement ) { shaderStatement = { type : SHADER_AST_NODE_TYPES . ReturnStatement , parent : parentNode , argument : null , } ; shaderStatement . argument = this . transformExpression ( stmt . argument , shaderStatement , ( parentNode . parent ?. parent as ShaderFunctionDeclaration ) . returnType ) ; } else if ( stmt . type === AST_NODE_TYPES . ExpressionStatement ) { shaderStatement = { type : SHADER_AST_NODE_TYPES . ExpressionStatement , parent : parentNode , } ; shaderStatement . expression = this . transformExpression ( stmt . expression , shaderStatement as ExpressionStatement ) ; } else if ( stmt . type === AST_NODE_TYPES . ImportDeclaration ) { stmt . specifiers . forEach ( ( specifier ) => { if ( specifier . type === AST_NODE_TYPES . ImportSpecifier ) { const result = exportFunctions [ this . target ] [ specifier . local . name ] ; if ( result ) { addIdentifierToScope ( scopeNode , { type : SHADER_AST_NODE_TYPES . Identifier , name : specifier . local . name , dataType : result . returnType , } , true ) ; shaderStatement = { type : SHADER_AST_NODE_TYPES . ImportedFunctionStatement , content : result . content , } ; } } } ) ; } else if ( stmt . type === AST_NODE_TYPES . ClassDeclaration ) { shaderStatement = this . transformClassDeclaration ( stmt , scopeNode , parentNode ) ; } return shaderStatement ! ; } private transformExpression ( expression , parentNode , type ? ) { let retExpression = null ; if ( expression ) { if ( expression . type === AST_NODE_TYPES . Identifier ) { const existedIdentifier = getIdentifierFromScope ( parentNode , expression . name ) ; retExpression = { type : SHADER_AST_NODE_TYPES . Identifier , parent : parentNode , name : expression . name , dataType : ( existedIdentifier && existedIdentifier . dataType ) || type || SHADER_AST_TOKEN_TYPES . Float , } ; } else if ( expression . type === AST_NODE_TYPES . Literal ) { retExpression = { type : type || SHADER_AST_TOKEN_TYPES . Float , parent : parentNode , value : expression . value as number , } ; } else if ( expression . type === AST_NODE_TYPES . AssignmentExpression ) { const leftDataType = this . inference . infer ( expression . left , parentNode ) ; retExpression = { type : SHADER_AST_NODE_TYPES . AssignmentExpression , parent : parentNode , left : null , right : null , operator : expression . operator , } ; retExpression . left = this . transformExpression ( expression . left , retExpression , leftDataType ) as ShaderExpression ; retExpression . right = this . transformExpression ( expression . right , retExpression , leftDataType ) as ShaderExpression ; } else if ( expression . type === AST_NODE_TYPES . BinaryExpression || expression . type === AST_NODE_TYPES . LogicalExpression ) { const leftRightDataType = type || this . inference . infer ( expression , parentNode ) ; retExpression = { type : SHADER_AST_NODE_TYPES . BinaryExpression , parent : parentNode , left : null , right : null , operator : expression . operator , } ; retExpression . left = this . transformExpression ( expression . left , retExpression , leftRightDataType ) as ShaderExpression ; retExpression . right = this . transformExpression ( expression . right , retExpression , leftRightDataType ) as ShaderExpression ; } else if ( expression . type === AST_NODE_TYPES . ArrayExpression ) { retExpression = { type : SHADER_AST_NODE_TYPES . ArrayExpression , parent : parentNode , elements : [ ] , dataType : type , } ; const componentDataType = getComponentDataTypeFromVector ( type || SHADER_AST_TOKEN_TYPES . Float ) ; retExpression . elements = expression . elements . map ( ( e ) => this . transformExpression ( e , retExpression ! , componentDataType ) ! ) ; } else if ( expression . type === AST_NODE_TYPES . UpdateExpression ) { retExpression = { type : SHADER_AST_NODE_TYPES . UpdateExpression , parent : parentNode , operator : expression . operator , argument : null , prefix : expression . prefix , } ; retExpression . argument = this . transformExpression ( expression . argument , retExpression ) ; } else if ( expression . type === AST_NODE_TYPES . CallExpression ) { retExpression = { type : SHADER_AST_NODE_TYPES . CallExpression , parent : parentNode , callee : null , arguments : [ ] , } ; retExpression . callee = this . transformExpression ( expression . callee , retExpression ) ; retExpression . arguments = expression . arguments . map ( ( arg ) => this . transformExpression ( arg , retExpression ) ) ; } else if ( expression . type === AST_NODE_TYPES . MemberExpression ) { if ( expression . object . type === AST_NODE_TYPES . Identifier && expression . object . name === 'this' ) { return this . transformExpression ( expression . property , parentNode ) ; } retExpression = { type : SHADER_AST_NODE_TYPES . MemberExpression , parent : parentNode , object : null , property : null , computed : expression . computed , } ; retExpression . object = this . transformExpression ( expression . object , retExpression ) ; retExpression . property = this . transformExpression ( expression . property , retExpression ) ; return retExpression ; } else if ( expression . type === AST_NODE_TYPES . ConditionalExpression ) { retExpression = { type : SHADER_AST_NODE_TYPES . ConditionalExpression , parent : parentNode , test : null , alternate : null , consequent : null , } ; retExpression . test = this . transformExpression ( expression . test , retExpression ) ; retExpression . alternate = this . transformExpression ( expression . alternate , retExpression ) ; retExpression . consequent = this . transformExpression ( expression . consequent , retExpression ) ; } else if ( expression . type === AST_NODE_TYPES . UnaryExpression ) { retExpression = { type : SHADER_AST_NODE_TYPES . UnaryExpression , parent : parentNode , operator : expression . operator , prefix : expression . prefix , } ; retExpression . argument = this . transformExpression ( expression . argument , retExpression ) ; } } return retExpression ; } private transformIdentifier ( id , parent ? ) { const { name , typeAnnotation } = id ; return { type : SHADER_AST_NODE_TYPES . Identifier , dataType : ( typeAnnotation && dataTypeMap [ typeAnnotation ] ) || SHADER_AST_TOKEN_TYPES . Float , name , parent , } ; } private transformDoWhileStatement ( stmt , parentNode ) { const doWhileStatement = { type : SHADER_AST_NODE_TYPES . DoWhileStatement , parent : parentNode , test : null , body : null , } ; doWhileStatement . test = this . transformExpression ( stmt . test , doWhileStatement ) ; doWhileStatement . body = this . transformStatement ( stmt . body , doWhileStatement , doWhileStatement ) ; return doWhileStatement ; } private transformWhileStatement ( stmt , parentNode ) { const whileStatement = { type : SHADER_AST_NODE_TYPES . WhileStatement , parent : parentNode , test : null , body : null , } ; whileStatement . test = this . transformExpression ( stmt . test , whileStatement ) ; whileStatement . body = this . transformStatement ( stmt . body , whileStatement , whileStatement ) ; return whileStatement ; } private transformForStatement ( stmt , parentNode ) { const forStatement = { type : SHADER_AST_NODE_TYPES . ForStatement , parent : parentNode , init : null , test : null , update : null , body : emptyShaderStatement , } ; if ( stmt . init ?. type === AST_NODE_TYPES . VariableDeclaration ) { forStatement . init = this . transformStatement ( stmt . init as VariableDeclaration , parentNode , forStatement ) as ShaderVariableDeclaration ; } else if ( stmt . init ?. type === AST_NODE_TYPES . AssignmentExpression ) { forStatement . init = this . transformExpression ( stmt . init , forStatement ) ; } forStatement . test = this . transformExpression ( stmt . test , forStatement , SHADER_AST_TOKEN_TYPES . Int32 ) ; forStatement . update = this . transformExpression ( stmt . update , forStatement ) ; forStatement . body = this . transformStatement ( stmt . body , parentNode , forStatement ) as ShaderStatement ; return forStatement ; } private transformVariableDeclaration ( stmt , scopeNode , parentNode ) { const shaderStatement = { type : SHADER_AST_NODE_TYPES . VariableDeclaration , parent : parentNode , declarations : stmt . declarations . map ( ( d ) => { const id = this . transformIdentifier ( d . id as Identifier ) ; if ( d . init && d . init . type === AST_NODE_TYPES . Identifier ) { const scope = getScopeByIdentifierName ( scopeNode , d . init . name ) ; if ( scope ) { scope . alias . push ( d . init . name ) ; } } addIdentifierToScope ( scopeNode , id , false ) ; if ( ! d . id . typeAnnotation ) { const inferredDataType = this . inference . infer ( d . init , scopeNode ) ; id . dataType = inferredDataType ; } let storageClass = STORAGE_CLASS . Private ; /**
                 * 全大写常量转译成 GLSL 常量
                 * const AAA = 100;
                 * ->
                 * #define AAA 100
                 */ if ( ! scopeNode . parent && stmt . kind === 'const' && id . name === id . name . toUpperCase ( ) ) { let definedValue ; let runtime = false ; storageClass = STORAGE_CLASS . UniformConstant ; if ( d ?. init ?. type === AST_NODE_TYPES . Literal && d . init . value !== null && isFinite ( Number ( d . init . value ) ) ) { definedValue = d . init . value ; } else { runtime = true ; } if ( this . context ) { this . context . defines . push ( { name : id . name , type : id . dataType , value : definedValue , runtime , } ) ; } } const declarator = { type : SHADER_AST_NODE_TYPES . VariableDeclarator , id , parent : shaderStatement as ShaderVariableDeclaration , storageClass , init : d . init && this . transformExpression ( d . init , parentNode , id . dataType ) , } ; id . parent = declarator ; return declarator ; } ) , } ; return shaderStatement ; } private transformFunctionDeclaration ( stmt , parentNode ) { const id = this . transformIdentifier ( stmt . id ! ) ; const returnType = ( stmt . returnType && dataTypeMap [ stmt . returnType ] ) || SHADER_AST_TOKEN_TYPES . Float ; const shaderStatement = { type : SHADER_AST_NODE_TYPES . FunctionDeclaration , parent : parentNode , id , body : { type : SHADER_AST_NODE_TYPES . BlockStatement , body : [ ] , scope : [ ] , } , params : [ ] , returnType , } ; id . parent = shaderStatement ; shaderStatement . body . parent = shaderStatement ; addIdentifierToScope ( parentNode , { ... id , dataType : returnType , } , true ) ; shaderStatement . params = stmt . params . map ( ( param ) => { const paramId = this . transformIdentifier ( param as Identifier , shaderStatement as ShaderFunctionDeclaration ) ; addIdentifierToScope ( shaderStatement . body , paramId , false ) ; return paramId ; } ) ; shaderStatement . body . body = flatten ( stmt . body . body . map ( ( b ) => this . transformStatement ( b , shaderStatement , shaderStatement . body ) ) ) ; return shaderStatement ; } private transformIfStatement ( stmt , parentNode ) { const shaderStatement = { type : SHADER_AST_NODE_TYPES . IfStatement , parent : parentNode , test : null , consequent : null , alternate : null , } ; shaderStatement . test = this . transformExpression ( stmt . test , shaderStatement ) ! ; shaderStatement . consequent = this . transformStatement ( stmt . consequent , shaderStatement , shaderStatement ) as ShaderStatement ; if ( stmt . alternate ) { shaderStatement . alternate = this . transformStatement ( stmt . alternate , shaderStatement , shaderStatement ) as ShaderStatement ; } return shaderStatement ; } private transformClassDeclaration ( stmt , scopeNode , parentNode ) { const shaderStatement = [ ] ; stmt . decorators ?. forEach ( ( decorator ) => { if ( decorator . expression . type === AST_NODE_TYPES . CallExpression && decorator . expression . callee . type === AST_NODE_TYPES . Identifier && decorator . expression . callee . name === 'numthreads' ) { const threadGroupSize = decorator . expression . arguments . slice ( 0 , 3 ) . map ( ( a ) => ( a . type === AST_NODE_TYPES . Literal && Number ( a . value ) ) || 1 ) as [ number , number , number ] ; ( shaderStatement as ShaderStatement [ ] ) . push ( { type : SHADER_AST_NODE_TYPES . NumThreadStatement , parent : parentNode , threadGroupSize , } ) ; this . context . threadGroupSize = threadGroupSize ; } } ) ; if ( stmt . body . type === AST_NODE_TYPES . ClassBody ) { /**
             * uniforms
             * eg. prop1: vec3[]
             */ const classProperties = stmt . body . body . filter ( ( e ) => e . type === AST_NODE_TYPES . ClassProperty ) as ClassProperty [ ] ; if ( classProperties . length ) { const declaration = { type : SHADER_AST_NODE_TYPES . VariableDeclaration , declarations : [ ] , } ; classProperties . forEach ( ( property ) => { if ( property . key . type === AST_NODE_TYPES . Identifier ) { const id = this . transformIdentifier ( property . key ) ; addIdentifierToScope ( scopeNode , id , false ) ; let storageClass = STORAGE_CLASS . Uniform ; if ( id . dataType === SHADER_AST_TOKEN_TYPES . FloatArray || id . dataType === SHADER_AST_TOKEN_TYPES . Vector4FloatArray ) { storageClass = STORAGE_CLASS . StorageBuffer ; } if ( property . decorators && this . context ) { this . analyzeDecorators ( property . key . name , id . dataType , storageClass , property . decorators ) ; } const declarator = { type : SHADER_AST_NODE_TYPES . VariableDeclarator , id , init : null , storageClass , } ; declaration . declarations . push ( declarator as ShaderVariableDeclarator ) ; id . parent = declarator ; } } ) ; shaderStatement . push ( declaration ) ; } const methodDefinitions = stmt . body . body . filter ( ( e ) => e . type === AST_NODE_TYPES . MethodDefinition ) as MethodDefinition [ ] ; if ( methodDefinitions . length ) { methodDefinitions . forEach ( ( method ) => { if ( method . value . type === AST_NODE_TYPES . FunctionExpression ) { if ( method . decorators && method . decorators . length === 1 ) { method . value . returnType = 'void' ; method . value . params = [ ] ; ( method . value . id as Identifier ) . name = 'main' ; } const functionDeclaration = { ... method . value , type : AST_NODE_TYPES . FunctionDeclaration , body : method . value . body ! , } ; ( shaderStatement as ShaderStatement [ ] ) . push ( this . transformFunctionDeclaration ( functionDeclaration , parentNode ) ) ; } } ) ; } } return shaderStatement ; } private analyzeDecorators ( propertyName , propertyType , storageClass , decorators ) { const analyzeDecorators = decorators . map ( ( d ) => this . extractDecorator ( d ) ) ; if ( analyzeDecorators . find ( ( d ) => d . name === PropertyDecorator . Out ) && analyzeDecorators . find ( ( d ) => d . name === PropertyDecorator . In ) ) { this . context . needPingpong = true ; } analyzeDecorators . forEach ( ( { name , params } ) => { if ( name === PropertyDecorator . Shared ) { this . context . globalDeclarations . push ( { name : propertyName , type : propertyType , shared : true , value : `${ params [ 0 ] }` , } ) ; } else if ( name === PropertyDecorator . Out || name === PropertyDecorator . In ) { let existed = this . context . uniforms . find ( ( u ) => u . name === propertyName ) ; if ( ! existed ) { existed = { name : propertyName , type : propertyType , storageClass , readonly : ! ! ! analyzeDecorators . find ( ( d ) => d . name === PropertyDecorator . Out ) , writeonly : ! ! ! analyzeDecorators . find ( ( d ) => d . name === PropertyDecorator . In ) , size : [ Number ( params [ 0 ] ) || 1 , Number ( params [ 1 ] ) || 1 ] as [ number , number ] , } ; this . context . uniforms . push ( existed ) ; } if ( name === PropertyDecorator . Out ) { this . context . output . name = propertyName ; this . context . output . size = existed . size ; this . context . output . length = existed . size ! [ 0 ] * existed . size ! [ 1 ] ; } } } ) ; } private extractDecorator ( decorator ) { let name = '' ; let params = [ ] ; if ( decorator . expression . type === AST_NODE_TYPES . Identifier ) { name = decorator . expression . name ; } else if ( decorator . expression . type === AST_NODE_TYPES . CallExpression && decorator . expression . callee . type === AST_NODE_TYPES . Identifier ) { name = decorator . expression . callee . name ; params = decorator . expression . arguments . map ( ( a ) => ( a . type === AST_NODE_TYPES . Literal && Number ( a . value ) ) || 0 ) ; } return { name , params , } ; } }