More GLSL generation

Author: iwoplaza

packages/typegpu-glsl/src/index.ts

@@ -1,13 +1,78 @@
 import { type ShaderGenerator, type Snippet, WgslGenerator } from 'typegpu';
 import type * as tinyest from 'tinyest';
+import type { AnyData } from 'typegpu/data';
+
+const wgslToGlslTypeMap = {
+  f32: 'float',
+  f64: 'double',
+  i32: 'int',
+  u32: 'uint',
+  void: 'void',
+};
+
+function resolveSchema(
+  ctx: ShaderGenerator.ResolutionCtx,
+  schema: AnyData,
+): { typePrefix: string; typeSuffix: string } {
+  if (schema.type in wgslToGlslTypeMap) {
+    return {
+      typePrefix:
+        wgslToGlslTypeMap[schema.type as keyof typeof wgslToGlslTypeMap],
+      typeSuffix: '',
+    };
+  }
+
+  if (schema.type === 'array') {
+    return {
+      typePrefix: resolveSchema(ctx, schema.elementType as AnyData).typePrefix,
+      typeSuffix: `[${schema.length}]`,
+    };
+  }
+
+  // TODO: Make struct declaration the responsibility of the
+  //       shader generator
+  // if (schema.type === 'struct') {
+  //   ctx.addDeclaration();
+
+  //   const fields = schema.fields.map((field) =>
+  //     `${resolveSchema(ctx, field.type).typePrefix} ${field.name}`
+  //   ).join(', ');
+  //   return {
+  //     typePrefix: `struct ${structName} { ${fields} }`,
+  //     typeSuffix: '',
+  //   };
+  // }
+
+  throw new Error(`Unsupported type: ${String(schema)}`);
+}
 
 export class GLSLShaderGenerator extends WgslGenerator
   implements ShaderGenerator {
   public override expression(expression: tinyest.Expression): Snippet {
     return super.expression(expression);
   }
 
-  public override functionBody(bodyNode: tinyest.Block): string {
-    return 'hello';
+  public override functionHeader(
+    { type, args, id, returnType }: ShaderGenerator.FunctionHeaderOptions,
+  ): string {
+    if (type === 'compute') {
+      throw new Error('Compute shaders are not supported in GLSL');
+    }
+
+    const argList = args.map((arg) => {
+      const segment = resolveSchema(this.ctx, arg.dataType as AnyData);
+      return `${segment.typePrefix} ${arg.value}${segment.typeSuffix}`;
+    }).join(', ');
+
+    const returnSegment = resolveSchema(this.ctx, returnType);
+
+    if (returnSegment.typeSuffix.length > 0) {
+      throw new Error(`Unsupported return type: ${String(returnType)}`);
+    }
+
+    // Entry functions are always named 'main' in GLSL
+    const realId = type === 'normal' ? id : 'main';
+
+    return `${returnSegment.typePrefix} ${realId}(${argList})`;
   }
 }

packages/typegpu-glsl/tests/functionDefinition.test.ts

@@ -1,11 +1,51 @@
 import { expect, test } from 'vitest';
+import * as d from 'typegpu/data';
 import { asGLSL } from './utils/asGLSL.ts';
 
-test('functionDefinition', () => {
+test('no args, f32 return', () => {
   const foo = () => {
     'use gpu';
-    return 1 + 2;
+    return d.f32(1 + 2);
   };
 
-  expect(asGLSL(foo)).toMatchInlineSnapshot(`"fn foo() hello"`);
+  expect(asGLSL(foo)).toMatchInlineSnapshot(`
+    "float foo() {
+      return 3;
+    }"
+  `);
+});
+
+test('no args, void return', () => {
+  const foo = () => {
+    'use gpu';
+    const a = 1 + 2;
+  };
+
+  expect(asGLSL(foo)).toMatchInlineSnapshot(`
+    "void foo() {
+      var a = 3;
+    }"
+  `);
+});
+
+test('primitive args, f32 return', () => {
+  const foo = (a: number, b: number) => {
+    'use gpu';
+    return d.f32(a + b + 2);
+  };
+
+  const bar = () => {
+    'use gpu';
+    foo(1, 2);
+  };
+
+  expect(asGLSL(bar)).toMatchInlineSnapshot(`
+    "float foo(int a, int b) {
+      return f32(((a + b) + 2));
+    }
+
+    void bar() {
+      foo(1, 2);
+    }"
+  `);
 });

packages/typegpu/src/tgsl/shaderGenerator.ts

@@ -2,6 +2,7 @@ import type { Block, Expression, Statement } from 'tinyest';
 import type { AnyData } from '../data/dataTypes.ts';
 import type { Snippet } from '../data/snippet.ts';
 import type { GenerationCtx } from './generationHelpers.ts';
+import type { ResolutionCtx as $ResolutionCtx } from '../types.ts';
 
 export interface FunctionDefinitionExtra {
   type: 'normal' | 'vertex' | 'fragment' | 'compute';
@@ -23,6 +24,7 @@ export interface ShaderGenerator {
 }
 
 export namespace ShaderGenerator {
+  export type ResolutionCtx = $ResolutionCtx;
   export interface FunctionHeaderOptions extends FunctionDefinitionExtra {
     id: string;
     args: Snippet[];

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -127,7 +127,7 @@ export class WgslGenerator implements ShaderGenerator {
     this.#ctx = ctx;
   }
 
-  private get ctx(): GenerationCtx {
+  public get ctx(): GenerationCtx {
     if (!this.#ctx) {
       throw new Error(
         'WGSL Generator has not yet been initialized. Please call initialize(ctx) before using the generator.',