feat(@typegpu/gl): Generate GLSL function signatures

iwoplaza · iwoplaza · commit bf31165ba168 · 2026-04-11T13:30:09.000+02:00
diff --git a/packages/typegpu-gl/src/glOptions.ts b/packages/typegpu-gl/src/glOptions.ts
@@ -0,0 +1,7 @@
+import glslGenerator from './glslGenerator.ts';
+
+export function GLOptions() {
+  return {
+    unstable_shaderGenerator: glslGenerator,
+  };
+}
diff --git a/packages/typegpu-gl/src/glslGenerator.ts b/packages/typegpu-gl/src/glslGenerator.ts
@@ -1,10 +1,16 @@
-import { d, ShaderGenerator, WgslGenerator } from 'typegpu';
+import { NodeTypeCatalog as NODE } from 'tinyest';
+import type { Return } from 'tinyest';
+import tgpu, { d, ShaderGenerator, WgslGenerator } from 'typegpu';
+import { setName } from '../../typegpu/src/shared/meta';
+
+const UnknownData: typeof ShaderGenerator.UnknownData = ShaderGenerator.UnknownData;
 
 // ----------
 // WGSL → GLSL type name mapping
 // ----------
 
 const WGSL_TO_GLSL_TYPE: Record<string, string> = {
+  void: 'void',
   f32: 'float',
   u32: 'uint',
   i32: 'int',
@@ -37,13 +43,18 @@ export function translateWgslTypeToGlsl(wgslType: string): string {
   return WGSL_TO_GLSL_TYPE[wgslType] ?? wgslType;
 }
 
+const versionMacro = tgpu['~unstable'].declare('#version 300 es');
+const gl_PositionSnippet = tgpu['~unstable'].rawCodeSnippet('gl_Position', d.vec4f, 'private');
+
 /**
  * A GLSL ES 3.0 shader generator that extends WgslGenerator.
  * Overrides `dataType` to emit GLSL type names instead of WGSL ones,
  * and overrides variable declaration emission to use `type name = rhs` syntax.
  */
 export class GlslGenerator extends WgslGenerator {
-  public override typeAnnotation(data: d.BaseData): string {
+  #functionType: ShaderGenerator.TgpuShaderStage | 'normal' | undefined;
+
+  override typeAnnotation(data: d.BaseData): string {
     // For WGSL identity types (scalars, vectors, common matrices), map to GLSL directly.
     if (!d.isLooseData(data)) {
       const glslName = WGSL_TO_GLSL_TYPE[data.type];
@@ -56,16 +67,81 @@ export class GlslGenerator extends WgslGenerator {
     return super.typeAnnotation(data);
   }
 
-  protected override emitVarDecl(
-    pre: string,
+  override _emitVarDecl(
     _keyword: 'var' | 'let' | 'const',
     name: string,
     dataType: d.BaseData | ShaderGenerator.UnknownData,
     rhsStr: string,
   ): string {
-    const glslTypeName =
-      dataType !== ShaderGenerator.UnknownData ? this.typeAnnotation(dataType) : 'auto';
-    return `${pre}${glslTypeName} ${name} = ${rhsStr};`;
+    const glslTypeName = dataType !== UnknownData ? this.typeAnnotation(dataType) : 'auto';
+    return `${this.ctx.pre}${glslTypeName} ${name} = ${rhsStr};`;
+  }
+
+  override _returnStatement(statement: Return): string {
+    if (this.#functionType !== 'normal') {
+      const exprNode = statement[1];
+      if (exprNode !== undefined) {
+        // Resolving the expression to inspect it's type
+        // We will resolve it again as part of the modifed statement
+        const expr = this._expression(exprNode);
+        if (expr.dataType !== UnknownData && expr.dataType.type.startsWith('vec')) {
+          const block = super._block(
+            [NODE.block, [[NODE.assignmentExpr, 'gl_Position', '=', exprNode], [NODE.return]]],
+            { gl_Position: gl_PositionSnippet.$ },
+          );
+
+          return `${this.ctx.pre}${block}`;
+
+          // return (
+          //   super._statement([NODE.assignmentExpr, 'gl_Position', '=', exprNode]) +
+          //   '\n' +
+          //   super._returnStatement([NODE.return])
+          // );
+        }
+      }
+    }
+
+    return super._returnStatement(statement);
+  }
+
+  override functionDefinition(options: ShaderGenerator.FunctionDefinitionOptions): string {
+    // Ensuring the version macro is present in the shader output
+    this.ctx.resolve(versionMacro);
+
+    if (options.functionType !== 'normal') {
+      // Reserving a global name for gl_Position to avoid conflicts
+      // TODO: This needs a better API, like a `reservedIdentifiers` array
+      // passed to tgpu.resolve
+      const gl_Position = {};
+      setName(gl_Position, 'gl_Position');
+      console.log(this.ctx.getUniqueName(gl_Position));
+    }
+
+    // Function body
+    let lastFunctionType = this.#functionType;
+    this.#functionType = options.functionType;
+    try {
+      const body = this._block(options.body);
+
+      // Only after generating the body can we determine the return type
+      const returnType = options.determineReturnType();
+
+      if (options.functionType !== 'normal') {
+        return `void main() ${body}`;
+      }
+
+      const argList = options.args
+        // Stripping out unused arguments in entry functions
+        .filter((arg) => arg.used || options.functionType === 'normal')
+        .map((arg) => {
+          return `${this.ctx.resolve(arg.decoratedType).value} ${arg.name}`;
+        })
+        .join(', ');
+
+      return `${this.ctx.resolve(returnType).value} ${options.name}(${argList}) ${body}`;
+    } finally {
+      this.#functionType = lastFunctionType;
+    }
   }
 }
 
diff --git a/packages/typegpu-gl/src/index.ts b/packages/typegpu-gl/src/index.ts
@@ -1,2 +1,3 @@
 export { initWithGL } from './initWithGL.ts';
 export { initWithGLFallback } from './initWithGLFallback.ts';
+export { GLOptions } from './glOptions.ts';
diff --git a/packages/typegpu-gl/tests/glslGenerator.test.ts b/packages/typegpu-gl/tests/glslGenerator.test.ts
@@ -1,5 +1,6 @@
 import { describe, expect, it } from 'vitest';
 import tgpu, { d } from 'typegpu';
+import { GLOptions } from '@typegpu/gl';
 import glslGenerator, { translateWgslTypeToGlsl } from '../src/glslGenerator.ts';
 
 describe('translateWgslTypeToGlsl', () => {
@@ -48,18 +49,14 @@ describe('translateWgslTypeToGlsl', () => {
 
 describe('GlslGenerator - variable declarations', () => {
   it('generates GLSL-style variable declarations for JS function', () => {
-    // 'use gpu' function - uses the TGSL code generation path, which calls functionDefinition
-    const fragFn = tgpu.fragmentFn({
-      in: { uv: d.vec2f },
-      out: d.vec4f,
-    })((input) => {
+    const main = () => {
       'use gpu';
       // A variable that uses a vector type
-      const color = d.vec4f(input.uv[0], input.uv[1], 0, 1);
+      const color = d.vec4f(1, 0, 0, 1);
       return color;
-    });
+    };
 
-    const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
+    const result = tgpu.resolveWithContext([main], GLOptions());
     // Should contain the resolved function code
     expect(result.code).toBeDefined();
     expect(result.code.length).toBeGreaterThan(0);
@@ -78,28 +75,38 @@ describe('GlslGenerator - variable declarations', () => {
       return d.vec4f(x, 0, 0, 1);
     });
 
-    const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
+    const result = tgpu.resolveWithContext([fragFn], GLOptions());
     expect(result.code).toBeDefined();
     // Variable declaration for f32 should be `float`
     expect(result.code).toContain('float ');
   });
 });
 
-describe('GlslGenerator - functionDefinition post-processing', () => {
-  it('translates WGSL type constructor calls in JS function body to GLSL', () => {
-    const fragFn = tgpu.fragmentFn({
-      in: { uv: d.vec2f },
-      out: d.vec4f,
-    })((input) => {
+describe('GlslGenerator - function definitions', () => {
+  it('generates proper function signatures', () => {
+    function add(a: number, b: number) {
       'use gpu';
-      return d.vec4f(input.uv[0], 0.0, 0.0, 1.0);
-    });
+      return a + b;
+    }
 
-    const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
-    // vec4f(...) in the body should become vec4(...)
-    expect(result.code).toContain('vec4(');
-    // Should not contain the WGSL-style constructor in the body
-    expect(result.code).not.toMatch(/\bvec4f\s*\(/);
+    function main() {
+      'use gpu';
+      return add(1.5, 1.2);
+    }
+
+    const result = tgpu.resolveWithContext([main], GLOptions());
+
+    expect(result.code).toMatchInlineSnapshot(`
+      "#version 300 es
+
+      float add(float a, float b) {
+        return (a + b);
+      }
+
+      float main() {
+        return add(1.5f, 1.2f);
+      }"
+    `);
   });
 
   it('translates vec3f to vec3 in function body', () => {
@@ -111,7 +118,7 @@ describe('GlslGenerator - functionDefinition post-processing', () => {
       return d.vec4f(color[0], color[1], color[2], 1.0);
     });
 
-    const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
+    const result = tgpu.resolveWithContext([fragFn], GLOptions());
     expect(result.code).toContain('vec3(');
     expect(result.code).not.toMatch(/\bvec3f\s*\(/);
     expect(result.code).toContain('vec4(');
@@ -127,41 +134,76 @@ describe('GlslGenerator - entry point generation with JS functions', () => {
       return Out({ pos: d.vec4f(0.0, 0.0, 0.0, 1.0) });
     });
 
-    const result = tgpu.resolveWithContext([vertFn], { unstable_shaderGenerator: glslGenerator });
+    const result = tgpu.resolveWithContext([vertFn], GLOptions());
     expect(result.code).toBeDefined();
     expect(result.code.length).toBeGreaterThan(0);
     // The body should have translated type names
     expect(result.code).toContain('vec4(');
     expect(result.code).not.toMatch(/\bvec4f\s*\(/);
 
     expect(result.code).toMatchInlineSnapshot(`
-      "struct vertFn_Output {
+      "#version 300 es
+
+      struct vertFn_Output {
         @builtin(position) pos: vec4,
       }
 
-      @vertex fn vertFn() -> vertFn_Output {
+      void main() {
         return vertFn_Output(vec4(0, 0, 0, 1));
       }"
     `);
   });
 
+  // it('resolves a vertex function returning a position', () => {
+  //   const vertFn = tgpu.vertexFn({
+  //     out: { position: d.builtin.position },
+  //   })(() => {
+  //     'use gpu';
+  //     return {
+  //       position: d.vec4f(1.0, 0.0, 0.0, 1.0),
+  //     };
+  //   });
+
+  //   const result = tgpu.resolveWithContext([vertFn], GLOptions());
+  //   expect(result.code).toBeDefined();
+  //   expect(result.code).toContain('vec4(');
+  //   expect(result.code).not.toMatch(/\bvec4f\s*\(/);
+
+  //   expect(result.code).toMatchInlineSnapshot(`
+  //     "#version 300 es
+
+  //     void main() {
+  //       gl_Position = vec4(1, 0, 0, 1);
+  //       return;
+  //     }"
+  //   `);
+  // });
+
   it('resolves a fragment function returning a color using GLSL generator', () => {
     const fragFn = tgpu.fragmentFn({
       out: d.vec4f,
     })(() => {
       'use gpu';
+      // This variable should get renamed to not conflict with
+      // the global.
+      const gl_Position = 1;
       return d.vec4f(1.0, 0.0, 0.0, 1.0);
     });
 
-    const result = tgpu.resolveWithContext([fragFn], { unstable_shaderGenerator: glslGenerator });
+    const result = tgpu.resolveWithContext([fragFn], GLOptions());
     expect(result.code).toBeDefined();
-    // The body should have translated vec4f → vec4
     expect(result.code).toContain('vec4(');
     expect(result.code).not.toMatch(/\bvec4f\s*\(/);
 
     expect(result.code).toMatchInlineSnapshot(`
-      "@fragment fn fragFn() -> @location(0) vec4 {
-        return vec4(1, 0, 0, 1);
+      "#version 300 es
+
+      void main() {
+        int gl_Position_1 = 1;
+        {
+          gl_Position = vec4(1, 0, 0, 1);
+          return;
+        }
       }"
     `);
   });
diff --git a/packages/typegpu/src/core/function/fnCore.ts b/packages/typegpu/src/core/function/fnCore.ts
@@ -201,8 +201,10 @@ export function createFnCore(
 
       // generate wgsl string
 
-      const { code, returnType: actualReturnType } = ctx.fnToWgsl({
+      const { code, returnType: actualReturnType } = ctx.resolveFunction({
         functionType,
+        name: id,
+        workgroupSize,
         argTypes,
         entryInput,
         params: ast.params,
@@ -211,16 +213,7 @@ export function createFnCore(
         externalMap,
       });
 
-      let attributes = '';
-      if (functionType === 'compute') {
-        attributes = `@compute @workgroup_size(${workgroupSize?.join(', ')}) `;
-      } else if (functionType === 'vertex') {
-        attributes = `@vertex `;
-      } else if (functionType === 'fragment') {
-        attributes = `@fragment `;
-      }
-
-      ctx.addDeclaration(`${attributes}fn ${id}${code}`);
+      ctx.addDeclaration(code);
 
       return snip(id, actualReturnType, /* origin */ 'runtime');
     },
diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts
@@ -39,7 +39,7 @@ import wgslGenerator from './tgsl/wgslGenerator.ts';
 import type {
   ExecMode,
   ExecState,
-  FnToWgslOptions,
+  ResolveFunctionOptions,
   FunctionArgumentAccess,
   FunctionScopeLayer,
   ItemLayer,
@@ -466,7 +466,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     return this.#logGenerator.logResources;
   }
 
-  fnToWgsl(options: FnToWgslOptions): { code: string; returnType: BaseData } {
+  resolveFunction(options: ResolveFunctionOptions): { code: string; returnType: BaseData } {
     let fnScopePushed = false;
 
     try {
@@ -584,6 +584,8 @@ export class ResolutionCtxImpl implements ResolutionCtx {
 
       const code = this.gen.functionDefinition({
         functionType: options.functionType,
+        name: options.name,
+        workgroupSize: options.workgroupSize,
         args,
         body: options.body,
         determineReturnType: () => {
diff --git a/packages/typegpu/src/tgsl/shaderGenerator_members.ts b/packages/typegpu/src/tgsl/shaderGenerator_members.ts
@@ -5,10 +5,12 @@ import type { FunctionArgument, TgpuShaderStage } from '../types.ts';
 export { UnknownData } from '../data/dataTypes.ts';
 
 // types
-export type { ResolutionCtx, FunctionArgument } from '../types.ts';
+export type { ResolutionCtx, FunctionArgument, TgpuShaderStage } from '../types.ts';
 
 export interface FunctionDefinitionOptions {
   readonly functionType: 'normal' | TgpuShaderStage;
+  readonly name: string;
+  readonly workgroupSize?: readonly number[] | undefined;
   readonly args: readonly FunctionArgument[];
   readonly body: Block;
 
diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts
diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`export { initWithGL } from './initWithGL.ts';`
`2`	`2`	`export { initWithGLFallback } from './initWithGLFallback.ts';`
	`3`	`+export { GLOptions } from './glOptions.ts';`