diff --git a/apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx b/apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx index 3ed119308..a7c2c4fc2 100644 --- a/apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx +++ b/apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx @@ -323,7 +323,7 @@ function manhattanDistance(a: d.v3f, b: d.v3f) { const dy = std.abs(a.y - b.y); const dz = std.abs(a.z - b.z); - return std.max(dx, std.max(dy, dz)); + return std.max(dx, dy, dz); } ``` diff --git a/apps/typegpu-docs/src/examples/rendering/jelly-slider/utils.ts b/apps/typegpu-docs/src/examples/rendering/jelly-slider/utils.ts index d4b96cf9b..a3e3d0e22 100644 --- a/apps/typegpu-docs/src/examples/rendering/jelly-slider/utils.ts +++ b/apps/typegpu-docs/src/examples/rendering/jelly-slider/utils.ts @@ -33,8 +33,8 @@ export const intersectBox = ( const tMinVec = std.min(t1, t2); const tMaxVec = std.max(t1, t2); - const tMin = std.max(std.max(tMinVec.x, tMinVec.y), tMinVec.z); - const tMax = std.min(std.min(tMaxVec.x, tMaxVec.y), tMaxVec.z); + const tMin = std.max(tMinVec.x, tMinVec.y, tMinVec.z); + const tMax = std.min(tMaxVec.x, tMaxVec.y, tMaxVec.z); const result = BoxIntersection(); result.hit = tMax >= tMin && tMax >= 0.0; diff --git a/apps/typegpu-docs/src/examples/rendering/jelly-switch/utils.ts b/apps/typegpu-docs/src/examples/rendering/jelly-switch/utils.ts index c7277a7f3..ed5b0fa35 100644 --- a/apps/typegpu-docs/src/examples/rendering/jelly-switch/utils.ts +++ b/apps/typegpu-docs/src/examples/rendering/jelly-switch/utils.ts @@ -32,8 +32,8 @@ export const intersectBox = ( const tMinVec = std.min(t1, t2); const tMaxVec = std.max(t1, t2); - const tMin = std.max(std.max(tMinVec.x, tMinVec.y), tMinVec.z); - const tMax = std.min(std.min(tMaxVec.x, tMaxVec.y), tMaxVec.z); + const tMin = std.max(tMinVec.x, tMinVec.y, tMinVec.z); + const tMax = std.min(tMaxVec.x, tMaxVec.y, tMaxVec.z); const result = BoxIntersection(); result.hit = tMax >= tMin && tMax >= 0.0; diff --git a/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts b/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts index 2a6e2cc00..995c5461a 100644 --- a/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts +++ b/apps/typegpu-docs/src/examples/simulation/slime-mold-3d/index.ts @@ -365,8 +365,8 @@ const rayBoxIntersection = ( const t1 = boxMax.sub(rayOrigin).mul(invDir); const tmin = std.min(t0, t1); const tmax = std.max(t0, t1); - const tNear = std.max(std.max(tmin.x, tmin.y), tmin.z); - const tFar = std.min(std.min(tmax.x, tmax.y), tmax.z); + const tNear = std.max(tmin.x, tmin.y, tmin.z); + const tFar = std.min(tmax.x, tmax.y, tmax.z); const hit = tFar >= tNear && tFar >= 0; return RayBoxResult({ tNear, tFar, hit }); }; diff --git a/packages/typegpu/src/std/numeric.ts b/packages/typegpu/src/std/numeric.ts index de02cf4e7..17526c9af 100644 --- a/packages/typegpu/src/std/numeric.ts +++ b/packages/typegpu/src/std/numeric.ts @@ -14,7 +14,7 @@ import { i32, u32, } from '../data/numeric.ts'; -import { snip } from '../data/snippet.ts'; +import { snip, Snippet } from '../data/snippet.ts'; import { abstruct } from '../data/struct.ts'; import { vec2f, @@ -56,6 +56,45 @@ import { mul, sub } from './operators.ts'; type NumVec = AnyNumericVecInstance; +// helpers + +const unaryIdentitySignature = (arg: AnyData) => { + return { + argTypes: [arg], + returnType: arg, + }; +}; + +const variadicUnifySignature = (...args: AnyData[]) => { + const uargs = unify(args) ?? args; + return ({ + argTypes: uargs, + returnType: uargs[0] as AnyData, + }); +}; + +function variadicReduce(fn: (a: T, b: T) => T) { + return (fst: T, ...rest: T[]): T => { + let acc = fst; + for (const r of rest) { + acc = fn(acc, r); + } + return acc; + }; +} + +function variadicStitch(wrapper: string) { + return (fst: Snippet, ...rest: Snippet[]): string => { + let acc = stitch`${fst}`; + for (const r of rest) { + acc = stitch`${wrapper}(${acc}, ${r})`; + } + return acc; + }; +} + +// std + function cpuAbs(value: number): number; function cpuAbs(value: T): T; function cpuAbs(value: T): T { @@ -65,13 +104,6 @@ function cpuAbs(value: T): T { return VectorOps.abs[value.kind](value) as T; } -const unaryIdentitySignature = (arg: AnyData) => { - return { - argTypes: [arg], - returnType: arg, - }; -}; - export const abs = dualImpl({ name: 'abs', signature: unaryIdentitySignature, @@ -231,10 +263,7 @@ function cpuClamp(value: T, low: T, high: T): T { export const clamp = dualImpl({ name: 'clamp', - signature: (...args) => { - const uargs = unify(args) ?? args; - return { argTypes: uargs, returnType: uargs[0] }; - }, + signature: variadicUnifySignature, normalImpl: cpuClamp, codegenImpl: (value, low, high) => stitch`clamp(${value}, ${low}, ${high})`, }); @@ -767,17 +796,16 @@ function cpuMax(a: T, b: T): T { return VectorOps.max[a.kind](a, b as NumVec) as T; } +type VariadicOverload = { + (fst: number, ...rest: number[]): number; + (fst: T, ...rest: T[]): T; +}; + export const max = dualImpl({ name: 'max', - signature: (...args) => { - const uargs = unify(args) ?? args; - return ({ - argTypes: uargs, - returnType: uargs[0], - }); - }, - normalImpl: cpuMax, - codegenImpl: (a, b) => stitch`max(${a}, ${b})`, + signature: variadicUnifySignature, + normalImpl: variadicReduce(cpuMax) as VariadicOverload, + codegenImpl: variadicStitch('max'), }); function cpuMin(a: number, b: number): number; @@ -791,15 +819,9 @@ function cpuMin(a: T, b: T): T { export const min = dualImpl({ name: 'min', - signature: (...args) => { - const uargs = unify(args) ?? args; - return ({ - argTypes: uargs, - returnType: uargs[0], - }); - }, - normalImpl: cpuMin, - codegenImpl: (a, b) => stitch`min(${a}, ${b})`, + signature: variadicUnifySignature, + normalImpl: variadicReduce(cpuMin) as VariadicOverload, + codegenImpl: variadicStitch('min'), }); function cpuMix(e1: number, e2: number, e3: number): number; @@ -828,13 +850,7 @@ function cpuMix( export const mix = dualImpl({ name: 'mix', - signature: (...args) => { - const uargs = unify(args) ?? args; - return ({ - argTypes: uargs, - returnType: uargs[0], - }); - }, + signature: variadicUnifySignature, normalImpl: cpuMix, codegenImpl: (e1, e2, e3) => stitch`mix(${e1}, ${e2}, ${e3})`, }); diff --git a/packages/typegpu/tests/std/numeric/max.test.ts b/packages/typegpu/tests/std/numeric/max.test.ts new file mode 100644 index 000000000..f39f66729 --- /dev/null +++ b/packages/typegpu/tests/std/numeric/max.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it } from 'vitest'; +import tgpu from '../../../src/index.ts'; +import * as d from '../../../src/data/index.ts'; +import * as std from '../../../src/std/index.ts'; + +describe('max', () => { + it('acts as identity when called with one argument', () => { + const myMax = tgpu.fn([d.f32], d.f32)((a: number) => { + 'use gpu'; + return std.max(a); + }); + + expect(myMax(6)).toBe(6); + expect(tgpu.resolve([myMax])).toMatchInlineSnapshot(` + "fn myMax(a: f32) -> f32 { + return a; + }" + `); + }); + + it('works with two arguments', () => { + const myMax = tgpu.fn([d.f32, d.f32], d.f32)((a, b) => { + 'use gpu'; + return std.max(a, b); + }); + + expect(myMax(1, 2)).toBe(2); + expect(tgpu.resolve([myMax])).toMatchInlineSnapshot(` + "fn myMax(a: f32, b: f32) -> f32 { + return max(a, b); + }" + `); + }); + + it('works with multiple arguments', () => { + const myMax = tgpu.fn([d.f32, d.f32, d.f32, d.f32], d.f32)( + (a, b, c, d) => { + 'use gpu'; + return std.max(a, b, c, d); + }, + ); + + expect(myMax(2, 1, 4, 5)).toBe(5); + expect(tgpu.resolve([myMax])).toMatchInlineSnapshot(` + "fn myMax(a: f32, b: f32, c: f32, d2: f32) -> f32 { + return max(max(max(a, b), c), d2); + }" + `); + }); + + it('unifies arguments', () => { + const myMax = tgpu.fn([], d.f32)(() => { + 'use gpu'; + const a = d.u32(9); + const b = d.i32(1); + const c = d.f32(4); + return std.max(a, b, 3.3, c, 7); + }); + + expect(myMax()).toBe(9); + expect(tgpu.resolve([myMax])).toMatchInlineSnapshot(` + "fn myMax() -> f32 { + const a = 9u; + const b = 1i; + const c = 4f; + return max(max(max(max(f32(a), f32(b)), 3.3f), c), 7f); + }" + `); + }); + + it('works with vectors', () => { + const myMax = tgpu.fn([d.vec3u, d.vec3u], d.vec3u)((a, b) => { + 'use gpu'; + return std.max(a, b); + }); + + expect(myMax(d.vec3u(1, 2, 3), d.vec3u(3, 2, 1))) + .toStrictEqual(d.vec3u(3, 2, 3)); + expect(tgpu.resolve([myMax])).toMatchInlineSnapshot(` + "fn myMax(a: vec3u, b: vec3u) -> vec3u { + return max(a, b); + }" + `); + }); + + it('cannot be called with invalid arguments', () => { + // @ts-expect-error + (() => std.max()); + // @ts-expect-error + (() => std.max(1, d.vec2f())); + // @ts-expect-error + (() => std.max(d.vec3f(), d.vec2f())); + }); +}); diff --git a/packages/typegpu/tests/std/numeric/min.test.ts b/packages/typegpu/tests/std/numeric/min.test.ts new file mode 100644 index 000000000..15248c375 --- /dev/null +++ b/packages/typegpu/tests/std/numeric/min.test.ts @@ -0,0 +1,94 @@ +import { describe, expect, it } from 'vitest'; +import tgpu from '../../../src/index.ts'; +import * as d from '../../../src/data/index.ts'; +import * as std from '../../../src/std/index.ts'; + +describe('min', () => { + it('acts as identity when called with one argument', () => { + const myMin = tgpu.fn([d.f32], d.f32)((a: number) => { + 'use gpu'; + return std.min(a); + }); + + expect(myMin(6)).toBe(6); + expect(tgpu.resolve([myMin])).toMatchInlineSnapshot(` + "fn myMin(a: f32) -> f32 { + return a; + }" + `); + }); + + it('works with two arguments', () => { + const myMin = tgpu.fn([d.f32, d.f32], d.f32)((a, b) => { + 'use gpu'; + return std.min(a, b); + }); + + expect(myMin(1, 2)).toBe(1); + expect(tgpu.resolve([myMin])).toMatchInlineSnapshot(` + "fn myMin(a: f32, b: f32) -> f32 { + return min(a, b); + }" + `); + }); + + it('works with multiple arguments', () => { + const myMin = tgpu.fn([d.f32, d.f32, d.f32, d.f32], d.f32)( + (a, b, c, d) => { + 'use gpu'; + return std.min(a, b, c, d); + }, + ); + + expect(myMin(2, 1, 4, 5)).toBe(1); + expect(tgpu.resolve([myMin])).toMatchInlineSnapshot(` + "fn myMin(a: f32, b: f32, c: f32, d2: f32) -> f32 { + return min(min(min(a, b), c), d2); + }" + `); + }); + + it('unifies arguments', () => { + const myMin = tgpu.fn([], d.f32)(() => { + 'use gpu'; + const a = d.u32(9); + const b = d.i32(1); + const c = d.f32(4); + return std.min(a, b, 3.3, c, 7); + }); + + expect(myMin()).toBe(1); + expect(tgpu.resolve([myMin])).toMatchInlineSnapshot(` + "fn myMin() -> f32 { + const a = 9u; + const b = 1i; + const c = 4f; + return min(min(min(min(f32(a), f32(b)), 3.3f), c), 7f); + }" + `); + }); + + it('works with vectors', () => { + const myMin = tgpu.fn([d.vec3u, d.vec3u], d.vec3u)((a, b) => { + 'use gpu'; + return std.min(a, b); + }); + + expect(myMin(d.vec3u(1, 2, 3), d.vec3u(3, 2, 1))) + .toStrictEqual(d.vec3u(1, 2, 1)); + expect(tgpu.resolve([myMin])).toMatchInlineSnapshot(` + "fn myMin(a: vec3u, b: vec3u) -> vec3u { + return min(a, b); + }" + `); + }); + + it('cannot be called with invalid arguments', () => { + // @ts-expect-error + (() => std.min()); + // @ts-expect-error + (() => std.min(1, d.vec2f())); + // @ts-expect-error + (() => std.min(d.vec3f(), d.vec2f())); + }); +});