Skip to content

Commit 5509beb

Browse files
iwoplazaiwoplaza
committed
feat(@typegpu/gl): Generate GLSL function signatures
1 parent 8a474d6 commit 5509beb

File tree

9 files changed

+507
-234
lines changed

9 files changed

+507
-234
lines changed

‎packages/typegpu-gl/src/glOptions.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import glslGenerator from './glslGenerator.ts';
2+
3+
export function GLOptions() {
4+
return {
5+
unstable_shaderGenerator: glslGenerator,
6+
};
7+
}

‎packages/typegpu-gl/src/glslGenerator.ts‎

Lines changed: 107 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
1-
import { d, ShaderGenerator, WgslGenerator } from 'typegpu';
1+
import { NodeTypeCatalog as NODE } from 'tinyest';
2+
import type { Return } from 'tinyest';
3+
import tgpu, { d, ShaderGenerator, WgslGenerator } from 'typegpu';
4+
5+
type ResolutionCtx = ShaderGenerator.ResolutionCtx;
6+
7+
const UnknownData: typeof ShaderGenerator.UnknownData = ShaderGenerator.UnknownData;
28

39
// ----------
410
// WGSL → GLSL type name mapping
511
// ----------
612

713
const WGSL_TO_GLSL_TYPE: Record<string, string> = {
14+
void: 'void',
815
f32: 'float',
916
u32: 'uint',
1017
i32: 'int',
@@ -37,13 +44,38 @@ export function translateWgslTypeToGlsl(wgslType: string): string {
3744
return WGSL_TO_GLSL_TYPE[wgslType] ?? wgslType;
3845
}
3946

47+
/**
48+
* Resolves a struct and adds its declaration to the resolution context.
49+
* @param ctx - The resolution context.
50+
* @param struct - The struct to resolve.
51+
*
52+
* @returns The resolved struct name.
53+
*/
54+
function resolveStruct(ctx: ResolutionCtx, struct: d.WgslStruct) {
55+
const id = ctx.getUniqueName(struct);
56+
57+
ctx.addDeclaration(`\
58+
struct ${id} {
59+
${Object.entries(struct.propTypes)
60+
.map(([prop, type]) => ` ${ctx.resolve(type).value} ${prop};\n`)
61+
.join('')}\
62+
};`);
63+
64+
return id;
65+
}
66+
67+
const versionMacro = tgpu['~unstable'].declare('#version 300 es');
68+
const gl_PositionSnippet = tgpu['~unstable'].rawCodeSnippet('gl_Position', d.vec4f, 'private');
69+
4070
/**
4171
* A GLSL ES 3.0 shader generator that extends WgslGenerator.
4272
* Overrides `dataType` to emit GLSL type names instead of WGSL ones,
4373
* and overrides variable declaration emission to use `type name = rhs` syntax.
4474
*/
4575
export class GlslGenerator extends WgslGenerator {
46-
public override typeAnnotation(data: d.BaseData): string {
76+
#functionType: ShaderGenerator.TgpuShaderStage | 'normal' | undefined;
77+
78+
override typeAnnotation(data: d.BaseData): string {
4779
// For WGSL identity types (scalars, vectors, common matrices), map to GLSL directly.
4880
if (!d.isLooseData(data)) {
4981
const glslName = WGSL_TO_GLSL_TYPE[data.type];
@@ -52,20 +84,88 @@ export class GlslGenerator extends WgslGenerator {
5284
}
5385
}
5486

87+
if (d.isWgslStruct(data)) {
88+
return resolveStruct(this.ctx, data);
89+
}
90+
5591
// For all other types (structs, arrays, etc.) delegate to WGSL resolution.
5692
return super.typeAnnotation(data);
5793
}
5894

59-
protected override emitVarDecl(
60-
pre: string,
95+
override _emitVarDecl(
6196
_keyword: 'var' | 'let' | 'const',
6297
name: string,
6398
dataType: d.BaseData | ShaderGenerator.UnknownData,
6499
rhsStr: string,
65100
): string {
66-
const glslTypeName =
67-
dataType !== ShaderGenerator.UnknownData ? this.typeAnnotation(dataType) : 'auto';
68-
return `${pre}${glslTypeName} ${name} = ${rhsStr};`;
101+
const glslTypeName = dataType !== UnknownData ? this.ctx.resolve(dataType).value : 'auto';
102+
return `${this.ctx.pre}${glslTypeName} ${name} = ${rhsStr};`;
103+
}
104+
105+
override _returnStatement(statement: Return): string {
106+
if (this.#functionType !== 'normal') {
107+
const exprNode = statement[1];
108+
if (exprNode !== undefined) {
109+
// Resolving the expression to inspect it's type
110+
// We will resolve it again as part of the modifed statement
111+
const expr = this._expression(exprNode);
112+
if (expr.dataType !== UnknownData && expr.dataType.type.startsWith('vec')) {
113+
const block = super._block(
114+
[NODE.block, [[NODE.assignmentExpr, 'gl_Position', '=', exprNode], [NODE.return]]],
115+
{ gl_Position: gl_PositionSnippet.$ },
116+
);
117+
118+
return `${this.ctx.pre}${block}`;
119+
120+
// return (
121+
// super._statement([NODE.assignmentExpr, 'gl_Position', '=', exprNode]) +
122+
// '\n' +
123+
// super._returnStatement([NODE.return])
124+
// );
125+
}
126+
}
127+
}
128+
129+
return super._returnStatement(statement);
130+
}
131+
132+
override functionDefinition(options: ShaderGenerator.FunctionDefinitionOptions): string {
133+
// Ensuring the version macro is present in the shader output
134+
this.ctx.resolve(versionMacro);
135+
136+
if (options.functionType !== 'normal') {
137+
// Reserving a global name for gl_Position to avoid conflicts
138+
// TODO: This needs a better API, like a `reservedIdentifiers` array
139+
// passed to tgpu.resolve
140+
const gl_Position = tgpu.const(d.f32, 0).$name('gl_Position');
141+
this.ctx.getUniqueName(gl_Position);
142+
}
143+
144+
// Function body
145+
let lastFunctionType = this.#functionType;
146+
this.#functionType = options.functionType;
147+
try {
148+
const body = this._block(options.body);
149+
150+
// Only after generating the body can we determine the return type
151+
const returnType = options.determineReturnType();
152+
153+
if (options.functionType !== 'normal') {
154+
return `void main() ${body}`;
155+
}
156+
157+
const argList = options.args
158+
// Stripping out unused arguments in entry functions
159+
.filter((arg) => arg.used || options.functionType === 'normal')
160+
.map((arg) => {
161+
return `${this.ctx.resolve(arg.decoratedType).value} ${arg.name}`;
162+
})
163+
.join(', ');
164+
165+
return `${this.ctx.resolve(returnType).value} ${options.name}(${argList}) ${body}`;
166+
} finally {
167+
this.#functionType = lastFunctionType;
168+
}
69169
}
70170
}
71171

‎packages/typegpu-gl/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
export { initWithGL } from './initWithGL.ts';
22
export { initWithGLFallback } from './initWithGLFallback.ts';
3+
export { GLOptions } from './glOptions.ts';

‎packages/typegpu-gl/tests/glslGenerator.test.ts‎

Lines changed: 110 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { describe, expect, it } from 'vitest';
22
import tgpu, { d } from 'typegpu';
3-
import glslGenerator, { translateWgslTypeToGlsl } from '../src/glslGenerator.ts';
3+
import { GLOptions } from '@typegpu/gl';
4+
import { translateWgslTypeToGlsl } from '../src/glslGenerator.ts';
45

56
describe('translateWgslTypeToGlsl', () => {
67
it('translates scalar types', () => {
@@ -48,18 +49,14 @@ describe('translateWgslTypeToGlsl', () => {
4849

4950
describe('GlslGenerator - variable declarations', () => {
5051
it('generates GLSL-style variable declarations for JS function', () => {
51-
// 'use gpu' function - uses the TGSL code generation path, which calls functionDefinition
52-
const fragFn = tgpu.fragmentFn({
53-
in: { uv: d.vec2f },
54-
out: d.vec4f,
55-
})((input) => {
52+
const main = () => {
5653
'use gpu';
5754
// A variable that uses a vector type
58-
const color = d.vec4f(input.uv[0], input.uv[1], 0, 1);
55+
const color = d.vec4f(1, 0, 0, 1);
5956
return color;
60-
});
57+
};
6158

62-
const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
59+
const result = tgpu.resolveWithContext([main], GLOptions());
6360
// Should contain the resolved function code
6461
expect(result.code).toBeDefined();
6562
expect(result.code.length).toBeGreaterThan(0);
@@ -78,28 +75,38 @@ describe('GlslGenerator - variable declarations', () => {
7875
return d.vec4f(x, 0, 0, 1);
7976
});
8077

81-
const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
78+
const result = tgpu.resolveWithContext([fragFn], GLOptions());
8279
expect(result.code).toBeDefined();
8380
// Variable declaration for f32 should be `float`
8481
expect(result.code).toContain('float ');
8582
});
8683
});
8784

88-
describe('GlslGenerator - functionDefinition post-processing', () => {
89-
it('translates WGSL type constructor calls in JS function body to GLSL', () => {
90-
const fragFn = tgpu.fragmentFn({
91-
in: { uv: d.vec2f },
92-
out: d.vec4f,
93-
})((input) => {
85+
describe('GlslGenerator - function definitions', () => {
86+
it('generates proper function signatures', () => {
87+
function add(a: number, b: number) {
9488
'use gpu';
95-
return d.vec4f(input.uv[0], 0.0, 0.0, 1.0);
96-
});
89+
return a + b;
90+
}
9791

98-
const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
99-
// vec4f(...) in the body should become vec4(...)
100-
expect(result.code).toContain('vec4(');
101-
// Should not contain the WGSL-style constructor in the body
102-
expect(result.code).not.toMatch(/\bvec4f\s*\(/);
92+
function main() {
93+
'use gpu';
94+
return add(1.5, 1.2);
95+
}
96+
97+
const result = tgpu.resolveWithContext([main], GLOptions());
98+
99+
expect(result.code).toMatchInlineSnapshot(`
100+
"#version 300 es
101+
102+
float add(float a, float b) {
103+
return (a + b);
104+
}
105+
106+
float main() {
107+
return add(1.5f, 1.2f);
108+
}"
109+
`);
103110
});
104111

105112
it('translates vec3f to vec3 in function body', () => {
@@ -111,11 +118,46 @@ describe('GlslGenerator - functionDefinition post-processing', () => {
111118
return d.vec4f(color[0], color[1], color[2], 1.0);
112119
});
113120

114-
const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
121+
const result = tgpu.resolveWithContext([fragFn], GLOptions());
115122
expect(result.code).toContain('vec3(');
116123
expect(result.code).not.toMatch(/\bvec3f\s*\(/);
117124
expect(result.code).toContain('vec4(');
118125
});
126+
127+
it('generates proper struct definition', () => {
128+
const Boid = d.struct({
129+
pos: d.vec3f,
130+
vel: d.vec3f,
131+
});
132+
133+
function createBoid() {
134+
'use gpu';
135+
return Boid({ pos: d.vec3f(), vel: d.vec3f(0, 1, 0) });
136+
}
137+
138+
function main() {
139+
'use gpu';
140+
const boid = createBoid();
141+
}
142+
143+
const result = tgpu.resolve([main], GLOptions());
144+
expect(result).toMatchInlineSnapshot(`
145+
"#version 300 es
146+
147+
struct Boid {
148+
vec3 pos;
149+
vec3 vel;
150+
};
151+
152+
Boid createBoid() {
153+
return Boid(vec3(), vec3(0, 1, 0));
154+
}
155+
156+
void main() {
157+
Boid boid = createBoid();
158+
}"
159+
`);
160+
});
119161
});
120162

121163
describe('GlslGenerator - entry point generation with JS functions', () => {
@@ -127,41 +169,76 @@ describe('GlslGenerator - entry point generation with JS functions', () => {
127169
return Out({ pos: d.vec4f(0.0, 0.0, 0.0, 1.0) });
128170
});
129171

130-
const result = tgpu.resolveWithContext([vertFn], { unstable_shaderGenerator: glslGenerator });
172+
const result = tgpu.resolveWithContext([vertFn], GLOptions());
131173
expect(result.code).toBeDefined();
132174
expect(result.code.length).toBeGreaterThan(0);
133175
// The body should have translated type names
134176
expect(result.code).toContain('vec4(');
135177
expect(result.code).not.toMatch(/\bvec4f\s*\(/);
136178

137179
expect(result.code).toMatchInlineSnapshot(`
138-
"struct vertFn_Output {
139-
@builtin(position) pos: vec4,
140-
}
180+
"#version 300 es
141181
142-
@vertex fn vertFn() -> vertFn_Output {
182+
struct vertFn_Output {
183+
vec4 pos;
184+
};
185+
186+
void main() {
143187
return vertFn_Output(vec4(0, 0, 0, 1));
144188
}"
145189
`);
146190
});
147191

192+
// it('resolves a vertex function returning a position', () => {
193+
// const vertFn = tgpu.vertexFn({
194+
// out: { position: d.builtin.position },
195+
// })(() => {
196+
// 'use gpu';
197+
// return {
198+
// position: d.vec4f(1.0, 0.0, 0.0, 1.0),
199+
// };
200+
// });
201+
202+
// const result = tgpu.resolveWithContext([vertFn], GLOptions());
203+
// expect(result.code).toBeDefined();
204+
// expect(result.code).toContain('vec4(');
205+
// expect(result.code).not.toMatch(/\bvec4f\s*\(/);
206+
207+
// expect(result.code).toMatchInlineSnapshot(`
208+
// "#version 300 es
209+
210+
// void main() {
211+
// gl_Position = vec4(1, 0, 0, 1);
212+
// return;
213+
// }"
214+
// `);
215+
// });
216+
148217
it('resolves a fragment function returning a color using GLSL generator', () => {
149218
const fragFn = tgpu.fragmentFn({
150219
out: d.vec4f,
151220
})(() => {
152221
'use gpu';
222+
// This variable should get renamed to not conflict with
223+
// the global.
224+
const gl_Position = 1;
153225
return d.vec4f(1.0, 0.0, 0.0, 1.0);
154226
});
155227

156-
const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
228+
const result = tgpu.resolveWithContext([fragFn], GLOptions());
157229
expect(result.code).toBeDefined();
158-
// The body should have translated vec4f → vec4
159230
expect(result.code).toContain('vec4(');
160231
expect(result.code).not.toMatch(/\bvec4f\s*\(/);
161232

162233
expect(result.code).toMatchInlineSnapshot(`
163-
"@fragment fn fragFn() -> @location(0) vec4 {
164-
return vec4(1, 0, 0, 1);
234+
"#version 300 es
235+
236+
void main() {
237+
int gl_Position_1 = 1;
238+
{
239+
gl_Position = vec4(1, 0, 0, 1);
240+
return;
241+
}
165242
}"
166243
`);
167244
});

0 commit comments

Comments
 (0)