diff --git a/packages/typegpu/src/std/boolean.ts b/packages/typegpu/src/std/boolean.ts index b96f8e3c89..b11e3ba520 100644 --- a/packages/typegpu/src/std/boolean.ts +++ b/packages/typegpu/src/std/boolean.ts @@ -212,6 +212,9 @@ export const ge = dualImpl({ function cpuNot(value: boolean): boolean; function cpuNot(value: T): T; +function cpuNot( + value: T, +): T extends boolean ? boolean : T; function cpuNot(value: T): T { assertKind(value, booleanKind); return generalizeBoolFn((a: boolean) => !a, [value]); diff --git a/packages/typegpu/src/std/derivative.ts b/packages/typegpu/src/std/derivative.ts index 9d880ffd17..46e1501914 100644 --- a/packages/typegpu/src/std/derivative.ts +++ b/packages/typegpu/src/std/derivative.ts @@ -6,7 +6,8 @@ import type { AnyFloat32VecInstance } from '../data/wgslTypes.ts'; import { unifyRestrictedSignature } from './numeric.ts'; type DerivativeSignature = ((value: number) => number) & - ((value: T) => T); + ((value: T) => T) & + ((value: T) => T extends number ? number : T); const derivativeNormalError = 'Derivative builtins are not allowed on the CPU'; // WGSL derivative builtins only accept f32 and vecN diff --git a/packages/typegpu/src/std/numeric.ts b/packages/typegpu/src/std/numeric.ts index 92ab11113e..3cfa0b875e 100644 --- a/packages/typegpu/src/std/numeric.ts +++ b/packages/typegpu/src/std/numeric.ts @@ -62,6 +62,8 @@ import { mul, sub } from './operators.ts'; type NumVec = AnyNumericVecInstance; +type WidenNumber = T extends number ? number : T; + // helpers const unaryIdentitySignature = (arg: BaseData) => { @@ -122,7 +124,8 @@ const anyConcreteInteger = [...anyConcreteIntegerPrimitive, ...anyConcreteIntege // std function cpuAbs(value: number): number; -function cpuAbs(value: T): T; +function cpuAbs(value: T): T; +function cpuAbs(value: T): WidenNumber; function cpuAbs(value: T): T { assertKind(value, numericKind); return generalizeFn(Math.abs, [value]); @@ -138,6 +141,7 @@ export const abs = dualImpl({ function cpuAcos(value: number): number; function cpuAcos(value: T): T; +function cpuAcos(value: T): WidenNumber; function cpuAcos(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.acos, [value]); @@ -153,6 +157,7 @@ export const acos = dualImpl({ function cpuAcosh(value: number): number; function cpuAcosh(value: T): T; +function cpuAcosh(value: T): WidenNumber; function cpuAcosh(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.acosh, [value]); @@ -168,6 +173,7 @@ export const acosh = dualImpl({ function cpuAsin(value: number): number; function cpuAsin(value: T): T; +function cpuAsin(value: T): WidenNumber; function cpuAsin(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.asin, [value]); @@ -183,6 +189,7 @@ export const asin = dualImpl({ function cpuAsinh(value: number): number; function cpuAsinh(value: T): T; +function cpuAsinh(value: T): WidenNumber; function cpuAsinh(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.asinh, [value]); @@ -198,6 +205,7 @@ export const asinh = dualImpl({ function cpuAtan(value: number): number; function cpuAtan(value: T): T; +function cpuAtan(value: T): WidenNumber; function cpuAtan(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.atan, [value]); @@ -213,6 +221,7 @@ export const atan = dualImpl({ function cpuAtanh(value: number): number; function cpuAtanh(value: T): T; +function cpuAtanh(value: T): WidenNumber; function cpuAtanh(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.atanh, [value]); @@ -244,6 +253,7 @@ export const atan2 = dualImpl({ function cpuCeil(value: number): number; function cpuCeil(value: T): T; +function cpuCeil(value: T): WidenNumber; function cpuCeil(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.ceil, [value]); @@ -275,6 +285,7 @@ export const clamp = dualImpl({ function cpuCos(value: number): number; function cpuCos(value: T): T; +function cpuCos(value: T): WidenNumber; function cpuCos(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.cos, [value]); @@ -290,6 +301,7 @@ export const cos = dualImpl({ function cpuCosh(value: number): number; function cpuCosh(value: T): T; +function cpuCosh(value: T): WidenNumber; function cpuCosh(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.cosh, [value]); @@ -305,6 +317,7 @@ export const cosh = dualImpl({ function cpuCountLeadingZeros(value: number): number; function cpuCountLeadingZeros(value: T): T; +function cpuCountLeadingZeros(value: T): WidenNumber; function cpuCountLeadingZeros(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -320,6 +333,7 @@ export const countLeadingZeros = dualImpl({ function cpuCountOneBits(value: number): number; function cpuCountOneBits(value: T): T; +function cpuCountOneBits(value: T): WidenNumber; function cpuCountOneBits(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -335,6 +349,7 @@ export const countOneBits = dualImpl({ function cpuCountTrailingZeros(value: number): number; function cpuCountTrailingZeros(value: T): T; +function cpuCountTrailingZeros(value: T): WidenNumber; function cpuCountTrailingZeros(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -362,6 +377,7 @@ export const cross = dualImpl({ function cpuDegrees(value: number): number; function cpuDegrees(value: T): T; +function cpuDegrees(value: T): WidenNumber; function cpuDegrees(value: T): T { assertKind(value, floatKind); if (typeof value === 'number') { @@ -457,6 +473,7 @@ export const dot4I8Packed = dualImpl<(e1: number, e2: number) => number>({ function cpuExp(value: number): number; function cpuExp(value: T): T; +function cpuExp(value: T): WidenNumber; function cpuExp(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.exp, [value]); @@ -472,6 +489,7 @@ export const exp = dualImpl({ function cpuExp2(value: number): number; function cpuExp2(value: T): T; +function cpuExp2(value: T): WidenNumber; function cpuExp2(value: T): T { assertKind(value, floatKind); return generalizeFn((val) => 2 ** val, [value]); @@ -524,6 +542,7 @@ export const faceForward = dualImpl<(e1: T, e2: T function cpuFirstLeadingBit(value: number): number; function cpuFirstLeadingBit(value: T): T; +function cpuFirstLeadingBit(value: T): WidenNumber; function cpuFirstLeadingBit(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -539,6 +558,7 @@ export const firstLeadingBit = dualImpl({ function cpuFirstTrailingBit(value: number): number; function cpuFirstTrailingBit(value: T): T; +function cpuFirstTrailingBit(value: T): WidenNumber; function cpuFirstTrailingBit(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -554,6 +574,8 @@ export const firstTrailingBit = dualImpl({ function cpuFloor(value: number): number; function cpuFloor(value: T): T; +function cpuFloor(value: T): WidenNumber; +function cpuFloor(value: T | number): T | number; function cpuFloor(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.floor, [value]); @@ -590,6 +612,7 @@ export const fma = dualImpl({ function cpuFract(value: number): number; function cpuFract(value: T): T; +function cpuFract(value: T): WidenNumber; function cpuFract(value: T): T { assertKind(value, floatKind); return generalizeFn((value) => value - Math.floor(value), [value]); @@ -618,6 +641,11 @@ const FrexpResults = { type FrexpOverload = { (value: number): Infer<(typeof FrexpResults)['f32']>; (value: T): Infer<(typeof FrexpResults)[T['kind']]>; + ( + value: T, + ): T extends AnyFloatVecInstance + ? Infer<(typeof FrexpResults)[T['kind']]> + : Infer<(typeof FrexpResults)['f32']>; }; export const frexp = dualImpl({ @@ -674,6 +702,7 @@ export const insertBits = dualImpl({ function cpuInverseSqrt(value: number): number; function cpuInverseSqrt(value: T): T; +function cpuInverseSqrt(value: T): WidenNumber; function cpuInverseSqrt(value: T): T { assertKind(value, floatKind); if (typeof value === 'number') { @@ -735,6 +764,7 @@ export const ldexp = dualImpl({ function cpuLength(value: number): number; function cpuLength(value: T): number; +function cpuLength(value: T): number; function cpuLength(value: T): number { assertKind(value, floatKind); if (typeof value === 'number') { @@ -762,6 +792,7 @@ export const length = dualImpl({ function cpuLog(value: number): number; function cpuLog(value: T): T; +function cpuLog(value: T): WidenNumber; function cpuLog(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.log, [value]); @@ -777,6 +808,7 @@ export const log = dualImpl({ function cpuLog2(value: number): number; function cpuLog2(value: T): T; +function cpuLog2(value: T): WidenNumber; function cpuLog2(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.log2, [value]); @@ -876,9 +908,19 @@ const ModfResult = { type ModfOverload = { (value: number): Infer<(typeof ModfResult)['f32']>; (value: T): Infer<(typeof ModfResult)[T['kind']]>; + ( + value: T, + ): T extends AnyFloatVecInstance + ? Infer<(typeof ModfResult)[T['kind']]> + : Infer<(typeof ModfResult)['f32']>; }; function cpuModf(e: number): Infer<(typeof ModfResult)['f32']>; function cpuModf(e: T): Infer<(typeof ModfResult)[T['kind']]>; +function cpuModf( + e: T, +): T extends AnyFloatVecInstance + ? Infer<(typeof ModfResult)[T['kind']]> + : Infer<(typeof ModfResult)['f32']>; function cpuModf( _value: T, ): Infer<(typeof ModfResult)[keyof typeof ModfResult]> { @@ -933,6 +975,7 @@ export const pow = dualImpl({ }); function cpuQuantizeToF16(value: number): number; function cpuQuantizeToF16(value: T): T; +function cpuQuantizeToF16(value: T): WidenNumber; function cpuQuantizeToF16(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -954,7 +997,8 @@ export const quantizeToF16 = dualImpl({ }); function cpuRadians(value: number): number; -function cpuRadians(value: T): T; +function cpuRadians(value: T): T; +function cpuRadians(value: T): WidenNumber; function cpuRadians(value: T): T { assertKind(value, floatKind); if (typeof value === 'number') { @@ -1007,6 +1051,7 @@ export const refract = dualImpl<(e1: T, e2: T, e3 }); function cpuReverseBits(value: number): number; function cpuReverseBits(value: T): T; +function cpuReverseBits(value: T): WidenNumber; function cpuReverseBits(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } @@ -1022,6 +1067,7 @@ export const reverseBits = dualImpl({ function cpuRound(value: number): number; function cpuRound(value: T): T; +function cpuRound(value: T): WidenNumber; function cpuRound(value: T): T { assertKind(value, floatKind); if (typeof value === 'number') { @@ -1049,6 +1095,7 @@ export const round = dualImpl({ function cpuSaturate(value: number): number; function cpuSaturate(value: T): T; +function cpuSaturate(value: T): WidenNumber; function cpuSaturate(value: T): T { assertKind(value, floatKind); if (typeof value === 'number') { @@ -1069,6 +1116,8 @@ export const saturate = dualImpl({ function cpuSign(e: number): number; function cpuSign(e: T): T; +function cpuSign(e: T): WidenNumber; +function cpuSign(e: T | number): T | number; function cpuSign(e: T): T { assertKind(e, signedKind); return generalizeFn(Math.sign, [e]); @@ -1091,6 +1140,7 @@ export const sign = dualImpl({ function cpuSin(value: number): number; function cpuSin(value: T): T; +function cpuSin(value: T): WidenNumber; function cpuSin(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.sin, [value]); @@ -1106,6 +1156,7 @@ export const sin = dualImpl({ function cpuSinh(value: number): number; function cpuSinh(value: T): T; +function cpuSinh(value: T): WidenNumber; function cpuSinh(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.sinh, [value]); @@ -1137,6 +1188,7 @@ export const smoothstep = dualImpl({ function cpuSqrt(value: number): number; function cpuSqrt(value: T): T; +function cpuSqrt(value: T): WidenNumber; function cpuSqrt(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.sqrt, [value]); @@ -1173,6 +1225,7 @@ export const step = dualImpl({ function cpuTan(value: number): number; function cpuTan(value: T): T; +function cpuTan(value: T): WidenNumber; function cpuTan(value: T): T { assertKind(value, floatKind); if (typeof value === 'number') { @@ -1193,6 +1246,7 @@ export const tan = dualImpl({ function cpuTanh(value: number): number; function cpuTanh(value: T): T; +function cpuTanh(value: T): WidenNumber; function cpuTanh(value: T): T { assertKind(value, floatKind); return generalizeFn(Math.tanh, [value]); @@ -1275,6 +1329,7 @@ export const transpose = dualImpl({ function cpuTrunc(value: number): number; function cpuTrunc(value: T): T; +function cpuTrunc(value: T): WidenNumber; function cpuTrunc(_value: T): T { throw new Error('Unreachable code. The function is only used for the type.'); } diff --git a/packages/typegpu/src/std/operators.ts b/packages/typegpu/src/std/operators.ts index cc43783e72..af5121cd9b 100644 --- a/packages/typegpu/src/std/operators.ts +++ b/packages/typegpu/src/std/operators.ts @@ -36,6 +36,8 @@ import { unify } from '../tgsl/conversion.ts'; type NumVec = AnyNumericVecInstance; type Mat = AnyMatInstance; +type WidenNumber = T extends number ? number : T; + const getPrimitive = (t: BaseData): BaseData => ('primitive' in t ? (t.primitive as BaseData) : t); const makeBinarySignature = @@ -268,6 +270,7 @@ export const mod = dualImpl({ function cpuNeg(value: number): number; function cpuNeg(value: T): T; +function cpuNeg(value: T): WidenNumber; function cpuNeg(value: NumVec | number): NumVec | number { assertKind(value, signedKind); return generalizeFn((value) => -value, [value]); diff --git a/packages/typegpu/src/std/subgroup.ts b/packages/typegpu/src/std/subgroup.ts index d203cbcab8..4462a5fd2f 100644 --- a/packages/typegpu/src/std/subgroup.ts +++ b/packages/typegpu/src/std/subgroup.ts @@ -13,11 +13,13 @@ import { unify } from '../tgsl/conversion.ts'; interface IdentityNumOrVec { (e: number): number; (e: T): T; + (e: T): T extends number ? number : T; } interface IdentityIntNumOrVec { (e: number): number; (e: T): T; + (e: T): T extends number ? number : T; } interface IdentityNumOrVecWithIdx { diff --git a/packages/typegpu/tests/std/numeric/scalarVectorUnions.test.ts b/packages/typegpu/tests/std/numeric/scalarVectorUnions.test.ts new file mode 100644 index 0000000000..1e7c2bc7a8 --- /dev/null +++ b/packages/typegpu/tests/std/numeric/scalarVectorUnions.test.ts @@ -0,0 +1,400 @@ +import { describe, expect, expectTypeOf, it } from 'vitest'; +import * as d from 'typegpu/data'; +import * as std from 'typegpu/std'; + +type FloatVector = d.v2f | d.v3f | d.v4f | d.v2h | d.v3h | d.v4h; +type IntegerVector = d.v2i | d.v3i | d.v4i | d.v2u | d.v3u | d.v4u; +type NumericVector = FloatVector | IntegerVector; +type SignedVector = FloatVector | d.v2i | d.v3i | d.v4i; +type WidenNumber = T extends number ? number : T; + +const floatUnaryNames = [ + 'acos', + 'acosh', + 'asin', + 'asinh', + 'atan', + 'atanh', + 'ceil', + 'cos', + 'cosh', + 'degrees', + 'exp', + 'exp2', + 'floor', + 'fract', + 'inverseSqrt', + 'log', + 'log2', + 'round', + 'saturate', + 'sin', + 'sinh', + 'sqrt', + 'tan', + 'tanh', + 'trunc', +] as const; + +const integerUnaryNames = [ + 'countLeadingZeros', + 'countOneBits', + 'countTrailingZeros', + 'firstLeadingBit', + 'firstTrailingBit', + 'reverseBits', +] as const; + +describe('numeric scalar-vector union overloads', () => { + it.each(floatUnaryNames)('%s preserves floating-point vector unions', (name) => { + const fn = std[name]; + expectTypeOf((value: number | d.v2f) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v3f) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v4f) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v2h) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v3h) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v4h) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 1.5 | d.v2f) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: -1 | d.v4h) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: d.v2f | d.v3h) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf(() => fn(1.5)).returns.toEqualTypeOf(); + expectTypeOf(() => fn(d.vec3h())).returns.toEqualTypeOf(); + const vector = (value: T): T => fn(value); + const union = (value: T | number): T | number => fn(value); + const literal = (value: T): WidenNumber => fn(value); + expectTypeOf(vector).returns.toEqualTypeOf(); + expectTypeOf(union).returns.toEqualTypeOf(); + expectTypeOf(literal<1.5 | d.v2f>).returns.toEqualTypeOf(); + // Compile-only invalid calls: several helpers have no CPU implementation. + if (false) { + // @ts-expect-error Floating-point helpers exclude signed integer vectors. + fn(d.vec2i()); + // @ts-expect-error Floating-point helpers exclude unsigned integer vectors. + fn(d.vec2u()); + // @ts-expect-error Non-numeric values remain unsupported. + fn(true); + } + }); + + it.each(integerUnaryNames)('%s preserves signed and unsigned integer unions', (name) => { + const fn = std[name]; + expectTypeOf((value: number | d.v2i) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v3i) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v4i) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v2u) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v3u) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v4u) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 4 | d.v2i) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 4 | d.v4u) => fn(value)).returns.toEqualTypeOf(); + const vector = (value: T): T => fn(value); + const union = (value: T | number): T | number => fn(value); + const literal = (value: T): WidenNumber => fn(value); + expectTypeOf(vector).returns.toEqualTypeOf(); + expectTypeOf(union).returns.toEqualTypeOf(); + expectTypeOf(literal<4 | d.v2i>).returns.toEqualTypeOf(); + if (false) { + // @ts-expect-error Integer helpers exclude floating-point vectors. + fn(d.vec2f()); + // @ts-expect-error Half-precision vectors are not integer vectors. + fn(d.vec2h()); + } + }); + + it('retains the abs, sign, and quantizeToF16 vector domains', () => { + expectTypeOf((value: number | NumericVector) => std.abs(value)).returns.toEqualTypeOf< + number | NumericVector + >(); + expectTypeOf((value: -2 | d.v2i) => std.abs(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 2 | d.v4u) => std.abs(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | SignedVector) => std.sign(value)).returns.toEqualTypeOf< + number | SignedVector + >(); + expectTypeOf((value: number | d.v2f) => std.quantizeToF16(value)).returns.toEqualTypeOf< + number | d.v2f + >(); + expectTypeOf((value: number | d.v3f) => std.quantizeToF16(value)).returns.toEqualTypeOf< + number | d.v3f + >(); + expectTypeOf((value: 1.5 | d.v4f) => std.quantizeToF16(value)).returns.toEqualTypeOf< + number | d.v4f + >(); + const abs = (value: T): WidenNumber => std.abs(value); + const sign = (value: T): WidenNumber => std.sign(value); + const quantize = (value: T): WidenNumber => + std.quantizeToF16(value); + expectTypeOf(abs<-2 | d.v4u>).returns.toEqualTypeOf(); + expectTypeOf(sign<-2 | d.v3i>).returns.toEqualTypeOf(); + expectTypeOf(quantize<1.5 | d.v4f>).returns.toEqualTypeOf(); + if (false) { + // @ts-expect-error Sign excludes unsigned vectors, including union members. + std.sign(null as unknown as 2 | d.v4u); + // @ts-expect-error quantizeToF16 excludes f16 vectors. + std.quantizeToF16(d.vec2h()); + // @ts-expect-error quantizeToF16 excludes integer vectors. + std.quantizeToF16(d.vec2i()); + } + }); + + it('retains explicit scalar type arguments and widens generic numeric literals', () => { + expectTypeOf(() => std.abs(-2)).returns.toEqualTypeOf(); + expectTypeOf(() => std.radians(180)).returns.toEqualTypeOf(); + expectTypeOf(() => std.clamp(1, 0, 2)).returns.toEqualTypeOf(); + expectTypeOf(() => std.step(0, 1)).returns.toEqualTypeOf(); + expectTypeOf((value: 180 | d.v3h) => std.floor(value)).returns.toEqualTypeOf< + number | d.v3h + >(); + expectTypeOf(() => std.sign(180)).returns.toEqualTypeOf(); + const radians = (value: T): WidenNumber => + std.radians(value); + const absScalar = (value: T): number => std.abs(value); + const radiansScalar = (value: T): number => std.radians(value); + expectTypeOf(radians<180 | d.v3h>).returns.toEqualTypeOf(); + expectTypeOf(absScalar<-2>).returns.toEqualTypeOf(); + expectTypeOf(radiansScalar<180>).returns.toEqualTypeOf(); + }); + + it('length retains a scalar return for scalar-vector unions', () => { + expectTypeOf((value: number | d.v2f) => std.length(value)).returns.toEqualTypeOf(); + expectTypeOf((value: -2 | d.v4h) => std.length(value)).returns.toEqualTypeOf(); + expectTypeOf((value: FloatVector) => std.length(value)).returns.toEqualTypeOf(); + const measure = (value: T): number => std.length(value); + expect(measure(-3)).toBe(3); + expect(measure(d.vec2f(3, 4))).toBe(5); + if (false) { + // @ts-expect-error Length excludes integer vectors. + std.length(d.vec2i()); + } + }); + + it('frexp and modf return unions of complete mapped result structs', () => { + const frexp = (value: T) => std.frexp(value); + const modf = (value: T) => std.modf(value); + expectTypeOf(frexp).returns.toEqualTypeOf<{ fract: number; exp: number }>(); + expectTypeOf(frexp).returns.toEqualTypeOf<{ fract: d.v2f; exp: d.v2i }>(); + expectTypeOf(frexp).returns.toEqualTypeOf<{ fract: d.v3h; exp: d.v3i }>(); + expectTypeOf(frexp<1.5 | d.v2f>).returns.toEqualTypeOf< + { fract: number; exp: number } | { fract: d.v2f; exp: d.v2i } + >(); + expectTypeOf(modf).returns.toEqualTypeOf<{ fract: number; whole: number }>(); + expectTypeOf(modf).returns.toEqualTypeOf<{ fract: d.v2f; whole: d.v2f }>(); + expectTypeOf(modf).returns.toEqualTypeOf<{ fract: d.v3h; whole: d.v3h }>(); + expectTypeOf(modf<1.5 | d.v2f>).returns.toEqualTypeOf< + { fract: number; whole: number } | { fract: d.v2f; whole: d.v2f } + >(); + if (false) { + // @ts-expect-error frexp excludes integer vectors. + std.frexp(d.vec2i()); + // @ts-expect-error modf excludes integer vectors. + std.modf(d.vec2u()); + } + }); + + it('max and min still reject explicitly supplied scalar-vector union type arguments', () => { + expect(std.max(d.vec2f(1, 4), d.vec2f(3, 2))).toEqual(d.vec2f(3, 4)); + expect(std.min(d.vec2f(1, 4), d.vec2f(3, 2))).toEqual(d.vec2f(1, 2)); + if (false) { + // @ts-expect-error A union type argument must not permit mismatched operand shapes. + std.max(d.vec2f(), 1); + // @ts-expect-error Min has the same matching-shape requirement. + std.min(d.vec2f(), 1); + } + }); +}); + +const cpuUnaryApply = { + abs: (value) => std.abs(value), + acos: (value) => std.acos(value), + acosh: (value) => std.acosh(value), + asin: (value) => std.asin(value), + asinh: (value) => std.asinh(value), + atan: (value) => std.atan(value), + atanh: (value) => std.atanh(value), + ceil: (value) => std.ceil(value), + cos: (value) => std.cos(value), + cosh: (value) => std.cosh(value), + exp: (value) => std.exp(value), + exp2: (value) => std.exp2(value), + floor: (value) => std.floor(value), + fract: (value) => std.fract(value), + log: (value) => std.log(value), + log2: (value) => std.log2(value), + sign: (value) => std.sign(value), + sin: (value) => std.sin(value), + sinh: (value) => std.sinh(value), + sqrt: (value) => std.sqrt(value), + tanh: (value) => std.tanh(value), +} satisfies Record number | d.v2f>; + +const cpuUnaryCases = [ + ['abs', -2, 2], + ['acos', 1, 0], + ['acosh', 1, 0], + ['asin', 0, 0], + ['asinh', 0, 0], + ['atan', 0, 0], + ['atanh', 0, 0], + ['ceil', 1.2, 2], + ['cos', 0, 1], + ['cosh', 0, 1], + ['exp', 0, 1], + ['exp2', 2, 4], + ['floor', 1.7, 1], + ['fract', 1.2, 0.2], + ['log', 1, 0], + ['log2', 8, 3], + ['sign', -2, -1], + ['sin', 0, 0], + ['sinh', 0, 0], + ['sqrt', 4, 2], + ['tanh', 0, 0], +] as const; + +const scalarOnlyUnaryApply = { + degrees: (value) => std.degrees(value), + inverseSqrt: (value) => std.inverseSqrt(value), + radians: (value) => std.radians(value), + round: (value) => std.round(value), + saturate: (value) => std.saturate(value), + tan: (value) => std.tan(value), +} satisfies Record number | d.v2f>; + +describe('numeric union CPU behavior', () => { + it.each(cpuUnaryCases)('%s computes scalar and vector branches', (name, input, expected) => { + const apply = cpuUnaryApply[name]; + expect(apply(input)).toBeCloseTo(expected); + const result = apply(d.vec2f(input, input)); + expect(typeof result).toBe('object'); + if (typeof result !== 'number') { + expect(result.x).toBeCloseTo(expected); + expect(result.y).toBeCloseTo(expected); + } + }); + it.each([ + ['degrees', Math.PI, 180], + ['inverseSqrt', 4, 0.5], + ['radians', 180, Math.PI], + ['round', 2.5, 2], + ['saturate', 2, 1], + ['tan', 0, 0], + ] as const)('%s computes its scalar union branch', (name, input, expected) => { + expect(scalarOnlyUnaryApply[name](input)).toBeCloseTo(expected); + }); +}); + +const derivativeNames = [ + 'dpdx', + 'dpdxCoarse', + 'dpdxFine', + 'dpdy', + 'dpdyCoarse', + 'dpdyFine', + 'fwidth', + 'fwidthCoarse', + 'fwidthFine', +] as const; +const numericSubgroupNames = [ + 'subgroupAdd', + 'subgroupExclusiveAdd', + 'subgroupInclusiveAdd', + 'subgroupBroadcastFirst', + 'subgroupMax', + 'subgroupMin', + 'subgroupMul', + 'subgroupExclusiveMul', + 'subgroupInclusiveMul', + 'quadSwapDiagonal', + 'quadSwapX', + 'quadSwapY', +] as const; +const integerSubgroupNames = ['subgroupAnd', 'subgroupOr', 'subgroupXor'] as const; + +describe('GPU-only scalar-vector union overloads', () => { + it.each(derivativeNames)('%s preserves f32 vector unions', (name) => { + const fn = std[name]; + expectTypeOf((value: number | d.v2f) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v3f) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 1.5 | d.v4f) => fn(value)).returns.toEqualTypeOf(); + const generic = (value: T): WidenNumber => fn(value); + expectTypeOf(generic<1.5 | d.v2f>).returns.toEqualTypeOf(); + if (false) { + // @ts-expect-error Derivatives exclude f16 vectors. + fn(d.vec2h()); + // @ts-expect-error Derivatives exclude integer vectors. + fn(d.vec3i()); + } + }); + it.each(numericSubgroupNames)('%s preserves numeric union shapes', (name) => { + const fn = std[name]; + expectTypeOf((value: number | NumericVector) => fn(value)).returns.toEqualTypeOf< + number | NumericVector + >(); + expectTypeOf((value: 1.5 | d.v3h) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 4 | d.v4u) => fn(value)).returns.toEqualTypeOf(); + const generic = (value: T): WidenNumber => fn(value); + expectTypeOf(generic<1.5 | d.v2f>).returns.toEqualTypeOf(); + if (false) { + // @ts-expect-error Numeric subgroup operations exclude booleans. + fn(true); + // @ts-expect-error Numeric subgroup operations exclude boolean vectors. + fn(d.vec2b()); + } + }); + it.each(integerSubgroupNames)('%s retains its integer-only vector domain', (name) => { + const fn = std[name]; + expectTypeOf((value: number | IntegerVector) => fn(value)).returns.toEqualTypeOf< + number | IntegerVector + >(); + expectTypeOf((value: number | d.v3i) => fn(value)).returns.toEqualTypeOf(); + expectTypeOf((value: 4 | d.v4u) => fn(value)).returns.toEqualTypeOf(); + if (false) { + // @ts-expect-error Bitwise subgroup operations exclude floating-point vectors. + fn(d.vec2f()); + // @ts-expect-error Bitwise subgroup operations exclude half-precision vectors. + fn(d.vec2h()); + } + }); +}); + +describe('other unary scalar-vector union overloads', () => { + it('not widens boolean literals and preserves boolean vector shapes', () => { + expectTypeOf((value: boolean | d.v2b) => std.not(value)).returns.toEqualTypeOf< + boolean | d.v2b + >(); + expectTypeOf((value: boolean | d.v3b) => std.not(value)).returns.toEqualTypeOf< + boolean | d.v3b + >(); + expectTypeOf((value: true | d.v4b) => std.not(value)).returns.toEqualTypeOf(); + const generic = (value: T): T extends boolean ? boolean : T => + std.not(value); + expectTypeOf(generic).returns.toEqualTypeOf(); + const apply = (value: true | d.v3b) => std.not(value); + expect(apply(true)).toBe(false); + expect(apply(d.vec3b(true, false, true))).toEqual(d.vec3b(false, true, false)); + if (false) { + // @ts-expect-error Logical negation excludes numbers. + std.not(1); + // @ts-expect-error Logical negation excludes numeric vectors. + std.not(d.vec2f()); + } + }); + it('neg widens scalar literals and preserves its signed-vector domain', () => { + expectTypeOf((value: number | SignedVector) => std.neg(value)).returns.toEqualTypeOf< + number | SignedVector + >(); + expectTypeOf((value: -2 | d.v3i) => std.neg(value)).returns.toEqualTypeOf(); + expectTypeOf((value: number | d.v4h) => std.neg(value)).returns.toEqualTypeOf(); + const generic = (value: T): WidenNumber => std.neg(value); + expectTypeOf(generic<-2 | d.v2f>).returns.toEqualTypeOf(); + const apply = (value: -2 | d.v2i) => std.neg(value); + expect(apply(-2)).toBe(2); + expect(apply(d.vec2i(-2, 2))).toEqual(d.vec2i(2, -2)); + if (false) { + // @ts-expect-error Negation excludes unsigned vectors. + std.neg(d.vec2u()); + // @ts-expect-error Negation excludes boolean vectors. + std.neg(d.vec2b()); + // @ts-expect-error Negation excludes matrices. + std.neg(d.mat2x2f()); + } + }); +}); diff --git a/packages/typegpu/tests/std/numeric/sign.test.ts b/packages/typegpu/tests/std/numeric/sign.test.ts index 00dc5857e6..ed43ff20b2 100644 --- a/packages/typegpu/tests/std/numeric/sign.test.ts +++ b/packages/typegpu/tests/std/numeric/sign.test.ts @@ -1,6 +1,6 @@ -import { describe, expect, it } from 'vitest'; -import { vec3f } from 'typegpu/data'; -import { isCloseTo, sign } from 'typegpu/std'; +import { describe, expect, expectTypeOf, it } from 'vitest'; +import { vec2f, vec2i, vec2u, vec3f, type v2f, type v2i } from 'typegpu/data'; +import { floor, isCloseTo, sign } from 'typegpu/std'; describe('sign', () => { it('computes sign of numeric value', () => { @@ -12,4 +12,55 @@ describe('sign', () => { it('computes sign of a numeric vector', () => { expect(isCloseTo(sign(vec3f(-1000, 0, 2000)), vec3f(-1, 0, 1))).toBe(true); }); + + it('accepts scalar-vector unions without preserving scalar literal types', () => { + const apply = (value: number | v2f) => { + expectTypeOf(sign(value)).toEqualTypeOf(); + expectTypeOf(floor(value)).toEqualTypeOf(); + return [sign(value), floor(value)]; + }; + const [signed, floored] = apply(vec2f(-1.5, 2.5)); + + expect(signed).toEqual(vec2f(-1, 1)); + expect(floored).toEqual(vec2f(-2, 2)); + expect(apply(-1.5)).toEqual([-1, -2]); + expectTypeOf(sign(2.5 as const)).toEqualTypeOf(); + expectTypeOf(floor(1.5 as const)).toEqualTypeOf(); + }); + + it('widens numeric literals within scalar-vector unions', () => { + const apply = (value: 1.5 | v2f) => { + expectTypeOf(sign(value)).toEqualTypeOf(); + expectTypeOf(floor(value)).toEqualTypeOf(); + return [sign(value), floor(value)]; + }; + + expect(apply(1.5)).toEqual([1, 1]); + expect(apply(vec2f(-1.5, 2.5))).toEqual([vec2f(-1, 1), vec2f(-2, 2)]); + }); + + it('preserves generic vector return types', () => { + const apply = (value: T): [T, T] => [sign(value), floor(value)]; + const signedInteger = (value: T): T => sign(value); + + expect(apply(vec2f(-1.5, 2.5))).toEqual([vec2f(-1, 1), vec2f(-2, 2)]); + expect(signedInteger(vec2i(-2, 2))).toEqual(vec2i(-1, 1)); + }); + + it('widens generic scalar return types', () => { + const apply = (value: T) => { + expectTypeOf(sign(value)).toEqualTypeOf(); + expectTypeOf(floor(value)).toEqualTypeOf(); + return [sign(value), floor(value)]; + }; + + expect(apply(1.5)).toEqual([1, 1]); + }); + + it('rejects unsupported vector types', () => { + // @ts-expect-error floor only accepts floating-point vectors + expect(() => floor(vec2i())).toThrow(); + // @ts-expect-error sign only accepts signed vectors + expect(() => sign(vec2u())).toThrow(); + }); });