feat: Initial @typegpu/glsl POC

packages/typegpu-glsl/package.json

@@ -0,0 +1,43 @@
+{
+  "name": "@typegpu/glsl",
+  "type": "module",
+  "version": "0.9.0",
+  "description": "An alternate shader generator for TypeGPU, producing GLSL code.",
+  "exports": {
+    ".": "./src/index.ts",
+    "./package.json": "./package.json"
+  },
+  "publishConfig": {
+    "directory": "dist",
+    "linkDirectory": false,
+    "main": "./dist/index.mjs",
+    "types": "./dist/index.d.mts",
+    "exports": {
+      "./package.json": "./dist/package.json",
+      ".": {
+        "types": "./dist/index.d.mts",
+        "default": "./dist/index.mjs"
+      }
+    }
+  },
+  "sideEffects": false,
+  "scripts": {
+    "build": "unbuild",
+    "test:types": "pnpm tsc --p ./tsconfig.json --noEmit",
+    "prepublishOnly": "tgpu-dev-cli prepack"
+  },
+  "keywords": [],
+  "license": "MIT",
+  "peerDependencies": {
+    "typegpu": "^0.9.0",
+    "tinyest": "^0.2.0"
+  },
+  "devDependencies": {
+    "@typegpu/tgpu-dev-cli": "workspace:*",
+    "@webgpu/types": "catalog:types",
+    "typegpu": "workspace:*",
+    "typescript": "catalog:types",
+    "unbuild": "catalog:build",
+    "unplugin-typegpu": "workspace:*"
+  }
+}

packages/typegpu-glsl/src/glslGenerator.ts

@@ -0,0 +1,78 @@
+import type * as tinyest from 'tinyest';
+import { type ShaderGenerator, type Snippet, WgslGenerator } from 'typegpu';
+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 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/src/index.ts

@@ -0,0 +1 @@
+export { GLSLShaderGenerator } from './glslGenerator.ts';

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

@@ -0,0 +1,51 @@
+import * as d from 'typegpu/data';
+import { expect, test } from 'vitest';
+import { asGLSL } from './utils/asGLSL.ts';
+
+test('no args, f32 return', () => {
+  const foo = () => {
+    'use gpu';
+    return d.f32(1 + 2);
+  };
+
+  expect(asGLSL(foo)).toMatchInlineSnapshot(`
+    "float foo() {
+      return 3f;
+    }"
+  `);
+});
+
+test('no args, void return', () => {
+  const foo = () => {
+    'use gpu';
+    const a = 1 + 2;
+  };
+
+  expect(asGLSL(foo)).toMatchInlineSnapshot(`
+    "void foo() {
+      const 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) + 2i));
+    }
+
+    void bar() {
+      foo(1i, 2i);
+    }"
+  `);
+});

packages/typegpu-glsl/tests/utils/asGLSL.ts

@@ -0,0 +1,12 @@
+import tgpu from 'typegpu';
+import { GLSLShaderGenerator } from '../../src/index.ts';
+
+type ResolvableArray = Parameters<typeof tgpu.resolve>[0];
+
+/**
+ * Just a shorthand for tgpu.resolve, with a custom generator
+ */
+export function asGLSL(...values: ResolvableArray): string {
+  const shaderGenerator = new GLSLShaderGenerator();
+  return tgpu.resolve(values, { shaderGenerator });
+}

packages/typegpu/src/core/function/fnCore.ts

@@ -124,6 +124,13 @@ export function createFnCore(
           body = replacedImpl.slice(providedArgs.range.end);
         }
 
+        const fnAttribute = extra.type === 'vertex'
+          ? '@vertex '
+          : extra.type === 'fragment'
+          ? '@fragment '
+          : extra.type === 'compute'
+          ? `@compute @workgroup_size(${extra.workgroupSize?.join(', ')}) `
+          : '';
         ctx.addDeclaration(`${fnAttribute}fn ${id}${header}${body}`);
         return snip(id, returnType, /* origin */ 'runtime');
       }
@@ -177,16 +184,12 @@ export function createFnCore(
         );
       }
 
-      // generate wgsl string
+      // generate shader string
 
-      const { head, body, returnType: actualReturnType } = ctx.fnToWgsl({
-        functionType: fnAttribute.includes('@compute')
-          ? 'compute'
-          : fnAttribute.includes('@vertex')
-          ? 'vertex'
-          : fnAttribute.includes('@fragment')
-          ? 'fragment'
-          : 'normal',
+      const { code, returnType: actualReturnType } = ctx.fnToShaderCode({
+        type: extra.type,
+        workgroupSize: extra.workgroupSize,
+        id,
         argTypes,
         params: ast.params,
         returnType,

packages/typegpu/src/resolutionCtx.ts

@@ -20,7 +20,6 @@ import {
   type TgpuLazy,
   type TgpuSlot,
 } from './core/slot/slotTypes.ts';
-import { getAttributesString } from './data/attributes.ts';
 import {
   type AnyData,
   isData,
@@ -70,8 +69,8 @@ import type {
   ItemLayer,
   ItemStateStack,
   ResolutionCtx,
+  ShaderStage,
   StackLayer,
-  TgpuShaderStage,
   Wgsl,
 } from './types.ts';
 import { CodegenState, isSelfResolvable, NormalState } from './types.ts';
@@ -134,7 +133,7 @@ class ItemStateStackImpl implements ItemStateStack {
   }
 
   pushFunctionScope(
-    functionType: 'normal' | TgpuShaderStage,
+    functionType: 'normal' | ShaderStage,
     args: Snippet[],
     argAliases: Record<string, Snippet>,
     returnType: AnyData | undefined,
@@ -430,9 +429,9 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     return this.#logGenerator.logResources;
   }
 
-  fnToWgsl(
-    options: FnToWgslOptions,
-  ): { head: Wgsl; body: Wgsl; returnType: AnyData } {
+  fnToShaderCode(
+    options: FnToShaderCodeOptions,
+  ): { code: string; returnType: AnyData } {
     let fnScopePushed = false;
 
     try {
@@ -492,7 +491,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
       }
 
       const scope = this._itemStateStack.pushFunctionScope(
-        options.functionType,
+        options.type,
         args,
         Object.fromEntries(argAliases),
         options.returnType,
@@ -501,7 +500,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
       fnScopePushed = true;
 
       this.#shaderGenerator.initGenerator(this);
-      const body = this.#shaderGenerator.functionDefinition(options.body);
+      const body = this.#shaderGenerator.functionBody(options.body);
 
       let returnType = options.returnType;
       if (!returnType) {
@@ -526,6 +525,14 @@ export class ResolutionCtxImpl implements ResolutionCtx {
         returnType = concretize(returnType);
       }
 
+      const header = this.#shaderGenerator.functionHeader({
+        id: options.id,
+        type: options.type,
+        workgroupSize: options.workgroupSize,
+        returnType,
+        args,
+      });
+
       return {
         head: resolveFunctionHeader(this, args, returnType),
         body,
@@ -1013,19 +1020,3 @@ export function resolve(
     logResources: ctx.logResources,
   };
 }
-
-export function resolveFunctionHeader(
-  ctx: ResolutionCtx,
-  args: Snippet[],
-  returnType: AnyData,
-) {
-  const argList = args
-    .map((arg) => `${arg.value}: ${ctx.resolve(arg.dataType as AnyData).value}`)
-    .join(', ');
-
-  return returnType.type !== 'void'
-    ? `(${argList}) -> ${getAttributesString(returnType)}${
-      ctx.resolve(returnType).value
-    } `
-    : `(${argList}) `;
-}

packages/typegpu/src/tgpuBindGroupLayout.ts

@@ -66,7 +66,7 @@ import type {
   NullableToOptional,
   Prettify,
 } from './shared/utilityTypes.ts';
-import type { TgpuShaderStage } from './types.ts';
+import type { ShaderStage } from './types.ts';
 import type { Unwrapper } from './unwrapper.ts';
 import type { WgslComparisonSampler, WgslSampler } from './data/sampler.ts';
 
@@ -90,7 +90,7 @@ export type TgpuLayoutEntryBase = {
    * @default ['compute','fragment'] for mutable resources
    * @default ['compute','vertex','fragment'] for everything else
    */
-  visibility?: TgpuShaderStage[];
+  visibility?: ShaderStage[];
 };
 
 export type TgpuLayoutUniform = TgpuLayoutEntryBase & {
@@ -500,8 +500,8 @@ export class MissingBindingError extends Error {
 // Implementation
 // --------------
 
-const DEFAULT_MUTABLE_VISIBILITY: TgpuShaderStage[] = ['compute', 'fragment'];
-const DEFAULT_READONLY_VISIBILITY: TgpuShaderStage[] = [
+const DEFAULT_MUTABLE_VISIBILITY: ShaderStage[] = ['compute', 'fragment'];
+const DEFAULT_READONLY_VISIBILITY: ShaderStage[] = [
   'compute',
   'vertex',
   'fragment',

packages/typegpu/src/tgsl/shaderGenerator.ts

@@ -2,6 +2,16 @@ 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, ShaderStage } from '../types.ts';
+
+export interface FunctionDefinitionExtra {
+  type: 'normal' | ShaderStage;
+  /**
+   * Only applicable to 'compute' entry functions
+   */
+  workgroupSize?: number[] | undefined;
+  bodyNode: Block;
+}
 
 export interface ShaderGenerator {
   initGenerator(ctx: GenerationCtx): void;
@@ -10,5 +20,20 @@ export interface ShaderGenerator {
   typedExpression(expression: Expression, expectedType: AnyData): Snippet;
   expression(expression: Expression): Snippet;
   statement(statement: Statement): string;
-  functionDefinition(body: Block): string;
+  functionDefinition(
+    options: ShaderGenerator.FunctionDefinitionOptions,
+  ): string;
+  functionBody(body: Block): string;
+}
+
+export declare namespace ShaderGenerator {
+  type ResolutionCtx = $ResolutionCtx;
+  interface FunctionDefinitionOptions extends FunctionDefinitionExtra {
+    id: string;
+    args: Snippet[];
+    /**
+     * The return type of the function.
+     */
+    returnType?: AnyData | undefined;
+  }
 }