feat: Tweaks to feat/tgpu-three

apps/typegpu-docs/astro.config.mjs

@@ -211,6 +211,10 @@ export default defineConfig({
               label: '@typegpu/noise',
               slug: 'ecosystem/typegpu-noise',
             },
+            {
+              label: '@typegpu/three',
+              slug: 'ecosystem/typegpu-three',
+            },
             DEV && {
               label: '@typegpu/color',
               slug: 'ecosystem/typegpu-color',

apps/typegpu-docs/src/content/docs/ecosystem/typegpu-three.mdx

@@ -0,0 +1,91 @@
+---
+title: "@typegpu/three"
+---
+
+TSL (Three.js Shading Language) is a node-based shader composition system for Three.js. Shader logic and control flow is built up by composing special functions,
+with a focus on composability, intuitive sharing of logic across modules and customizability. TypeGPU fits naturally into this system thanks to the `@typegpu/three` package. You can choose to write your TSL building blocks in TypeGPU, which has a few benefits:
+- Control-flow like `if` statements and `for` loops makes use of familiar JavaScript syntax instead of special functions.
+- The code you write is semantically valid JavaScript, with types flowing through each expression.
+- Unit-testability, since you can call these functions on the CPU
+
+Below are a select few cases comparing TSL and TypeGPU:
+
+import { Card, CardGrid } from '@astrojs/starlight/components';
+
+## Node definition
+
+TSL:
+```ts
+const simulate = Fn(() => {
+  //
+  // ... TSL code ...
+  //
+});
+```
+
+TypeGPU:
+```ts
+const simulate = t3.toTSL(() => {
+  'use gpu';
+  //
+  // ... TypeGPU code ...
+  // 
+});
+```
+
+## Function definition
+
+TSL:
+```ts
+const oscSine = Fn(([t = time]) => {
+  return t.add(0.75).mul(Math.PI * 2).sin().mul(0.5).add(0.5);
+});
+```
+
+TypeGPU:
+```ts
+const oscSine = (t: number) => {
+  'use gpu';
+  return std.sin((t + 0.75) * Math.PI * 2) * 0.5 + 0.5;
+};
+```
+
+## If statements
+
+TSL:
+```ts
+If(instanceIndex.greaterThanEqual(uint(vertexCount)), () => {
+  Return();
+});
+```
+
+TypeGPU:
+```ts
+if (t3.instanceIndex.$ >= vertexCount) {
+  return;
+}
+```
+
+## For loops
+
+TSL:
+```ts
+Loop({ start: ptrStart, end: ptrEnd, type: 'uint', condition: '<' }, ({ i }) => {
+	const springId = springListBuffer.element( i ).toVar( 'springId' );
+	const springForce = springForceBuffer.element( springId );
+	const springVertexIds = springVertexIdBuffer.element( springId );
+	const factor = select( springVertexIds.x.equal( instanceIndex ), 1.0, - 1.0 );
+	force.addAssign( springForce.mul( factor ) );
+});
+```
+
+TypeGPU:
+```ts
+for (let i = ptrStart; i < ptrEnd; i++) {
+  const springId = springListBuffer.$[i];
+  const springForce = springForceBuffer.$[springId];
+  const springVertexIds = springVertexIdBuffer.$[springId];
+  const factor = std.select(-1, 1, springVertexIds.x === idx);
+  force = force.add(springForce.mul(d.f32(factor)));
+}
+```

apps/typegpu-docs/src/examples/threejs/compute-cloth/index.ts

@@ -4,7 +4,7 @@
 
 import * as d from 'typegpu/data';
 import * as THREE from 'three/webgpu';
-import { fromTSL, toTSL, uv } from '@typegpu/three';
+import { access, fromTSL, toTSL } from '@typegpu/three';
 import * as std from 'typegpu/std';
 
 import * as TSL from 'three/tsl';
@@ -248,7 +248,8 @@ function setupClothMesh(): THREE.Mesh {
 
   clothMaterial.colorNode = toTSL(() => {
     'use gpu';
-    const pattern = checkerBoard(uv().$.mul(5));
+    const uv = access.uv().$;
+    const pattern = checkerBoard(uv.mul(5));
     return std.mix(d.vec4f(0.4, 0.3, 0.3, 1), d.vec4f(1, 0.5, 0.4, 1), pattern);
   });
 

apps/typegpu-docs/src/examples/threejs/compute-cloth/verlet.ts

@@ -1,6 +1,6 @@
-import * as THREE from 'three/webgpu';
+import * as t3 from '@typegpu/three';
 import * as TSL from 'three/tsl';
-import { fromTSL, toTSL, type TSLAccessor } from '@typegpu/three';
+import * as THREE from 'three/webgpu';
 import * as d from 'typegpu/data';
 import * as std from 'typegpu/std';
 import { triNoise3D } from './triNoise.ts';
@@ -25,62 +25,50 @@ export const clothNumSegmentsY = 30;
 
 interface VerletSimulationOptions {
   sphereRadius: number;
-  sphereUniform: TSLAccessor<d.F32, THREE.UniformNode<number>>;
-  spherePositionUniform: TSLAccessor<d.Vec3f, THREE.UniformNode<THREE.Vector3>>;
+  sphereUniform: t3.TSLAccessor<d.F32, THREE.UniformNode<number>>;
+  spherePositionUniform: t3.TSLAccessor<
+    d.Vec3f,
+    THREE.UniformNode<THREE.Vector3>
+  >;
 }
 
+type TSLStorageAccessor<T extends d.AnyWgslData> = t3.TSLAccessor<
+  T,
+  THREE.StorageBufferNode
+>;
+
 export class VerletSimulation {
   readonly vertices: Vertex[];
   readonly springs: Spring[];
   readonly vertexColumns: Vertex[][];
 
-  readonly stiffnessUniform: TSLAccessor<d.F32, THREE.UniformNode<number>>;
-  readonly windUniform: TSLAccessor<d.F32, THREE.UniformNode<number>>;
-  readonly dampeningUniform: TSLAccessor<d.F32, THREE.UniformNode<number>>;
+  readonly stiffnessUniform: t3.TSLAccessor<d.F32, THREE.UniformNode<number>>;
+  readonly windUniform: t3.TSLAccessor<d.F32, THREE.UniformNode<number>>;
+  readonly dampeningUniform: t3.TSLAccessor<d.F32, THREE.UniformNode<number>>;
 
-  readonly vertexPositionBuffer: TSLAccessor<
-    d.WgslArray<d.Vec3f>,
-    THREE.StorageBufferNode
-  >;
-  readonly vertexForceBuffer: TSLAccessor<
-    d.WgslArray<d.Vec3f>,
-    THREE.StorageBufferNode
-  >;
-  readonly vertexParamsBuffer: TSLAccessor<
-    d.WgslArray<d.Vec3u>,
-    THREE.StorageBufferNode
-  >;
-  readonly springListBuffer: TSLAccessor<
-    d.WgslArray<d.U32>,
-    THREE.StorageBufferNode
-  >;
-  readonly springVertexIdBuffer: TSLAccessor<
-    d.WgslArray<d.Vec2u>,
-    THREE.StorageBufferNode
-  >;
-  readonly springRestLengthBuffer: TSLAccessor<
-    d.WgslArray<d.F32>,
-    THREE.StorageBufferNode
-  >;
-  readonly springForceBuffer: TSLAccessor<
-    d.WgslArray<d.Vec3f>,
-    THREE.StorageBufferNode
-  >;
+  readonly vertexPositionBuffer: TSLStorageAccessor<d.WgslArray<d.Vec3f>>;
+  readonly vertexForceBuffer: TSLStorageAccessor<d.WgslArray<d.Vec3f>>;
+  readonly vertexParamsBuffer: TSLStorageAccessor<d.WgslArray<d.Vec3u>>;
+  readonly springListBuffer: TSLStorageAccessor<d.WgslArray<d.U32>>;
+  readonly springVertexIdBuffer: TSLStorageAccessor<d.WgslArray<d.Vec2u>>;
+  readonly springRestLengthBuffer: TSLStorageAccessor<d.WgslArray<d.F32>>;
+  readonly springForceBuffer: TSLStorageAccessor<d.WgslArray<d.Vec3f>>;
 
   readonly computeSpringForces: TSL.ShaderNodeObject<THREE.ComputeNode>;
   readonly computeVertexForces: TSL.ShaderNodeObject<THREE.ComputeNode>;
 
-  constructor(
-    { sphereRadius, sphereUniform, spherePositionUniform }:
-      VerletSimulationOptions,
-  ) {
+  constructor({
+    sphereRadius,
+    sphereUniform,
+    spherePositionUniform,
+  }: VerletSimulationOptions) {
     this.vertices = [];
     this.springs = [];
     this.vertexColumns = [];
 
-    this.stiffnessUniform = fromTSL(TSL.uniform(0.2), { type: d.f32 });
-    this.windUniform = fromTSL(TSL.uniform(1.0), { type: d.f32 });
-    this.dampeningUniform = fromTSL(TSL.uniform(0.99), { type: d.f32 });
+    this.stiffnessUniform = t3.fromTSL(TSL.uniform(0.2), { type: d.f32 });
+    this.windUniform = t3.fromTSL(TSL.uniform(1.0), { type: d.f32 });
+    this.dampeningUniform = t3.fromTSL(TSL.uniform(0.99), { type: d.f32 });
 
     // this function sets up the geometry of the verlet system, a grid of vertices connected by springs
 
@@ -120,7 +108,7 @@ export class VerletSimulation {
       for (let y = 0; y <= clothNumSegmentsY; y++) {
         const posX = x * (clothWidth / clothNumSegmentsX) - clothWidth * 0.5;
         const posZ = y * (clothHeight / clothNumSegmentsY);
-        const isFixed = (y === 0) && ((x % 5) === 0); // make some of the top vertices' positions fixed
+        const isFixed = y === 0 && x % 5 === 0; // make some of the top vertices' positions fixed
         const vertex = addVerletVertex(posX, clothHeight * 0.5, posZ, isFixed);
         column.push(vertex);
       }
@@ -176,21 +164,24 @@ export class VerletSimulation {
       }
     }
 
-    this.vertexPositionBuffer = fromTSL(
+    this.vertexPositionBuffer = t3.fromTSL(
       TSL.instancedArray(vertexPositionArray, 'vec3'),
       { type: d.arrayOf(d.vec3f) },
     );
 
-    this.vertexForceBuffer = fromTSL(TSL.instancedArray(vertexCount, 'vec3'), {
-      type: d.arrayOf(d.vec3f),
-    });
+    this.vertexForceBuffer = t3.fromTSL(
+      TSL.instancedArray(vertexCount, 'vec3'),
+      {
+        type: d.arrayOf(d.vec3f),
+      },
+    );
 
-    this.vertexParamsBuffer = fromTSL(
+    this.vertexParamsBuffer = t3.fromTSL(
       TSL.instancedArray(vertexParamsArray, 'uvec3'),
       { type: d.arrayOf(d.vec3u) },
     );
 
-    this.springListBuffer = fromTSL(
+    this.springListBuffer = t3.fromTSL(
       TSL.instancedArray(new Uint32Array(springListArray), 'uint').setPBO(true),
       { type: d.arrayOf(d.u32) },
     );
@@ -211,50 +202,49 @@ export class VerletSimulation {
       );
     }
 
-    this.springVertexIdBuffer = fromTSL(
+    this.springVertexIdBuffer = t3.fromTSL(
       TSL.instancedArray(springVertexIdArray, 'uvec2').setPBO(true),
       { type: d.arrayOf(d.vec2u) },
     );
 
-    this.springRestLengthBuffer = fromTSL(
+    this.springRestLengthBuffer = t3.fromTSL(
       TSL.instancedArray(springRestLengthArray, 'float'),
       { type: d.arrayOf(d.f32) },
     );
 
-    this.springForceBuffer = fromTSL(
+    this.springForceBuffer = t3.fromTSL(
       TSL.instancedArray(springCount * 3, 'vec3').setPBO(true),
       { type: d.arrayOf(d.vec3f) },
     );
 
     // This sets up the compute shaders for the verlet simulation
     // There are two shaders that are executed for each simulation step
 
-    const instanceIndex = fromTSL(TSL.instanceIndex, { type: d.u32 });
-
     // 1. computeSpringForces:
     // This shader computes a force for each spring, depending on the distance between the two vertices connected by that spring and the targeted rest length
 
-    this.computeSpringForces = toTSL(() => {
+    this.computeSpringForces = t3.toTSL(() => {
       'use gpu';
 
-      if (instanceIndex.$ >= springCount) {
+      const idx = t3.instanceIndex.$;
+      if (idx >= springCount) {
         // compute Shaders are executed in groups of 64, so instanceIndex might be bigger than the amount of springs.
         // in that case, return.
         return;
       }
 
-      const vertexId = this.springVertexIdBuffer.$[instanceIndex.$];
-      const restLength = this.springRestLengthBuffer.$[instanceIndex.$];
+      const vertexId = this.springVertexIdBuffer.$[idx];
+      const restLength = this.springRestLengthBuffer.$[idx];
 
       const vertex0Position = this.vertexPositionBuffer.$[vertexId.x];
       const vertex1Position = this.vertexPositionBuffer.$[vertexId.y];
 
       const delta = vertex1Position.sub(vertex0Position);
       const dist = std.max(std.length(delta), 0.000001);
       const force = delta.mul(
-        (dist - restLength) * this.stiffnessUniform.$ * 0.5 / dist,
+        ((dist - restLength) * this.stiffnessUniform.$ * 0.5) / dist,
       );
-      this.springForceBuffer.$[instanceIndex.$] = d.vec3f(force);
+      this.springForceBuffer.$[idx] = d.vec3f(force);
     }).compute(springCount);
 
     // 2. computeVertexForces:
@@ -263,11 +253,9 @@ export class VerletSimulation {
     // Then it adds a gravital force, wind force, and the collision with the sphere.
     // In the end it adds the force to the vertex' position.
 
-    const time = fromTSL(TSL.time, { type: d.f32 });
-
-    this.computeVertexForces = toTSL(() => {
+    this.computeVertexForces = t3.toTSL(() => {
       'use gpu';
-      const idx = instanceIndex.$;
+      const idx = t3.instanceIndex.$;
 
       if (idx >= vertexCount) {
         // compute shaders are executed in groups of 64, so instanceIndex might be bigger than the amount of vertices.
@@ -295,27 +283,24 @@ export class VerletSimulation {
         const springId = this.springListBuffer.$[i];
         const springForce = this.springForceBuffer.$[springId];
         const springVertexIds = this.springVertexIdBuffer.$[springId];
-        const factor = std.select(
-          d.f32(-1),
-          d.f32(1),
-          springVertexIds.x === idx,
-        );
-        force = force.add(springForce.mul(factor));
+        const factor = std.select(-1, 1, springVertexIds.x === idx);
+        force = force.add(springForce.mul(d.f32(factor)));
       }
 
       // gravity
       force.y -= 0.00005;
 
       // wind
-      const noise = (triNoise3D(position, 1, time.$) - 0.2) * 0.0001;
+      const time = t3.time.$;
+      const noise = (triNoise3D(position, 1, time) - 0.2) * 0.0001;
       const windForce = noise * this.windUniform.$;
       force.z -= windForce;
 
       // collision with sphere
       const deltaSphere = position.add(force).sub(spherePositionUniform.$);
       const dist = std.length(deltaSphere);
       const sphereForce = deltaSphere.mul(
-        std.max(0, sphereRadius - dist) / dist * sphereUniform.$,
+        (std.max(0, sphereRadius - dist) / dist) * sphereUniform.$,
       );
       force = force.add(sphereForce);
 

apps/typegpu-docs/src/examples/threejs/simple/index.ts

@@ -1,5 +1,5 @@
 import * as THREE from 'three/webgpu';
-import { time, toTSL, uv } from '@typegpu/three';
+import * as t3 from '@typegpu/three';
 import { perlin3d } from '@typegpu/noise';
 import * as d from 'typegpu/data';
 import { tanh } from 'typegpu/std';
@@ -25,10 +25,10 @@ camera.position.z = 5;
 
 const material = new THREE.MeshBasicNodeMaterial();
 
-material.colorNode = toTSL(() => {
+material.colorNode = t3.toTSL(() => {
   'use gpu';
-  const coords = uv().$.mul(2);
-  const pattern = perlin3d.sample(d.vec3f(coords, time.$ * 0.2));
+  const coords = t3.uv().$.mul(2);
+  const pattern = perlin3d.sample(d.vec3f(coords, t3.time.$ * 0.2));
   return d.vec4f(tanh(pattern * 5), 0.2, 0.4, 1);
 });