shader dsl: add discard() for cut-out fragments

packages/brometal/src/compiler/analyze.ts

@@ -719,6 +719,9 @@ function lowerMutation(
   expr: ts.Expression,
   options: { topLevel: boolean },
 ): IrStmt {
+  if (ts.isCallExpression(expr) && ts.isIdentifier(expr.expression) && expr.expression.text === 'discard') {
+    return lowerDiscard(ctx, expr);
+  }
   // storageWrite is the only call allowed in statement position — it is the
   // sole way a stage produces output other than by returning.
   if (
@@ -821,6 +824,29 @@ function lowerMutation(
   );
 }
 
+/**
+ * discard() is a statement and not a value. It is the only function call that can
+ * be an expression statement.
+ *
+ * Only fragment() can use it. The vertex stage can call a helper function, and a
+ * discard has no meaning there. Therefore a helper function cannot use it.
+ */
+function lowerDiscard(ctx: StageContext, expr: ts.CallExpression): IrStmt {
+  if (ctx.stage !== 'fragment') {
+    throw errorAt(
+      ctx.sourceFile,
+      expr,
+      ctx.stage === 'vertex'
+        ? `discard() is only valid in fragment() — there is no fragment to throw away in the vertex stage`
+        : `discard() is only valid in fragment(), not in ${ctx.ownerLabel} — a helper may be called from vertex(), where discarding is meaningless`,
+    );
+  }
+  if (expr.arguments.length > 0) {
+    throw errorAt(ctx.sourceFile, expr, `discard() takes no arguments`);
+  }
+  return { kind: 'discard' };
+}
+
 function requireMutableFloat(
   ctx: StageContext,
   scope: Scope,
@@ -1093,6 +1119,13 @@ function lowerCall(ctx: StageContext, scope: Scope, node: ts.CallExpression): Ir
     throw errorAt(ctx.sourceFile, node, `only vec constructors, intrinsics, and .add/.sub/.mul/.div/.scale method calls are supported`);
   }
   const callee = node.expression.text;
+  if (callee === 'discard') {
+    throw errorAt(
+      ctx.sourceFile,
+      node,
+      `discard() produces no value — call it as its own statement, e.g. \`if (alpha < 0.5) { discard(); }\``,
+    );
+  }
   if (scope.lookup(callee) !== undefined) {
     throw errorAt(ctx.sourceFile, node, `'${callee}' is not callable in shader code`);
   }

packages/brometal/src/compiler/emit-wgsl.ts

@@ -214,6 +214,9 @@ function emitStatements(lines: string[], statements: IrStmt[], ctx: EmitContext,
           `${indent}${statement.buffer}[u32(${emitExpr(statement.index, ctx, 0)})] = ${emitExpr(statement.value, ctx, 0)};`,
         );
         break;
+      case 'discard':
+        lines.push(`${indent}discard;`);
+        break;
       case 'return':
         if (ctx.stage === 'vertex') {
           lines.push(`${indent}bm_out.bm_position = ${emitExpr(statement.expr, ctx, 0)};`);

packages/brometal/src/compiler/ir.ts

@@ -45,7 +45,9 @@ export type IrStmt =
       body: IrStmt[];
     }
   | { kind: 'storageWrite'; buffer: string; index: IrExpr; value: IrExpr }
-  | { kind: 'return'; expr: IrExpr };
+  | { kind: 'return'; expr: IrExpr }
+  /** Fragment stage only. Discards the fragment. The GPU writes no colour and no depth. */
+  | { kind: 'discard' };
 
 export interface HelperParam {
   name: string;

packages/brometal/src/compiler/optimize.ts

@@ -25,6 +25,8 @@ function foldStmt(statement: IrStmt): IrStmt {
       return { ...statement, expr: foldExpr(statement.expr) };
     case 'assign':
       return { ...statement, expr: foldExpr(statement.expr) };
+    case 'discard':
+      return statement;
     case 'return':
       return { ...statement, expr: foldExpr(statement.expr) };
     case 'storageWrite':

packages/brometal/src/dsl/builtins.ts

@@ -53,6 +53,29 @@ export function texture(): Vec4 {
   return gpuOnly('texture');
 }
 
+/**
+ * Discards the current fragment. The GPU writes no colour and no depth for it.
+ *
+ * Use discard() only in the fragment stage. Put it in an `if` statement:
+ *
+ * ```ts
+ * fragment({ uAtlas, uCutoff }, { vUv }) {
+ *   const texel = texture(uAtlas, vUv);
+ *   if (texel.w < uCutoff) { discard(); }
+ *   return texel;
+ * }
+ * ```
+ *
+ * This function makes cut-out sprites possible. Each fragment that remains is
+ * fully opaque, so the program writes depth, and it can draw the sprites in any
+ * order. Without discard(), a sprite with a transparent edge must blend, a blended
+ * program cannot write depth, and the application must sort the sprites from back
+ * to front on the CPU in each frame.
+ */
+export function discard(): void {
+  gpuOnly('discard');
+}
+
 export function reflect<T extends Vec2 | Vec3 | Vec4>(incident: T, normal: T): T;
 export function reflect(): never {
   return gpuOnly('reflect');

packages/brometal/src/index.ts

@@ -22,6 +22,7 @@ export {
   clamp,
   cos,
   cross,
+  discard,
   distance,
   dot,
   exp,

packages/brometal/tests/discard.test.ts

@@ -0,0 +1,129 @@
+import { describe, expect, it } from 'vitest';
+
+/**
+ * Ported from the cut-out sprite branch, minus its GLSL assertions: WebGL2 was
+ * removed from the library after that branch was written, so WGSL is the only
+ * emitter left to check.
+ */
+import { compileShaderSource } from '../src/compiler/compile.js';
+
+function compile(source: string) {
+  return compileShaderSource('test.shader.ts', source);
+}
+
+/** A cut-out sprite shader: the whole point of `discard()`. */
+const CUTOUT_SHADER = `
+import { shader, discard, texture, vec4 } from 'brometal';
+
+export default shader({
+  attributes: { aPosition: 'vec3', aUv: 'vec2' },
+  uniforms: { uAtlas: 'sampler2D', uCutoff: 'float' },
+  varyings: { vUv: 'vec2' },
+
+  vertex({ aPosition, aUv }, _u, v) {
+    v.vUv = aUv;
+    return vec4(aPosition, 1);
+  },
+
+  fragment({ uAtlas, uCutoff }, { vUv }) {
+    const texel = texture(uAtlas, vUv);
+    if (texel.w < uCutoff) {
+      discard();
+    }
+    return vec4(texel.xyz, 1);
+  },
+});
+`;
+
+function shaderWith(fragment: string, vertex?: string): string {
+  return `
+import { shader, discard, vec4 } from 'brometal';
+export default shader({
+  attributes: { aPosition: 'vec3' },
+  uniforms: { uCutoff: 'float' },
+  varyings: { vFade: 'float' },
+  vertex: ${vertex ?? `({ aPosition }, _u, v) => { v.vFade = 1; return vec4(aPosition, 1); }`},
+  fragment: ${fragment},
+});
+`;
+}
+
+describe('discard()', () => {
+  it('emits a WGSL discard inside the branch', () => {
+    const wgsl = compile(CUTOUT_SHADER).wgslSrc ?? '';
+    expect(wgsl).toContain('if (texel.w < bm_u.uCutoff) {');
+    expect(wgsl).toContain('discard;');
+  });
+
+  it('survives a prod build with constant folding', () => {
+    const folded = compileShaderSource('test.shader.ts', CUTOUT_SHADER, { optimize: true });
+    expect(folded.wgslSrc).toContain('discard;');
+  });
+
+  it('keeps a varying alive when only the discard condition reads it', () => {
+    // vFade is never used for colour — only to decide whether to discard. It
+    // must not be pruned as a dead varying.
+    const source = shaderWith(
+      `(_u, { vFade }) => { if (vFade < 0.5) { discard(); } return vec4(1, 1, 1, 1); }`,
+    );
+    const compiled = compileShaderSource('test.shader.ts', source, { optimize: true });
+    expect(compiled.wgslSrc).toContain('vFade');
+    expect(compiled.wgslSrc).toContain('discard;');
+    expect(Object.keys(compiled.varyings)).toContain('vFade');
+  });
+
+  it('emits discard from an else branch', () => {
+    const source = shaderWith(
+      `(_u, { vFade }) => { let keep = 0; if (vFade < 0.5) { keep = 1; } else { discard(); } return vec4(keep, 0, 0, 1); }`,
+    );
+    const glsl = compile(source).wgslSrc;
+    expect(glsl).toContain('} else {');
+    expect(glsl).toContain('discard;');
+  });
+
+  it('rejects discard() in the vertex stage', () => {
+    const source = shaderWith(
+      `(_u, { vFade }) => vec4(vFade, vFade, vFade, 1)`,
+      `({ aPosition }, _u, v) => { v.vFade = 1; discard(); return vec4(aPosition, 1); }`,
+    );
+    expect(() => compile(source)).toThrow(/discard\(\) is only valid in fragment\(\)/);
+  });
+
+  it('rejects discard() inside a helper', () => {
+    const source = `
+import { shader, discard, vec4 } from 'brometal';
+
+function cut(a: number): number {
+  if (a < 0.5) {
+    discard();
+  }
+  return a;
+}
+
+export default shader({
+  attributes: { aPosition: 'vec3' },
+  varyings: { vFade: 'float' },
+  vertex({ aPosition }, _u, v) { v.vFade = 1; return vec4(aPosition, 1); },
+  fragment(_u, { vFade }) { return vec4(cut(vFade), 0, 0, 1); },
+});
+`;
+    expect(() => compile(source)).toThrow(/discard\(\) is only valid in fragment\(\)/);
+  });
+
+  it('rejects discard() used as a value', () => {
+    const source = shaderWith(`(_u, { vFade }) => vec4(discard(), vFade, vFade, 1)`);
+    expect(() => compile(source)).toThrow(/discard\(\) produces no value/);
+  });
+
+  it('rejects discard() with arguments', () => {
+    const source = shaderWith(
+      `(_u, { vFade }) => { if (vFade < 0.5) { discard(1); } return vec4(1, 1, 1, 1); }`,
+    );
+    expect(() => compile(source)).toThrow(/discard\(\) takes no arguments/);
+  });
+
+  it('allows discard() at the top level of fragment(), before the return', () => {
+    const source = shaderWith(`(_u, { vFade }) => { discard(); return vec4(vFade, 0, 0, 1); }`);
+    expect(compile(source).wgslSrc).toContain('discard;');
+  });
+});