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

Author: iwoplaza

packages/typegpu-gl/src/glOptions.ts

@@ -0,0 +1,7 @@
+import glslGenerator from './glslGenerator.ts';
+
+export function glOptions() {
+  return {
+    unstable_shaderGenerator: glslGenerator,
+  };
+}

packages/typegpu-gl/src/glslGenerator.ts

@@ -1,10 +1,17 @@
-import { d, ShaderGenerator, WgslGenerator } from 'typegpu';
+import { NodeTypeCatalog as NODE } from 'tinyest';
+import type { Return } from 'tinyest';
+import tgpu, { d, ShaderGenerator, WgslGenerator } from 'typegpu';
+
+type ResolutionCtx = ShaderGenerator.ResolutionCtx;
+
+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 +44,43 @@ export function translateWgslTypeToGlsl(wgslType: string): string {
   return WGSL_TO_GLSL_TYPE[wgslType] ?? wgslType;
 }
 
+/**
+ * Resolves a struct and adds its declaration to the resolution context.
+ * @param ctx - The resolution context.
+ * @param struct - The struct to resolve.
+ *
+ * @returns The resolved struct name.
+ */
+function resolveStruct(ctx: ResolutionCtx, struct: d.WgslStruct) {
+  const id = ctx.makeUniqueIdentifier(ShaderGenerator.getName(struct), 'global');
+
+  ctx.addDeclaration(`\
+struct ${id} {
+${Object.entries(struct.propTypes)
+  .map(([prop, type]) => `  ${ctx.resolve(type).value} ${prop};\n`)
+  .join('')}\
+};`);
+
+  return id;
+}
+
+const gl_PositionSnippet = tgpu['~unstable'].rawCodeSnippet('gl_Position', d.vec4f, 'private');
+
+interface EntryFnState {
+  structPropToVarMap: Record<string, string>;
+  outVars: { varName: string; propName: string }[];
+}
+
 /**
  * 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;
+  #entryFnState: EntryFnState | 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];
@@ -52,20 +89,153 @@ export class GlslGenerator extends WgslGenerator {
       }
     }
 
+    if (d.isWgslStruct(data)) {
+      return resolveStruct(this.ctx, data);
+    }
+
     // For all other types (structs, arrays, etc.) delegate to WGSL resolution.
     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.ctx.resolve(dataType).value : 'auto';
+    return `${this.ctx.pre}${glslTypeName} ${name} = ${rhsStr};`;
+  }
+
+  override _return(statement: Return): string {
+    const exprNode = statement[1];
+
+    if (exprNode === undefined) {
+      // Default behavior
+      return super._return(statement);
+    }
+
+    if (this.#functionType !== 'normal') {
+      // oxlint-disable-next-line no-non-null-assertion
+      const entryFnState = this.#entryFnState!;
+      const expectedReturnType = this.ctx.topFunctionReturnType;
+
+      if (typeof exprNode === 'object' && exprNode[0] === NODE.objectExpr) {
+        const transformed = Object.entries(exprNode[1]).map(([prop, rhsNode]) => {
+          let name: string | undefined = entryFnState.structPropToVarMap[prop];
+          if (name === undefined) {
+            if (
+              prop === '$position' ||
+              (expectedReturnType &&
+                d.isWgslStruct(expectedReturnType) &&
+                expectedReturnType.propTypes[prop] === d.builtin.position)
+            ) {
+              name = 'gl_Position';
+            } else {
+              name = this.ctx.makeUniqueIdentifier(prop, 'global');
+              entryFnState.outVars.push({ varName: name, propName: prop });
+            }
+            entryFnState.structPropToVarMap[prop] = name;
+          }
+          const rhsExpr = this._expression(rhsNode);
+          const type = rhsExpr.dataType as d.BaseData;
+
+          const snippet = tgpu['~unstable'].rawCodeSnippet(name, type as d.AnyData, 'private');
+
+          return {
+            name,
+            snippet,
+            assignment: [NODE.assignmentExpr, name, '=', rhsNode],
+          } as const;
+        });
+
+        const block = super._block(
+          [NODE.block, [...transformed.map((t) => t.assignment), [NODE.return]]],
+          Object.fromEntries(
+            transformed.map(({ name, snippet }) => {
+              return [name, snippet.$] as const;
+            }),
+          ),
+        );
+
+        return `${this.ctx.pre}${block}`;
+      } else {
+        // Resolving the expression to inspect it's type
+        // We will resolve it again as part of the modifed statement
+        const expr = expectedReturnType
+          ? this._typedExpression(exprNode, expectedReturnType)
+          : this._expression(exprNode);
+
+        if (expr.dataType === UnknownData) {
+          // Unknown data type, don't know what to do
+          return super._return(statement);
+        }
+
+        if (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._return(statement);
+  }
+
+  override functionDefinition(options: ShaderGenerator.FunctionDefinitionOptions): string {
+    if (options.functionType !== 'normal') {
+      this.ctx.reserveIdentifier('gl_Position', 'global');
+    }
+
+    // Function body
+    let lastFunctionType = this.#functionType;
+    this.#functionType = options.functionType;
+    if (options.functionType !== 'normal') {
+      if (this.#entryFnState) {
+        throw new Error('Cannot nest entry functions');
+      }
+      this.#entryFnState = { structPropToVarMap: {}, outVars: [] };
+    }
+
+    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') {
+        // oxlint-disable-next-line no-non-null-assertion
+        const entryFnState = this.#entryFnState!;
+        if (d.isWgslStruct(returnType)) {
+          for (const { varName, propName } of entryFnState.outVars) {
+            const dataType = returnType.propTypes[propName];
+            if (dataType && d.isDecorated(dataType)) {
+              const location = (dataType.attribs as d.AnyAttribute[]).find(
+                (a) => a.type === '@location',
+              )?.params[0];
+              this.ctx.addDeclaration(`layout(location = ${location}) out ${varName};`);
+            }
+          }
+        }
+        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;
+      this.#entryFnState = undefined;
+    }
   }
 }
 

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';