feat: Use @typegpu/gl as a fallback when @typegpu/three is made to generate GLSL

packages/typegpu-three/package.json

@@ -50,6 +50,7 @@
     "vitest": "catalog:test"
   },
   "peerDependencies": {
+    "@typegpu/gl": "workspace:^",
     "three": ">0.126.0",
     "typegpu": "workspace:^"
   }

packages/typegpu-three/src/typegpu-node.ts

@@ -5,6 +5,7 @@ import * as THREE from 'three/webgpu';
 import * as TSL from 'three/tsl';
 import { tgpu, d, type Namespace, type TgpuVar, type ResolvedDeclaration } from 'typegpu';
 import WGSLNodeBuilder from 'three/src/renderers/webgpu/nodes/WGSLNodeBuilder.js';
+import { glOptions } from '@typegpu/gl';
 
 /**
  * State held by the node, used during shader generation.
@@ -19,7 +20,7 @@ interface TgpuFnNodeData extends THREE.NodeData {
 }
 
 abstract class StageData {
-  declare readonly type: 'analyze' | 'generate';
+  declare readonly type: 'setup' | 'analyze' | 'generate';
   readonly stage: 'vertex' | 'fragment' | 'compute' | null;
   readonly namespace: Namespace;
 
@@ -49,13 +50,19 @@ class AnalyzeStageData extends StageData {
   readonly type = 'analyze';
 }
 
+class SetupStageData extends StageData {
+  readonly type = 'setup';
+}
+
 class BuilderData {
   generateStageDataMap: Map<'vertex' | 'fragment' | 'compute' | null, GenerateStageData>;
   analyzeStageDataMap: Map<'vertex' | 'fragment' | 'compute' | null, AnalyzeStageData>;
+  setupStageDataMap: Map<'vertex' | 'fragment' | 'compute' | null, SetupStageData>;
 
   constructor() {
     this.generateStageDataMap = new Map();
     this.analyzeStageDataMap = new Map();
+    this.setupStageDataMap = new Map();
   }
 
   getGenerateStageData(stage: 'vertex' | 'fragment' | 'compute' | null): GenerateStageData {
@@ -75,6 +82,15 @@ class BuilderData {
     }
     return stageData;
   }
+
+  getSetupStageData(stage: 'vertex' | 'fragment' | 'compute' | null): SetupStageData {
+    let stageData = this.setupStageDataMap.get(stage);
+    if (!stageData) {
+      stageData = new SetupStageData(stage);
+      this.setupStageDataMap.set(stage, stageData);
+    }
+    return stageData;
+  }
 }
 
 const builderDataMap = new WeakMap<THREE.NodeBuilder, BuilderData>();
@@ -142,6 +158,11 @@ function forceExplicitVoidReturn(codeIn: string) {
   return codeIn.substring(0, closingParen + 1) + '-> void' + codeIn.substring(closingParen + 1);
 }
 
+function isWebGL(builder: THREE.NodeBuilder): boolean {
+  const backend = builder.renderer?.backend;
+  return (backend as { isWebGLBackend?: boolean } | undefined)?.isWebGLBackend === true;
+}
+
 class TgpuFnNode<T> extends THREE.Node {
   #impl: () => T;
 
@@ -163,6 +184,7 @@ class TgpuFnNode<T> extends THREE.Node {
   }
 
   #getNodeFunction(builder: THREE.NodeBuilder) {
+    const webgl = isWebGL(builder);
     const nodeData = builder.getDataFromNode(this) as TgpuFnNodeData;
     let builderData = builderDataMap.get(builder);
 
@@ -183,6 +205,7 @@ class TgpuFnNode<T> extends THREE.Node {
       const resolved = withGeneratingFnNodeCtx(ctx, () => {
         const { code, declarations } = tgpu.resolveWithContext([this.#impl], {
           names: stageData.namespace,
+          ...(webgl ? glOptions() : {}),
         });
 
         // Resolving this.#impl as second time in the same
@@ -192,6 +215,7 @@ class TgpuFnNode<T> extends THREE.Node {
           template: 'impl',
           externals: { impl: this.#impl },
           unstable_minify: false, // TODO(#2826): investigate
+          ...(webgl ? glOptions() : {}),
         });
 
         return {
@@ -218,7 +242,7 @@ class TgpuFnNode<T> extends THREE.Node {
         functionId: resolved.functionId,
         nodeFunction: builder.parser.parseFunction(
           // TODO: Upstream a fix to Three.js that accepts functions with no return type
-          forceExplicitVoidReturn(fnDeclaration),
+          webgl ? fnDeclaration : forceExplicitVoidReturn(fnDeclaration),
         ),
         // Including code that was resolved before the function as another node
         // that this node depends on
@@ -231,6 +255,7 @@ class TgpuFnNode<T> extends THREE.Node {
   }
 
   #analyzeFunction(builder: THREE.NodeBuilder) {
+    const webgl = isWebGL(builder);
     let builderData = builderDataMap.get(builder);
 
     if (!builderData) {
@@ -252,6 +277,7 @@ class TgpuFnNode<T> extends THREE.Node {
         template: 'impl',
         externals: { impl: this.#impl },
         unstable_minify: false, // TODO(#2826): investigate
+        ...(webgl ? glOptions() : {}),
       }),
     );
   }
@@ -269,6 +295,7 @@ class TgpuFnNode<T> extends THREE.Node {
     builder: THREE.NodeBuilder,
     output: string | null | undefined,
   ): string | null | undefined {
+    const webgl = isWebGL(builder);
     this.#getNodeFunction(builder); // making sure the node function exists
 
     const nodeData = builder.getDataFromNode(this) as TgpuFnNodeData;
@@ -289,16 +316,17 @@ class TgpuFnNode<T> extends THREE.Node {
       }
 
       const varValue = dep.node.build(builder);
+
+      const code = tgpu.resolve({
+        names: stageData.namespace,
+        // oxlint-disable-next-line typescript/no-base-to-string
+        template: `$var$ = ${varValue};\n`,
+        externals: { $var$: bridgeVar },
+        ...(webgl ? glOptions() : {}),
+      });
+
       // @ts-expect-error: it's there
-      builder.addLineFlowCode(
-        tgpu.resolve({
-          names: stageData.namespace,
-          // oxlint-disable-next-line typescript/no-base-to-string
-          template: `$var$ = ${varValue};\n`,
-          externals: { $var$: bridgeVar },
-        }),
-        this,
-      );
+      builder.addLineFlowCode(code, this);
     }
 
     return output === 'property' ? nodeData.custom.functionId : `${nodeData.custom.functionId}()`;
@@ -406,7 +434,7 @@ export const fromTSL = tgpu.comptime(((node, type) => {
 
   // In THREE, the type of array buffers equals to the type of the element.
   const wgslTypeFromTgpu = convertTypeToExplicit(
-    `${d.isWgslArray(tgpuType) ? tgpuType.elementType : tgpuType}`,
+    d.isWgslArray(tgpuType) ? tgpuType.elementType.type : tgpuType.type,
   );
 
   if (!sharedBuilder) {

packages/typegpu-three/tests/typegpu-node.test.ts

@@ -1,7 +1,7 @@
 import * as THREE from 'three/webgpu';
 import * as TSL from 'three/tsl';
 import WGSLNodeBuilder from 'three/src/renderers/webgpu/nodes/WGSLNodeBuilder.js';
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 import { tgpu, d } from 'typegpu';
 import { fromTSL, toTSL } from '@typegpu/three';
 
@@ -23,6 +23,14 @@ class ObservableFloatNode extends THREE.Node {
   }
 }
 
+class THREEWebGPUBackendMock {
+  isWebGPUBackend = true;
+}
+
+class THREERendererMock {
+  backend = new THREEWebGPUBackendMock();
+}
+
 function observableAccessor() {
   const node = new ObservableFloatNode();
   return {
@@ -33,6 +41,7 @@ function observableAccessor() {
 
 function builderFor(stage: 'analyze' | 'generate') {
   const builder = new WGSLNodeBuilder();
+  builder.renderer = new THREERendererMock() as unknown as THREE.Renderer;
   builder.setShaderStage('fragment');
   builder.setBuildStage(stage);
   return builder;