feat: Allow passing raw wgsl implementation to shell à la template literal (#1128)

Author: aleksanderkatan

apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx

@@ -19,17 +19,27 @@ These functions can reference outside resources, like other user-defined or help
 ## Creating a function
 
 Functions are constructed by first defining their shells, which specify their inputs and outputs.
-Then the actual WGSL implementation is passed in as an argument to a shell invocation.
+Then the actual WGSL implementation is passed in as an argument to a shell invocation. If the code string is a template literal, you can omit the parentheses, which may result in a more compact Biome/Prettier formatting.
 
 The following code defines a function that accepts one argument and returns one value.
 
 ```ts
 const getGradientColor = tgpu['~unstable']
-  .fn({ ratio: d.f32 }, d.vec4f)(`{
-    let color = mix(vec4f(0.769, 0.392, 1.0, 1), d.vec4f(0.114, 0.447, 0.941, 1), ratio);
+  .fn(
+    { ratio: d.f32 },
+    d.vec4f,
+  )(`{
+    let color = mix(vec4f(0.769, 0.392, 1.0, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
     return color;
   }`)
   .$name('getGradientColor');
+
+// or
+
+const getGradientColor = tgpu['~unstable'].fn({ ratio: d.f32 }, d.vec4f)`{
+  let color = mix(vec4f(0.769, 0.392, 1.0, 1), vec4f(0.114, 0.447, 0.941, 1), ratio);
+  return color;
+}`.$name('getGradientColor');
 ```
 
 If you're using Visual Studio Code, you can use an [extension](https://marketplace.visualstudio.com/items?itemName=ggsimm.wgsl-literal) that brings syntax highlighting to the code fragments marked with `/* wgsl */` comments.
@@ -40,18 +50,16 @@ Functions can use external resources passed inside a record via the `$uses` meth
 Externals can be any value or TypeGPU resource that can be resolved to WGSL (functions, buffer usages, slots, accessors, constants, variables, declarations, vectors, matrices, textures, samplers etc.).
 
 ```ts
-const getBlue = tgpu['~unstable']
-  .fn({}, d.vec4f)(`{ 
-    return vec4f(0.114, 0.447, 0.941, 1); 
-  }`);
+const getBlue = tgpu['~unstable'].fn({}, d.vec4f)`{ 
+  return vec4f(0.114, 0.447, 0.941, 1); 
+}`;
 
 const purple = d.vec4f(0.769, 0.392, 1.0, 1);
 
-const getGradientColor = tgpu['~unstable']
-  .fn({ ratio: d.f32 }, d.vec4f)(`{
-    let color = mix(purple, getBlue(), ratio);
-    return color;
-  }`)
+const getGradientColor = tgpu['~unstable'].fn({ ratio: d.f32 }, d.vec4f)`{
+  let color = mix(purple, getBlue(), ratio);
+  return color;
+}`
   .$uses({ purple, getBlue })
   .$name('getGradientColor');
 ```
@@ -79,11 +87,10 @@ Defining entry functions is similar to regular ones, but is done through dedicat
 They can be passed to root-defined pipelines and they accept special arguments like builtins (`d.builtin`) and decorated data (`d.location`).
 
 ```ts
-const mainVertex = tgpu['~unstable']
-  .vertexFn({
-    in: { vertexIndex: d.builtin.vertexIndex },
-    out: { outPos: d.builtin.position, uv: d.vec2f },
-  })(/* wgsl */ `{
+const mainVertex = tgpu['~unstable'].vertexFn({
+  in: { vertexIndex: d.builtin.vertexIndex },
+  out: { outPos: d.builtin.position, uv: d.vec2f },
+}) /* wgsl */`{
     var pos = array<vec2f, 3>(
       vec2(0.0, 0.5),
       vec2(-0.5, -0.5),
@@ -97,13 +104,14 @@ const mainVertex = tgpu['~unstable']
     );
 
     return Out(vec4f(pos[in.vertexIndex], 0.0, 1.0), uv[in.vertexIndex]);
-  }`);
+  }`;
 
-const mainFragment = tgpu['~unstable']
-  .fragmentFn({ in: { uv: d.vec2f }, out: d.vec4f })(/* wgsl */ `{
+const mainFragment = tgpu['~unstable'].fragmentFn({
+  in: { uv: d.vec2f },
+  out: d.vec4f,
+}) /* wgsl */`{
     return getGradientColor((in.uv[0] + in.uv[1]) / 2);
-  }`)
-  .$uses({ getGradientColor });
+  }`.$uses({ getGradientColor });
 ```
 
 When entry function inputs or outputs are specified as objects containing builtins and inter-stage variables, the WGSL implementations need to access these arguments as passed in via structs. 

apps/typegpu-docs/src/content/examples/rendering/3d-fish/render.ts

@@ -86,13 +86,12 @@ export const vertexShader = tgpu['~unstable']
   })
   .$name('vertexShader');
 
-const sampleTexture = tgpu['~unstable']
-  .fn(
-    { uv: d.vec2f },
-    d.vec4f,
-  )(/* wgsl */ `{
-    return textureSample(layout.$.modelTexture, layout.$.sampler, uv);
-  }`)
+const sampleTexture = tgpu['~unstable'].fn(
+  { uv: d.vec2f },
+  d.vec4f,
+) /* wgsl */`{
+  return textureSample(layout.$.modelTexture, layout.$.sampler, uv);
+}`
   .$uses({ layout })
   .$name('sampleTexture');
 

apps/typegpu-docs/src/content/examples/rendering/box-raytracing/index.ts

@@ -106,11 +106,10 @@ const renderBindGroup = root.createBindGroup(renderLayout, {
 
 // functions
 
-const getBoxIntersection = tgpu['~unstable']
-  .fn(
-    { boundMin: d.vec3f, boundMax: d.vec3f, ray: RayStruct },
-    IntersectionStruct,
-  )(/* wgsl */ `{
+const getBoxIntersection = tgpu['~unstable'].fn(
+  { boundMin: d.vec3f, boundMax: d.vec3f, ray: RayStruct },
+  IntersectionStruct,
+) /* wgsl */`{
   var tMin: f32;
   var tMax: f32;
   var tMinY: f32;
@@ -167,15 +166,14 @@ const getBoxIntersection = tgpu['~unstable']
   }
 
   return IntersectionStruct(tMin > 0 && tMax > 0, tMin, tMax);
-}`)
+}`
   .$uses({ IntersectionStruct })
   .$name('box_intersection');
 
-const vertexFunction = tgpu['~unstable']
-  .vertexFn({
-    in: { vertexIndex: d.builtin.vertexIndex },
-    out: { outPos: d.builtin.position },
-  })(/* wgsl */ `{
+const vertexFunction = tgpu['~unstable'].vertexFn({
+  in: { vertexIndex: d.builtin.vertexIndex },
+  out: { outPos: d.builtin.position },
+}) /* wgsl */`{
   var pos = array<vec2f, 6>(
     vec2<f32>( 1,  1),
     vec2<f32>( 1, -1),
@@ -186,17 +184,15 @@ const vertexFunction = tgpu['~unstable']
   );
 
   return Out(vec4f(pos[in.vertexIndex], 0, 1));
-}`)
-  .$name('vertex_main');
+}`.$name('vertex_main');
 
 const boxSizeAccessor = tgpu['~unstable'].accessor(d.u32);
 const canvasDimsAccessor = tgpu['~unstable'].accessor(CanvasDimsStruct);
 
-const fragmentFunction = tgpu['~unstable']
-  .fragmentFn({
-    in: { position: d.builtin.position },
-    out: d.vec4f,
-  })(/* wgsl */ `{
+const fragmentFunction = tgpu['~unstable'].fragmentFn({
+  in: { position: d.builtin.position },
+  out: d.vec4f,
+}) /* wgsl */`{
   let minDim = f32(min(canvasDims.width, canvasDims.height));
 
   var ray: RayStruct;
@@ -246,7 +242,7 @@ const fragmentFunction = tgpu['~unstable']
   }
 
   return color;
-}`)
+}`
   .$uses({
     ...renderLayout.bound,
     RayStruct,
@@ -266,12 +262,9 @@ const fragmentFunction = tgpu['~unstable']
 const pipeline = root['~unstable']
   .with(
     boxSizeAccessor,
-    tgpu['~unstable']
-      .fn(
-        [],
-        d.u32,
-      )('() -> u32 { return boxSize; }')
-      .$uses({ boxSize: boxSizeUniform }),
+    tgpu['~unstable'].fn([], d.u32)`() -> u32 { return boxSize; }`.$uses({
+      boxSize: boxSizeUniform,
+    }),
   )
   .with(canvasDimsAccessor, canvasDimsUniform)
   .withVertex(vertexFunction, {})

apps/typegpu-docs/src/content/examples/simple/triangle/index.ts

@@ -16,21 +16,20 @@ context.configure({
   alphaMode: 'premultiplied',
 });
 
-const getGradientColor = tgpu['~unstable']
-  .fn(
-    { ratio: d.f32 },
-    d.vec4f,
-  )(/* wgsl */ `{
+const getGradientColor = tgpu['~unstable'].fn(
+  { ratio: d.f32 },
+  d.vec4f,
+) /* wgsl */`{
     let color = mix(purple, blue, ratio);
     return color;
-  }`)
+  }`
   .$uses({ purple, blue })
   .$name('getGradientColor');
 
 const mainVertex = tgpu['~unstable'].vertexFn({
   in: { vertexIndex: d.builtin.vertexIndex },
   out: { outPos: d.builtin.position, uv: d.vec2f },
-})(/* wgsl */ `{
+}) /* wgsl */`{
     var pos = array<vec2f, 3>(
       vec2(0.0, 0.5),
       vec2(-0.5, -0.5),
@@ -44,16 +43,14 @@ const mainVertex = tgpu['~unstable'].vertexFn({
     );
 
     return Out(vec4f(pos[in.vertexIndex], 0.0, 1.0), uv[in.vertexIndex]);
-  }`);
+  }`;
 
-const mainFragment = tgpu['~unstable']
-  .fragmentFn({
-    in: { uv: d.vec2f },
-    out: d.vec4f,
-  })(/* wgsl */ `{
+const mainFragment = tgpu['~unstable'].fragmentFn({
+  in: { uv: d.vec2f },
+  out: d.vec4f,
+}) /* wgsl */`{
     return getGradientColor((in.uv[0] + in.uv[1]) / 2);
-  }`)
-  .$uses({ getGradientColor });
+  }`.$uses({ getGradientColor });
 
 const pipeline = root['~unstable']
   .withVertex(mainVertex, {})

apps/typegpu-docs/src/content/examples/simulation/confetti/index.ts

@@ -90,75 +90,67 @@ const dataLayout = tgpu
 const rotate = tgpu['~unstable'].fn(
   { v: d.vec2f, angle: d.f32 },
   d.vec2f,
-)(/* wgsl */ `{
-    let pos = vec2(
-      (v.x * cos(angle)) - (v.y * sin(angle)),
-      (v.x * sin(angle)) + (v.y * cos(angle))
-    );
+) /* wgsl */`{
+  let pos = vec2(
+    (v.x * cos(angle)) - (v.y * sin(angle)),
+    (v.x * sin(angle)) + (v.y * cos(angle))
+  );
+
+  return pos;
+}`;
 
-    return pos;
+const mainVert = tgpu['~unstable'].vertexFn({
+  in: {
+    tilt: d.f32,
+    angle: d.f32,
+    color: d.vec4f,
+    center: d.vec2f,
+    index: d.builtin.vertexIndex,
+  },
+  out: VertexOutput,
+}) /* wgsl */`{
+  let width = in.tilt;
+  let height = in.tilt / 2;
+
+  var pos = rotate(array<vec2f, 4>(
+    vec2f(0, 0),
+    vec2f(width, 0),
+    vec2f(0, height),
+    vec2f(width, height),
+  )[in.index] / 350, in.angle) + in.center;
+
+  if (canvasAspectRatio < 1) {
+    pos.x /= canvasAspectRatio;
+  } else {
+    pos.y *= canvasAspectRatio;
   }
-`);
-
-const mainVert = tgpu['~unstable']
-  .vertexFn({
-    in: {
-      tilt: d.f32,
-      angle: d.f32,
-      color: d.vec4f,
-      center: d.vec2f,
-      index: d.builtin.vertexIndex,
-    },
-    out: VertexOutput,
-  })(
-    /* wgsl */ `{
-    let width = in.tilt;
-    let height = in.tilt / 2;
-
-    var pos = rotate(array<vec2f, 4>(
-      vec2f(0, 0),
-      vec2f(width, 0),
-      vec2f(0, height),
-      vec2f(width, height),
-    )[in.index] / 350, in.angle) + in.center;
-
-    if (canvasAspectRatio < 1) {
-      pos.x /= canvasAspectRatio;
-    } else {
-      pos.y *= canvasAspectRatio;
-    }
 
-    return Out(vec4f(pos, 0.0, 1.0), in.color);
-  }`,
-  )
-  .$uses({
-    rotate,
-    canvasAspectRatio: canvasAspectRatioUniform,
-  });
+  return Out(vec4f(pos, 0.0, 1.0), in.color);
+}`.$uses({
+  rotate,
+  canvasAspectRatio: canvasAspectRatioUniform,
+});
 
 const mainFrag = tgpu['~unstable'].fragmentFn({
   in: VertexOutput,
   out: d.vec4f,
-})(/* wgsl */ `{
-    return in.color;
-  }`);
-
-const mainCompute = tgpu['~unstable']
-  .computeFn({ in: { gid: d.builtin.globalInvocationId }, workgroupSize: [1] })(
-    /* wgsl */ `{
-    let index = in.gid.x;
-    if index == 0 {
-      time += deltaTime;
-    }
-    let phase = (time / 300) + particleData[index].seed;
-    particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
-  }`,
-  )
-  .$uses({
-    particleData: particleDataStorage,
-    deltaTime: deltaTimeUniform,
-    time: timeStorage,
-  });
+}) /* wgsl */`{ return in.color; }`;
+
+const mainCompute = tgpu['~unstable'].computeFn({
+  in: { gid: d.builtin.globalInvocationId },
+  workgroupSize: [1],
+}) /* wgsl */`{
+  let index = in.gid.x;
+  if index == 0 {
+    time += deltaTime;
+  }
+  let phase = (time / 300) + particleData[index].seed;
+  particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
+}`.$uses({
+  particleData: particleDataStorage,
+  deltaTime: deltaTimeUniform,
+  time: timeStorage,
+});
 
 // pipelines