From 60716f1626e7f7dec1fc380ea83926672bcd337a Mon Sep 17 00:00:00 2001 From: shadowcodex <1348053+shadowcodex@users.noreply.github.com> Date: Tue, 11 Aug 2026 19:34:45 -0500 Subject: [PATCH] shader dsl: add discard() for cut-out fragments --- packages/brometal/src/compiler/analyze.ts | 33 +++++ packages/brometal/src/compiler/emit-wgsl.ts | 3 + packages/brometal/src/compiler/ir.ts | 4 +- packages/brometal/src/compiler/optimize.ts | 2 + packages/brometal/src/dsl/builtins.ts | 23 ++++ packages/brometal/src/index.ts | 1 + packages/brometal/tests/discard.test.ts | 129 ++++++++++++++++++++ 7 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 packages/brometal/tests/discard.test.ts diff --git a/packages/brometal/src/compiler/analyze.ts b/packages/brometal/src/compiler/analyze.ts index ab136a9..fd759dc 100644 --- a/packages/brometal/src/compiler/analyze.ts +++ b/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`); } diff --git a/packages/brometal/src/compiler/emit-wgsl.ts b/packages/brometal/src/compiler/emit-wgsl.ts index c60aac2..11c51b4 100644 --- a/packages/brometal/src/compiler/emit-wgsl.ts +++ b/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)};`); diff --git a/packages/brometal/src/compiler/ir.ts b/packages/brometal/src/compiler/ir.ts index 15e3206..41dd6ca 100644 --- a/packages/brometal/src/compiler/ir.ts +++ b/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; diff --git a/packages/brometal/src/compiler/optimize.ts b/packages/brometal/src/compiler/optimize.ts index 6d64161..e3df2b4 100644 --- a/packages/brometal/src/compiler/optimize.ts +++ b/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': diff --git a/packages/brometal/src/dsl/builtins.ts b/packages/brometal/src/dsl/builtins.ts index 94689e3..07ac392 100644 --- a/packages/brometal/src/dsl/builtins.ts +++ b/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(incident: T, normal: T): T; export function reflect(): never { return gpuOnly('reflect'); diff --git a/packages/brometal/src/index.ts b/packages/brometal/src/index.ts index 98f9573..cd7f02e 100644 --- a/packages/brometal/src/index.ts +++ b/packages/brometal/src/index.ts @@ -22,6 +22,7 @@ export { clamp, cos, cross, + discard, distance, dot, exp, diff --git a/packages/brometal/tests/discard.test.ts b/packages/brometal/tests/discard.test.ts new file mode 100644 index 0000000..6c0c891 --- /dev/null +++ b/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;'); + }); +});