Skip to content

fix: accept scalar-vector unions across std helpers #2962

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
oxura wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix: accept scalar-vector unions across std helpers #2962

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
40e3649
fix: accept scalar-vector unions in floor and sign
oxura Sep 2, 2026
afb81e4
fix: preserve scalar numeric return types
oxura Sep 2, 2026
d71b337
fix: widen scalar-vector union return types
oxura Oct 2, 2026
df5d524
fix: accept scalar-vector unions across std helpers
oxura Oct 6, 2026
d133ad8
fix: restrict scalar-vector union updates to unary helpers
oxura Oct 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions packages/typegpu/src/std/boolean.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,9 @@ export const ge = dualImpl({

function cpuNot(value: boolean): boolean;
function cpuNot<T extends AnyBooleanVecInstance>(value: T): T;
function cpuNot<T extends AnyBooleanVecInstance | boolean>(
value: T,
): T extends boolean ? boolean : T;
function cpuNot<T extends AnyBooleanVecInstance | boolean>(value: T): T {
assertKind(value, booleanKind);
return generalizeBoolFn((a: boolean) => !a, [value]);
Expand Down
3 changes: 2 additions & 1 deletion packages/typegpu/src/std/derivative.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@ import type { AnyFloat32VecInstance } from '../data/wgslTypes.ts';
import { unifyRestrictedSignature } from './numeric.ts';

type DerivativeSignature = ((value: number) => number) &
(<T extends AnyFloat32VecInstance>(value: T) => T);
(<T extends AnyFloat32VecInstance>(value: T) => T) &
(<T extends AnyFloat32VecInstance | number>(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<f32>
Expand Down
59 changes: 57 additions & 2 deletions packages/typegpu/src/std/numeric.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ import { mul, sub } from './operators.ts';

type NumVec = AnyNumericVecInstance;

type WidenNumber<T> = T extends number ? number : T;

// helpers

const unaryIdentitySignature = (arg: BaseData) => {
Expand Down Expand Up @@ -122,7 +124,8 @@ const anyConcreteInteger = [...anyConcreteIntegerPrimitive, ...anyConcreteIntege
// std

function cpuAbs(value: number): number;
function cpuAbs<T extends NumVec | number>(value: T): T;
function cpuAbs<T extends NumVec>(value: T): T;
function cpuAbs<T extends NumVec | number>(value: T): WidenNumber<T>;
function cpuAbs<T extends NumVec | number>(value: T): T {
assertKind(value, numericKind);
return generalizeFn(Math.abs, [value]);
Expand All @@ -138,6 +141,7 @@ export const abs = dualImpl({

function cpuAcos(value: number): number;
function cpuAcos<T extends AnyFloatVecInstance>(value: T): T;
function cpuAcos<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuAcos<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.acos, [value]);
Expand All @@ -153,6 +157,7 @@ export const acos = dualImpl({

function cpuAcosh(value: number): number;
function cpuAcosh<T extends AnyFloatVecInstance>(value: T): T;
function cpuAcosh<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuAcosh<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.acosh, [value]);
Expand All @@ -168,6 +173,7 @@ export const acosh = dualImpl({

function cpuAsin(value: number): number;
function cpuAsin<T extends AnyFloatVecInstance>(value: T): T;
function cpuAsin<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuAsin<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.asin, [value]);
Expand All @@ -183,6 +189,7 @@ export const asin = dualImpl({

function cpuAsinh(value: number): number;
function cpuAsinh<T extends AnyFloatVecInstance>(value: T): T;
function cpuAsinh<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuAsinh<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.asinh, [value]);
Expand All @@ -198,6 +205,7 @@ export const asinh = dualImpl({

function cpuAtan(value: number): number;
function cpuAtan<T extends AnyFloatVecInstance>(value: T): T;
function cpuAtan<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuAtan<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.atan, [value]);
Expand All @@ -213,6 +221,7 @@ export const atan = dualImpl({

function cpuAtanh(value: number): number;
function cpuAtanh<T extends AnyFloatVecInstance>(value: T): T;
function cpuAtanh<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuAtanh<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.atanh, [value]);
Expand Down Expand Up @@ -244,6 +253,7 @@ export const atan2 = dualImpl({

function cpuCeil(value: number): number;
function cpuCeil<T extends AnyFloatVecInstance>(value: T): T;
function cpuCeil<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuCeil<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.ceil, [value]);
Expand Down Expand Up @@ -275,6 +285,7 @@ export const clamp = dualImpl({

function cpuCos(value: number): number;
function cpuCos<T extends AnyFloatVecInstance>(value: T): T;
function cpuCos<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuCos<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.cos, [value]);
Expand All @@ -290,6 +301,7 @@ export const cos = dualImpl({

function cpuCosh(value: number): number;
function cpuCosh<T extends AnyFloatVecInstance>(value: T): T;
function cpuCosh<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuCosh<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.cosh, [value]);
Expand All @@ -305,6 +317,7 @@ export const cosh = dualImpl({

function cpuCountLeadingZeros(value: number): number;
function cpuCountLeadingZeros<T extends AnyIntegerVecInstance>(value: T): T;
function cpuCountLeadingZeros<T extends AnyIntegerVecInstance | number>(value: T): WidenNumber<T>;
function cpuCountLeadingZeros<T extends AnyIntegerVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand All @@ -320,6 +333,7 @@ export const countLeadingZeros = dualImpl<typeof cpuCountLeadingZeros>({

function cpuCountOneBits(value: number): number;
function cpuCountOneBits<T extends AnyIntegerVecInstance>(value: T): T;
function cpuCountOneBits<T extends AnyIntegerVecInstance | number>(value: T): WidenNumber<T>;
function cpuCountOneBits<T extends AnyIntegerVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand All @@ -335,6 +349,7 @@ export const countOneBits = dualImpl<typeof cpuCountOneBits>({

function cpuCountTrailingZeros(value: number): number;
function cpuCountTrailingZeros<T extends AnyIntegerVecInstance>(value: T): T;
function cpuCountTrailingZeros<T extends AnyIntegerVecInstance | number>(value: T): WidenNumber<T>;
function cpuCountTrailingZeros<T extends AnyIntegerVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand Down Expand Up @@ -362,6 +377,7 @@ export const cross = dualImpl({

function cpuDegrees(value: number): number;
function cpuDegrees<T extends AnyFloatVecInstance>(value: T): T;
function cpuDegrees<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuDegrees<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand Down Expand Up @@ -457,6 +473,7 @@ export const dot4I8Packed = dualImpl<(e1: number, e2: number) => number>({

function cpuExp(value: number): number;
function cpuExp<T extends AnyFloatVecInstance>(value: T): T;
function cpuExp<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuExp<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.exp, [value]);
Expand All @@ -472,6 +489,7 @@ export const exp = dualImpl({

function cpuExp2(value: number): number;
function cpuExp2<T extends AnyFloatVecInstance>(value: T): T;
function cpuExp2<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuExp2<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn((val) => 2 ** val, [value]);
Expand Down Expand Up @@ -524,6 +542,7 @@ export const faceForward = dualImpl<<T extends AnyFloatVecInstance>(e1: T, e2: T

function cpuFirstLeadingBit(value: number): number;
function cpuFirstLeadingBit<T extends AnyIntegerVecInstance>(value: T): T;
function cpuFirstLeadingBit<T extends AnyIntegerVecInstance | number>(value: T): WidenNumber<T>;
function cpuFirstLeadingBit<T extends AnyIntegerVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand All @@ -539,6 +558,7 @@ export const firstLeadingBit = dualImpl<typeof cpuFirstLeadingBit>({

function cpuFirstTrailingBit(value: number): number;
function cpuFirstTrailingBit<T extends AnyIntegerVecInstance>(value: T): T;
function cpuFirstTrailingBit<T extends AnyIntegerVecInstance | number>(value: T): WidenNumber<T>;
function cpuFirstTrailingBit<T extends AnyIntegerVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand All @@ -554,6 +574,8 @@ export const firstTrailingBit = dualImpl<typeof cpuFirstTrailingBit>({

function cpuFloor(value: number): number;
function cpuFloor<T extends AnyFloatVecInstance>(value: T): T;
function cpuFloor<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuFloor<T extends AnyFloatVecInstance>(value: T | number): T | number;
function cpuFloor<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.floor, [value]);
Expand Down Expand Up @@ -590,6 +612,7 @@ export const fma = dualImpl({

function cpuFract(value: number): number;
function cpuFract<T extends AnyFloatVecInstance>(value: T): T;
function cpuFract<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuFract<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn((value) => value - Math.floor(value), [value]);
Expand Down Expand Up @@ -618,6 +641,11 @@ const FrexpResults = {
type FrexpOverload = {
(value: number): Infer<(typeof FrexpResults)['f32']>;
<T extends AnyFloatVecInstance>(value: T): Infer<(typeof FrexpResults)[T['kind']]>;
<T extends AnyFloatVecInstance | number>(
value: T,
): T extends AnyFloatVecInstance
? Infer<(typeof FrexpResults)[T['kind']]>
: Infer<(typeof FrexpResults)['f32']>;
};

export const frexp = dualImpl<FrexpOverload>({
Expand Down Expand Up @@ -674,6 +702,7 @@ export const insertBits = dualImpl<typeof cpuInsertBits>({

function cpuInverseSqrt(value: number): number;
function cpuInverseSqrt<T extends AnyFloatVecInstance>(value: T): T;
function cpuInverseSqrt<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuInverseSqrt<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand Down Expand Up @@ -735,6 +764,7 @@ export const ldexp = dualImpl<typeof cpuLdexp>({

function cpuLength(value: number): number;
function cpuLength<T extends AnyFloatVecInstance>(value: T): number;
function cpuLength<T extends AnyFloatVecInstance | number>(value: T): number;
function cpuLength<T extends AnyFloatVecInstance | number>(value: T): number {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand Down Expand Up @@ -762,6 +792,7 @@ export const length = dualImpl({

function cpuLog(value: number): number;
function cpuLog<T extends AnyFloatVecInstance>(value: T): T;
function cpuLog<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuLog<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.log, [value]);
Expand All @@ -777,6 +808,7 @@ export const log = dualImpl({

function cpuLog2(value: number): number;
function cpuLog2<T extends AnyFloatVecInstance>(value: T): T;
function cpuLog2<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuLog2<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.log2, [value]);
Expand Down Expand Up @@ -876,9 +908,19 @@ const ModfResult = {
type ModfOverload = {
(value: number): Infer<(typeof ModfResult)['f32']>;
<T extends AnyFloatVecInstance>(value: T): Infer<(typeof ModfResult)[T['kind']]>;
<T extends AnyFloatVecInstance | number>(
value: T,
): T extends AnyFloatVecInstance
? Infer<(typeof ModfResult)[T['kind']]>
: Infer<(typeof ModfResult)['f32']>;
};
function cpuModf(e: number): Infer<(typeof ModfResult)['f32']>;
function cpuModf<T extends AnyFloatVecInstance>(e: T): Infer<(typeof ModfResult)[T['kind']]>;
function cpuModf<T extends AnyFloatVecInstance | number>(
e: T,
): T extends AnyFloatVecInstance
? Infer<(typeof ModfResult)[T['kind']]>
: Infer<(typeof ModfResult)['f32']>;
function cpuModf<T extends AnyFloatVecInstance | number>(
_value: T,
): Infer<(typeof ModfResult)[keyof typeof ModfResult]> {
Expand Down Expand Up @@ -933,6 +975,7 @@ export const pow = dualImpl({
});
function cpuQuantizeToF16(value: number): number;
function cpuQuantizeToF16<T extends AnyFloat32VecInstance>(value: T): T;
function cpuQuantizeToF16<T extends AnyFloat32VecInstance | number>(value: T): WidenNumber<T>;
function cpuQuantizeToF16<T extends AnyFloat32VecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand All @@ -954,7 +997,8 @@ export const quantizeToF16 = dualImpl<typeof cpuQuantizeToF16>({
});

function cpuRadians(value: number): number;
function cpuRadians<T extends AnyFloatVecInstance | number>(value: T): T;
function cpuRadians<T extends AnyFloatVecInstance>(value: T): T;
function cpuRadians<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuRadians<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand Down Expand Up @@ -1007,6 +1051,7 @@ export const refract = dualImpl<<T extends AnyFloatVecInstance>(e1: T, e2: T, e3
});
function cpuReverseBits(value: number): number;
function cpuReverseBits<T extends AnyIntegerVecInstance>(value: T): T;
function cpuReverseBits<T extends AnyIntegerVecInstance | number>(value: T): WidenNumber<T>;
function cpuReverseBits<T extends AnyIntegerVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand All @@ -1022,6 +1067,7 @@ export const reverseBits = dualImpl<typeof cpuReverseBits>({

function cpuRound(value: number): number;
function cpuRound<T extends AnyFloatVecInstance>(value: T): T;
function cpuRound<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuRound<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand Down Expand Up @@ -1049,6 +1095,7 @@ export const round = dualImpl({

function cpuSaturate(value: number): number;
function cpuSaturate<T extends AnyFloatVecInstance>(value: T): T;
function cpuSaturate<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuSaturate<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand All @@ -1069,6 +1116,8 @@ export const saturate = dualImpl({

function cpuSign(e: number): number;
function cpuSign<T extends AnySignedVecInstance>(e: T): T;
function cpuSign<T extends AnySignedVecInstance | number>(e: T): WidenNumber<T>;
function cpuSign<T extends AnySignedVecInstance>(e: T | number): T | number;
function cpuSign<T extends AnySignedVecInstance | number>(e: T): T {
assertKind(e, signedKind);
return generalizeFn(Math.sign, [e]);
Expand All @@ -1091,6 +1140,7 @@ export const sign = dualImpl({

function cpuSin(value: number): number;
function cpuSin<T extends AnyFloatVecInstance>(value: T): T;
function cpuSin<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuSin<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.sin, [value]);
Expand All @@ -1106,6 +1156,7 @@ export const sin = dualImpl({

function cpuSinh(value: number): number;
function cpuSinh<T extends AnyFloatVecInstance>(value: T): T;
function cpuSinh<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuSinh<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.sinh, [value]);
Expand Down Expand Up @@ -1137,6 +1188,7 @@ export const smoothstep = dualImpl({

function cpuSqrt(value: number): number;
function cpuSqrt<T extends AnyFloatVecInstance>(value: T): T;
function cpuSqrt<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuSqrt<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.sqrt, [value]);
Expand Down Expand Up @@ -1173,6 +1225,7 @@ export const step = dualImpl({

function cpuTan(value: number): number;
function cpuTan<T extends AnyFloatVecInstance>(value: T): T;
function cpuTan<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuTan<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
if (typeof value === 'number') {
Expand All @@ -1193,6 +1246,7 @@ export const tan = dualImpl({

function cpuTanh(value: number): number;
function cpuTanh<T extends AnyFloatVecInstance>(value: T): T;
function cpuTanh<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuTanh<T extends AnyFloatVecInstance | number>(value: T): T {
assertKind(value, floatKind);
return generalizeFn(Math.tanh, [value]);
Expand Down Expand Up @@ -1275,6 +1329,7 @@ export const transpose = dualImpl({

function cpuTrunc(value: number): number;
function cpuTrunc<T extends AnyFloatVecInstance>(value: T): T;
function cpuTrunc<T extends AnyFloatVecInstance | number>(value: T): WidenNumber<T>;
function cpuTrunc<T extends AnyFloatVecInstance | number>(_value: T): T {
throw new Error('Unreachable code. The function is only used for the type.');
}
Expand Down
3 changes: 3 additions & 0 deletions packages/typegpu/src/std/operators.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ import { unify } from '../tgsl/conversion.ts';
type NumVec = AnyNumericVecInstance;
type Mat = AnyMatInstance;

type WidenNumber<T> = T extends number ? number : T;

const getPrimitive = (t: BaseData): BaseData => ('primitive' in t ? (t.primitive as BaseData) : t);

const makeBinarySignature =
Expand Down Expand Up @@ -268,6 +270,7 @@ export const mod = dualImpl({

function cpuNeg(value: number): number;
function cpuNeg<T extends AnySignedVecInstance>(value: T): T;
function cpuNeg<T extends AnySignedVecInstance | number>(value: T): WidenNumber<T>;
function cpuNeg(value: NumVec | number): NumVec | number {
assertKind(value, signedKind);
return generalizeFn((value) => -value, [value]);
Expand Down
2 changes: 2 additions & 0 deletions packages/typegpu/src/std/subgroup.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,13 @@ import { unify } from '../tgsl/conversion.ts';
interface IdentityNumOrVec {
(e: number): number;
<T extends AnyNumericVecInstance>(e: T): T;
<T extends AnyNumericVecInstance | number>(e: T): T extends number ? number : T;
}

interface IdentityIntNumOrVec {
(e: number): number;
<T extends AnyIntegerVecInstance>(e: T): T;
<T extends AnyIntegerVecInstance | number>(e: T): T extends number ? number : T;
}

interface IdentityNumOrVecWithIdx {
Expand Down
Loading
Loading