import { AST_NODE_TYPES , AST_TOKEN_TYPES , STORAGE_CLASS , } from '../ast/glsl-ast-node-types' ; import { ArrayExpression , AssignmentExpression , BaseNode , BinaryExpression , BlockStatement , CallExpression , ConditionalExpression , DataType , DoWhileStatement , Expression , ExpressionStatement , ForStatement , FunctionDeclaration , FunctionExpression , Identifier , IfStatement , MemberExpression , NumThreadStatement , Program , ReturnStatement , Scalar , Scope , Statement , UnaryExpression , UpdateExpression , VariableDeclaration , VariableDeclarator , WhileStatement , } from '../ast/glsl-tree' ; import { DefineValuePlaceholder , GLSLContext } from '../Compiler' ; import { dataTypeMap as transformDataTypeMap } from '../Transformer' ; import { builtinFunctions , componentWiseFunctions , importFunctions , typeCastFunctions , } from '../utils/builtin-functions' ; import { compareDataTypePriority , getByteCount , getComponentDataTypeFromVector , getComponentSize , isArray , isLeft , isMatrix , isScalar , isVector , } from '../utils/data-type' ; import { getIdentifierFromScope , traverseUpwards } from '../utils/node' ; import { ICodeGenerator , Target } from './ICodeGenerator' ; const dataTypeMap = { [ AST_TOKEN_TYPES . Void ] : 'void' , [ AST_TOKEN_TYPES . Boolean ] : 'bool' , [ AST_TOKEN_TYPES . Int32 ] : 'i32' , [ AST_TOKEN_TYPES . Uint32 ] : 'u32' , [ AST_TOKEN_TYPES . Float ] : 'f32' , [ AST_TOKEN_TYPES . Vector2Float ] : 'vec2<f32>' , [ AST_TOKEN_TYPES . Vector3Float ] : 'vec3<f32>' , [ AST_TOKEN_TYPES . Vector4Float ] : 'vec4<f32>' , [ AST_TOKEN_TYPES . Vector2Int ] : 'vec2<i32>' , [ AST_TOKEN_TYPES . Vector3Int ] : 'vec3<i32>' , [ AST_TOKEN_TYPES . Vector4Int ] : 'vec4<i32>' , [ AST_TOKEN_TYPES . Vector2Uint ] : 'vec2<u32>' , [ AST_TOKEN_TYPES . Vector3Uint ] : 'vec3<u32>' , [ AST_TOKEN_TYPES . Vector4Uint ] : 'vec4<u32>' , [ AST_TOKEN_TYPES . Vector2Boolean ] : 'vec2<bool>' , [ AST_TOKEN_TYPES . Vector3Boolean ] : 'vec3<bool>' , [ AST_TOKEN_TYPES . Vector4Boolean ] : 'vec4<bool>' , [ AST_TOKEN_TYPES . Matrix3x3Float ] : 'mat3x3<f32>' , [ AST_TOKEN_TYPES . Matrix4x4Float ] : 'mat4x4<f32>' , [ AST_TOKEN_TYPES . FloatArray ] : 'array<f32>' , [ AST_TOKEN_TYPES . Vector4FloatArray ] : 'array<vec4<f32>>' , } ; export class CodeGeneratorWGSL implements ICodeGenerator { public static GWEBGPU_UNIFORM_PARAMS = 'gWebGPUUniformParams' ; public static GWEBGPU_BUFFER = 'gWebGPUBuffer' ; private context : GLSLContext ; private bufferBindingIndex = - 1 ; public clear ( ) { this . bufferBindingIndex = - 1 ; } public generate ( program , context ? ) { if ( context ) { this . context = context ; } const body = program . body . map ( ( stmt ) => this . generateStatement ( stmt ) ) . join ( '\n' ) ; return `import "GLSL.std.450" as std;
${ importFunctions [ Target . WGSL ] }

# var gWebGPUDebug : bool = false;
# var gWebGPUDebugOutput : vec4<f32> = vec4<f32>(0.0);

[[builtin global_invocation_id]] var<in> globalInvocationID : vec3<u32>;
# [[builtin work_group_size]] var<in> workGroupSize : vec3<u32>;
# [[builtin work_group_id]] var<in> workGroupID : vec3<u32>;
[[builtin local_invocation_id]] var<in> localInvocationID : vec3<u32>;
# [[builtin num_work_groups]] var<in> numWorkGroups : vec3<u32>;
[[builtin local_invocation_idx]] var<in> localInvocationIndex : u32;

${ this . generateUniforms ( ) }
${ this . generateBuffers ( ) }
${ this . generateGlobalVariableDeclarations ( ) }

${ body }

entry_point compute as "main" = main;
` ; } public generateStatement ( stmt ) { if ( stmt . type === AST_NODE_TYPES . VariableDeclaration ) { return this . generateVariableDeclaration ( stmt ) ; } else if ( stmt . type === AST_NODE_TYPES . FunctionDeclaration ) { return this . generateFunctionExpression ( stmt ) ; } else if ( stmt . type === AST_NODE_TYPES . BlockStatement ) { return this . generateBlockStatement ( stmt ) ; } else if ( stmt . type === AST_NODE_TYPES . NumThreadStatement ) { return this . generateNumThreadStatement ( stmt ) ; } else { console . warn ( `[WGSL compiler]: unknown statement: ${ stmt . type }` ) ; } return '' ; } public generateNumThreadStatement ( stmt ) { return '' ; } public generateVariableDeclaration ( stmt ) { return stmt . declarations . map ( ( declarator ) => { if ( declarator . storageClass === STORAGE_CLASS . UniformConstant ) { } else if ( declarator . storageClass === STORAGE_CLASS . Uniform || declarator . storageClass === STORAGE_CLASS . StorageBuffer ) { } else if ( declarator . storageClass === STORAGE_CLASS . Private ) { if ( declarator . init ) { return `var ${ declarator . id . name } : ${ this . generateDataType ( declarator . id . dataType ) } = ${ this . generateExpression ( declarator . init , declarator . id . dataType ) };` ; } else { if ( stmt . parent && ( stmt . parent as Program ) . type === AST_NODE_TYPES . Program ) { return '' ; } return `var ${ declarator . id . name } : ${ this . generateDataType ( declarator . id . dataType ) };` ; } } return '' ; } ) . join ( '\n' ) ; } public generateBlockStatement ( stmt ) { return stmt . body . map ( ( s ) => { if ( s . type === AST_NODE_TYPES . VariableDeclaration ) { return this . generateVariableDeclaration ( s ) ; } else if ( s . type === AST_NODE_TYPES . ExpressionStatement ) { return this . generateExpression ( s . expression ) + ';' ; } else if ( s . type === AST_NODE_TYPES . ReturnStatement ) { return this . generateReturnStatement ( s ) ; } else if ( s . type === AST_NODE_TYPES . IfStatement ) { return this . generateIfStatement ( s ) ; } else if ( s . type === AST_NODE_TYPES . ForStatement ) { return this . generateForStatement ( s ) ; } else if ( s . type === AST_NODE_TYPES . BreakStatement ) { return 'break;' ; } else if ( s . type === AST_NODE_TYPES . ContinueStatement ) { return 'continue;' ; } else if ( s . type === AST_NODE_TYPES . WhileStatement ) { return this . generateWhileStatement ( s ) ; } else if ( s . type === AST_NODE_TYPES . DoWhileStatement ) { return this . generateDoWhileStatement ( s ) ; } else { console . warn ( `[WGSL compiler]: unknown statement: ${ stmt }` ) ; } return '' ; } ) . join ( '\n' ) ; } public generateReturnStatement ( stmt ) { const returnType = traverseUpwards < DataType > ( stmt , ( currentNode ) => { if ( ( currentNode as FunctionDeclaration ) . type === AST_NODE_TYPES . FunctionDeclaration && ( currentNode as FunctionDeclaration ) . returnType ) { return ( currentNode as FunctionDeclaration ) . returnType ; } } ) ; return `return ${ ( stmt . argument && this . generateExpression ( stmt . argument , returnType ) ) || '' };` ; } public generateWhileStatement ( stmt ) { return `loop {
  if (${ this . generateExpression ( stmt . test ) }) { break; }
${ this . generateStatement ( stmt . body ) }
}` ; } public generateDoWhileStatement ( stmt ) { return `loop {
${ this . generateStatement ( stmt . body ) }
  if (${ this . generateExpression ( stmt . test ) }) { break; }
}` ; } public generateForStatement ( node ) { let init = '' ; if ( node . init ?. type === AST_NODE_TYPES . VariableDeclaration ) { node . init . declarations . forEach ( ( d ) => { d . id . dataType = AST_TOKEN_TYPES . Int32 ; } ) ; init = this . generateVariableDeclaration ( node . init ) ; } else if ( node . init ?. type === AST_NODE_TYPES . AssignmentExpression ) { init = this . generateExpression ( node . init ) ; } return `for (${ init } ${ node . test && this . generateExpression ( node . test ) }; ${ node . update && this . generateExpression ( node . update ) }) {${ this . generateBlockStatement ( node . body as BlockStatement ) }}` ; } public generateIfStatement ( node ) { let consequent = '' ; if ( node . consequent . type === AST_NODE_TYPES . ExpressionStatement ) { consequent = this . generateExpression ( node . consequent . expression ) ; } else if ( node . consequent . type === AST_NODE_TYPES . BlockStatement ) { consequent = this . generateBlockStatement ( node . consequent ) ; } else if ( node . consequent . type === AST_NODE_TYPES . ReturnStatement ) { consequent = this . generateReturnStatement ( node . consequent ) ; } else if ( node . consequent . type === AST_NODE_TYPES . BreakStatement ) { consequent = 'break;' ; } let alternate = '' ; if ( node . alternate ) { if ( node . alternate . type === AST_NODE_TYPES . ExpressionStatement ) { alternate = this . generateExpression ( node . alternate . expression ) ; } else if ( node . alternate . type === AST_NODE_TYPES . BlockStatement ) { alternate = this . generateBlockStatement ( node . alternate ) ; } else if ( node . alternate . type === AST_NODE_TYPES . ReturnStatement ) { alternate = this . generateReturnStatement ( node . alternate ) ; } else if ( node . alternate . type === AST_NODE_TYPES . BreakStatement ) { alternate = 'break;' ; } } return `if (${ this . generateExpression ( node . test ) }) {${ consequent }}${ alternate ? `else {${ alternate }}` : '' }` ; } public generateFunctionExpression ( stmt ) { if ( stmt . body ) { const attributes = '' ; let append = '' ; if ( this . context && stmt . id ?. name === 'main' ) { const [ x , y , z ] = this . context . threadGroupSize ; } const lastStatment = stmt . body . body [ stmt . body . body . length - 1 ] ; if ( lastStatment . type !== AST_NODE_TYPES . ReturnStatement ) { append = '\nreturn;' ; } return `${ attributes }fn ${ stmt . id ?. name }(${ ( stmt . params || [ ] ) . map ( ( p ) => p . type === AST_NODE_TYPES . Identifier && `${ p . name } : ${ this . generateDataType ( p . dataType ) }` ) . join ( ', ' ) }) -> ${ this . generateDataType ( stmt . returnType || AST_TOKEN_TYPES . Void ) } {${ this . generateBlockStatement ( stmt . body ) }${ append }}` ; } return '' ; } public generateExpression ( expression , dataType ? ) { if ( isScalar ( expression ?. type ) ) { return this . generateScalar ( dataType || expression ?. type , expression . value ) ; } else if ( isVector ( expression ?. type ) ) { return this . generateVector ( dataType || expression ?. type , expression . value ) ; } else if ( isMatrix ( expression ?. type ) ) { return this . generateMatrix ( dataType || expression ?. type , expression . value ) ; } else if ( expression ?. type === AST_NODE_TYPES . ArrayExpression ) { return this . generateArrayExpression ( expression , dataType ) ; } else if ( expression ?. type === AST_NODE_TYPES . Identifier ) { const id = getIdentifierFromScope ( expression , expression . name ) ; const storageClass = ( id ?. parent as VariableDeclarator ) ?. storageClass ; if ( id && storageClass ) { if ( storageClass === STORAGE_CLASS . StorageBuffer ) { const bufferIndex = this . context . uniforms . filter ( ( u ) => u . storageClass === STORAGE_CLASS . StorageBuffer ) . findIndex ( ( u ) => u . name === id . name ) ; if ( bufferIndex > - 1 ) { return `${ CodeGeneratorWGSL . GWEBGPU_BUFFER }${ bufferIndex }.${ expression . name }` ; } } else if ( storageClass === STORAGE_CLASS . Uniform ) { const uniform = this . context . uniforms . filter ( ( u ) => u . storageClass === STORAGE_CLASS . Uniform ) . find ( ( u ) => u . name === id . name ) ; if ( uniform ) { return `${ CodeGeneratorWGSL . GWEBGPU_UNIFORM_PARAMS }.${ expression . name }` ; } } } if ( componentWiseFunctions . indexOf ( expression . name ) > - 1 ) { return `std::${ expression . name }` ; } const define = this . context ?. defines . find ( ( d ) => d . name === expression . name ) ; if ( define ) { const placeHolder = `${ DefineValuePlaceholder }${ define ?. name }` ; return define . runtime ? placeHolder : `${ define . value }` ; } if ( typeCastFunctions . indexOf ( expression . name ) > - 1 && ( expression ?. parent as CallExpression ) . type === AST_NODE_TYPES . CallExpression ) { return this . generateDataType ( transformDataTypeMap [ expression . name ] ) ; } return this . castingDataType ( dataType , expression . dataType , expression . name ) ; } else if ( expression ?. type === AST_NODE_TYPES . BinaryExpression ) { return this . generateBinaryExpression ( expression , dataType ) ; } else if ( expression ?. type === AST_NODE_TYPES . CallExpression ) { return this . generateCallExpression ( expression , dataType ) ; } else if ( expression ?. type === AST_NODE_TYPES . AssignmentExpression ) { return this . generateAssignmentExpression ( expression , dataType ) ; } else if ( expression ?. type === AST_NODE_TYPES . UpdateExpression ) { return this . generateUpdateExpression ( expression ) ; } else if ( expression ?. type === AST_NODE_TYPES . MemberExpression ) { return this . generateMemberExpression ( expression ) ; } else if ( expression ?. type === AST_NODE_TYPES . ConditionalExpression ) { return this . generateConditionalExpression ( expression ) ; } else if ( expression ?. type === AST_NODE_TYPES . UnaryExpression ) { return this . generateUnaryExpression ( expression , dataType ) ; } else { console . warn ( `[WGSL compiler]: unknown expression: ${ expression }` , dataType ) ; } return '' ; } public generateUnaryExpression ( expression , dataType ? ) { return `${ expression . operator }${ this . generateExpression ( expression . argument , dataType ) }` ; } public generateConditionalExpression ( expression ) { return `select(${ this . generateExpression ( expression . consequent ) }, ${ this . generateExpression ( expression . alternate ) }, ${ this . generateExpression ( expression . test ) })` ; } public generateMemberExpression ( expression ) { const objectStr = this . generateExpression ( expression . object ) ; if ( expression . property . type === AST_NODE_TYPES . Identifier ) { if ( ! expression . computed ) { return `${ objectStr }.${ ( expression . property as Identifier ) . name }` ; } else { return `${ objectStr }[${ ( expression . property as Identifier ) . name }]` ; } } else if ( ( expression . property . type === AST_TOKEN_TYPES . Float || expression . property . type === AST_TOKEN_TYPES . Int32 || expression . property . type === AST_TOKEN_TYPES . Uint32 ) && isFinite ( Number ( expression . property . value ) ) ) { const index = Number ( expression . property . value ) ; if ( expression . object . type === AST_NODE_TYPES . Identifier && isArray ( expression . object . dataType ) ) { return `${ objectStr }[${ index }]` ; } else { let swizzlingComponent = 'x' ; switch ( index ) { case 0 : swizzlingComponent = 'x' ; break ; case 1 : swizzlingComponent = 'y' ; break ; case 2 : swizzlingComponent = 'z' ; break ; case 3 : swizzlingComponent = 'w' ; break ; } return `${ objectStr }.${ swizzlingComponent }` ; } } else { return `${ objectStr }[${ this . generateExpression ( expression . property ) }]` ; } } public generateUpdateExpression ( expression ) { const { name } = expression . argument as Identifier ; if ( expression . operator === '++' ) { return `${ name } = ${ name } + 1` ; } else if ( expression . operator === '--' ) { return `${ name } = ${ name } - 1` ; } return expression . prefix ? `${ expression . operator }${ name }` : `${ name }${ expression . operator }` ; } public generateAssignmentExpression ( expression , dataType ? ) { const left = this . generateExpression ( expression . left , dataType ) ; const right = this . generateExpression ( expression . right , dataType ) ; if ( expression . operator . length > 1 && expression . operator . endsWith ( '=' ) ) { const realOperator = expression . operator . substr ( 0 , expression . operator . length - 1 ) ; return `${ left } = ${ left } ${ realOperator } ${ right }` ; } return `${ left } ${ expression . operator } ${ right }` ; } public generateArrayExpression ( expression , dataType ? ) { const dt = expression . dataType || dataType || AST_TOKEN_TYPES . Vector4Float ; return `${ this . generateDataType ( dt ) }(${ expression . elements . map ( ( e ) => this . generateExpression ( e , getComponentDataTypeFromVector ( dt ) ) ) . join ( ', ' ) })` ; } public generateBinaryExpression ( expression , dataType ? ) { const needBracket = expression . parent ?. type === AST_NODE_TYPES . BinaryExpression ; return `${ needBracket ? '(' : '' }${ this . generateExpression ( expression . left , dataType ) } ${ expression . operator } ${ this . generateExpression ( expression . right , dataType ) }${ needBracket ? ')' : '' }` ; } public generateCallExpression ( expression , dataType ? ) { const isComponentWise = expression . callee . type === AST_NODE_TYPES . Identifier && componentWiseFunctions . indexOf ( expression . callee . name ) > - 1 ; let params = [ ] ; if ( expression . callee . type === AST_NODE_TYPES . Identifier ) { const functionId = getIdentifierFromScope ( expression , expression . callee . name , true ) ; if ( functionId && ( functionId . parent as FunctionDeclaration ) . type === AST_NODE_TYPES . FunctionDeclaration ) { params = ( functionId . parent as FunctionDeclaration ) . params ; } } return `${ this . generateExpression ( expression . callee ) }(${ expression . arguments . map ( ( e , i ) => { return this . generateExpression ( e , isComponentWise ? dataType : params [ i ] && params [ i ] . dataType ) ; } ) . join ( ', ' ) })` ; } public generateDataType ( dt ) { return dataTypeMap [ dt ] ; } public generateScalar ( dt , scalar ) { if ( dt === AST_TOKEN_TYPES . Boolean ) { return `${ ! ! scalar }` ; } else if ( dt === AST_TOKEN_TYPES . Float ) { return this . wrapFloat ( `${ Number ( scalar ) }` ) ; } else if ( dt === AST_TOKEN_TYPES . Int32 || dt === AST_TOKEN_TYPES . Uint32 ) { return parseInt ( `${ Number ( scalar ) }` , 10 ) . toString ( ) ; } return `${ scalar }` ; } public generateVector ( dt , vector ) { const componentDataType = getComponentDataTypeFromVector ( dt ) ; return this . castingDataType ( dt , componentDataType , this . generateScalar ( componentDataType , vector ) ) ; } public generateMatrix ( dt , vector ) { const componentDataType = getComponentDataTypeFromVector ( dt ) ; return this . castingDataType ( dt , componentDataType , this . generateScalar ( componentDataType , vector ) ) ; } private wrapFloat ( float ) { return float . indexOf ( '.' ) === - 1 ? `${ float }.0` : float ; } private castingDataType ( dt1 , dt2 , content ) { if ( dt1 && dt1 !== dt2 ) { const castDataType = compareDataTypePriority ( dt1 , dt2 ) ; if ( castDataType !== dt2 ) { return `${ this . generateDataType ( castDataType ) }(${ content })` ; } } return content ; } private generateSwizzling ( dt ) { const size = getComponentSize ( dt ) ; if ( size === 3 ) { return '.rgb' ; } else if ( size === 2 ) { return '.rg' ; } else if ( size === 4 ) { return '.rgba' ; } else { return '.r' ; } } private generateGlobalVariableDeclarations ( ) { if ( ! this . context ) { return '' ; } return `${ this . context . globalDeclarations . map ( ( gd ) => `${ ( gd . shared && 'shared' ) || '' } ${ this . generateDataType ( getComponentDataTypeFromVector ( gd . type ) ) . replace ( '[]' , '' ) } ${ gd . name }[${ gd . value }];` ) . join ( '\n' ) }
` ; } private generateUniforms ( ) { if ( ! this . context ) { return '' ; } let offset = 0 ; const uniformDeclarations = this . context . uniforms . map ( ( uniform ) => { if ( uniform . storageClass === STORAGE_CLASS . Uniform ) { const byteCount = getByteCount ( uniform . type ) ; const ret = `[[offset ${ offset }]] ${ uniform . name } : ${ this . generateDataType ( uniform . type ) };` ; offset += byteCount ; return ret ; } return '' ; } ) . join ( '\n  ' ) . trim ( ) ; return ( uniformDeclarations && `type GWebGPUParams = [[block]] struct {
  ${ uniformDeclarations }
};
[[binding ${ ++ this . bufferBindingIndex }, set 0]] var<uniform> ${ CodeGeneratorWGSL . GWEBGPU_UNIFORM_PARAMS } : GWebGPUParams;` ) ; } private generateBuffers ( ) { if ( ! this . context ) { return '' ; } let bufferIndex = - 1 ; return this . context . uniforms . map ( ( u ) => { if ( u . storageClass === STORAGE_CLASS . StorageBuffer ) { bufferIndex ++ ; ++ this . bufferBindingIndex ; const componentDataType = getComponentDataTypeFromVector ( u . type ) ; return `type GWebGPUBuffer${ bufferIndex } = [[block]] struct {
  [[offset 0]] ${ u . name } : [[stride ${ getByteCount ( componentDataType ) }]] array<${ this . generateDataType ( componentDataType ) }>;
};
[[binding ${ this . bufferBindingIndex }, set 0]] var<storage_buffer> ${ CodeGeneratorWGSL . GWEBGPU_BUFFER }${ bufferIndex } : GWebGPUBuffer${ bufferIndex };` ; } return '' ; } ) . filter ( ( line ) => line ) . join ( '\n' ) . trim ( ) ; } }