Merge branch 'main' into fix/tgpu-functions-declarations-order

Author: aleksanderkatan

packages/typegpu/src/data/numberOps.ts

@@ -0,0 +1,21 @@
+export const smoothstepScalar = (
+  edge0: number,
+  edge1: number,
+  x: number,
+): number => {
+  if (edge0 === edge1) {
+    return 0; // WGSL spec says this is an indeterminate value
+  }
+  const t = clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0);
+  return t * t * (3 - 2 * t);
+};
+
+export const clamp = (value: number, low: number, high: number) =>
+  Math.min(Math.max(low, value), high);
+
+export const divInteger = (lhs: number, rhs: number) => {
+  if (rhs === 0) {
+    return lhs;
+  }
+  return Math.trunc(lhs / rhs);
+};

packages/typegpu/src/data/vectorOps.ts

@@ -1,4 +1,5 @@
 import { mat2x2f, mat3x3f, mat4x4f } from './matrix.ts';
+import { clamp, divInteger, smoothstepScalar } from './numberOps.ts';
 import {
   vec2b,
   vec2f,
@@ -38,9 +39,6 @@ const dotVec3 = (lhs: v3, rhs: v3) =>
 const dotVec4 = (lhs: v4, rhs: v4) =>
   lhs.x * rhs.x + lhs.y * rhs.y + lhs.z * rhs.z + lhs.w * rhs.w;
 
-const clamp = (value: number, low: number, high: number) =>
-  Math.min(Math.max(low, value), high);
-
 type UnaryOp = (a: number) => number;
 type BinaryOp = (a: number, b: number) => number;
 
@@ -162,14 +160,41 @@ const binaryComponentWise4x4f =
     );
   };
 
-export const NumberOps = {
-  divInteger: (lhs: number, rhs: number) => {
-    if (rhs === 0) {
-      return lhs;
-    }
-    return Math.trunc(lhs / rhs);
-  },
-};
+type TernaryOp = (a: number, b: number, c: number) => number;
+
+const ternaryComponentWise2f =
+  (op: TernaryOp) => (a: wgsl.v2f, b: wgsl.v2f, c: wgsl.v2f) =>
+    vec2f(op(a.x, b.x, c.x), op(a.y, b.y, c.y));
+
+const ternaryComponentWise2h =
+  (op: TernaryOp) => (a: wgsl.v2h, b: wgsl.v2h, c: wgsl.v2h) =>
+    vec2h(op(a.x, b.x, c.x), op(a.y, b.y, c.y));
+
+const ternaryComponentWise3f =
+  (op: TernaryOp) => (a: wgsl.v3f, b: wgsl.v3f, c: wgsl.v3f) =>
+    vec3f(op(a.x, b.x, c.x), op(a.y, b.y, c.y), op(a.z, b.z, c.z));
+
+const ternaryComponentWise3h =
+  (op: TernaryOp) => (a: wgsl.v3h, b: wgsl.v3h, c: wgsl.v3h) =>
+    vec3h(op(a.x, b.x, c.x), op(a.y, b.y, c.y), op(a.z, b.z, c.z));
+
+const ternaryComponentWise4f =
+  (op: TernaryOp) => (a: wgsl.v4f, b: wgsl.v4f, c: wgsl.v4f) =>
+    vec4f(
+      op(a.x, b.x, c.x),
+      op(a.y, b.y, c.y),
+      op(a.z, b.z, c.z),
+      op(a.w, b.w, c.w),
+    );
+
+const ternaryComponentWise4h =
+  (op: TernaryOp) => (a: wgsl.v4h, b: wgsl.v4h, c: wgsl.v4h) =>
+    vec4h(
+      op(a.x, b.x, c.x),
+      op(a.y, b.y, c.y),
+      op(a.z, b.z, c.z),
+      op(a.w, b.w, c.w),
+    );
 
 export const VectorOps = {
   eq: {
@@ -446,6 +471,25 @@ export const VectorOps = {
     <T extends vBase | mBase>(lhs: T, rhs: T) => T
   >,
 
+  smoothstep: {
+    vec2f: ternaryComponentWise2f(smoothstepScalar),
+    vec2h: ternaryComponentWise2h(smoothstepScalar),
+    vec3f: ternaryComponentWise3f(smoothstepScalar),
+    vec3h: ternaryComponentWise3h(smoothstepScalar),
+    vec4f: ternaryComponentWise4f(smoothstepScalar),
+    vec4h: ternaryComponentWise4h(smoothstepScalar),
+  } as Record<
+    VecKind,
+    <T extends vBase>(
+      edge0: T,
+      edge1: T,
+      x: T,
+    ) => T extends wgsl.AnyVec2Instance ? wgsl.v2f
+      : T extends wgsl.AnyVec3Instance ? wgsl.v3f
+      : T extends wgsl.AnyVec4Instance ? wgsl.v4f
+      : wgsl.AnyVecInstance
+  >,
+
   addMixed: {
     vec2f: (a: wgsl.v2f, b: number) => unary2f((e) => e + b)(a),
     vec2h: (a: wgsl.v2h, b: number) => unary2h((e) => e + b)(a),
@@ -688,41 +732,35 @@ export const VectorOps = {
   div: {
     vec2f: binaryComponentWise2f((a, b) => a / b),
     vec2h: binaryComponentWise2h((a, b) => a / b),
-    vec2i: binaryComponentWise2i(NumberOps.divInteger),
-    vec2u: binaryComponentWise2u(NumberOps.divInteger),
+    vec2i: binaryComponentWise2i(divInteger),
+    vec2u: binaryComponentWise2u(divInteger),
 
     vec3f: binaryComponentWise3f((a, b) => a / b),
     vec3h: binaryComponentWise3h((a, b) => a / b),
-    vec3i: binaryComponentWise3i(NumberOps.divInteger),
-    vec3u: binaryComponentWise3u(NumberOps.divInteger),
+    vec3i: binaryComponentWise3i(divInteger),
+    vec3u: binaryComponentWise3u(divInteger),
 
     vec4f: binaryComponentWise4f((a, b) => a / b),
     vec4h: binaryComponentWise4h((a, b) => a / b),
-    vec4i: binaryComponentWise4i(NumberOps.divInteger),
-    vec4u: binaryComponentWise4u(NumberOps.divInteger),
+    vec4i: binaryComponentWise4i(divInteger),
+    vec4u: binaryComponentWise4u(divInteger),
   } as Record<VecKind, <T extends vBase>(a: T, b: T) => T>,
 
   divMixed: {
     vec2f: (a: wgsl.v2f, b: number) => unary2f((e) => e / b)(a),
     vec2h: (a: wgsl.v2h, b: number) => unary2h((e) => e / b)(a),
-    vec2i: (a: wgsl.v2i, b: number) =>
-      unary2i((e) => NumberOps.divInteger(e, b))(a),
-    vec2u: (a: wgsl.v2u, b: number) =>
-      unary2u((e) => NumberOps.divInteger(e, b))(a),
+    vec2i: (a: wgsl.v2i, b: number) => unary2i((e) => divInteger(e, b))(a),
+    vec2u: (a: wgsl.v2u, b: number) => unary2u((e) => divInteger(e, b))(a),
 
     vec3f: (a: wgsl.v3f, b: number) => unary3f((e) => e / b)(a),
     vec3h: (a: wgsl.v3h, b: number) => unary3h((e) => e / b)(a),
-    vec3i: (a: wgsl.v3i, b: number) =>
-      unary3i((e) => NumberOps.divInteger(e, b))(a),
-    vec3u: (a: wgsl.v3u, b: number) =>
-      unary3u((e) => NumberOps.divInteger(e, b))(a),
+    vec3i: (a: wgsl.v3i, b: number) => unary3i((e) => divInteger(e, b))(a),
+    vec3u: (a: wgsl.v3u, b: number) => unary3u((e) => divInteger(e, b))(a),
 
     vec4f: (a: wgsl.v4f, b: number) => unary4f((e) => e / b)(a),
     vec4h: (a: wgsl.v4h, b: number) => unary4h((e) => e / b)(a),
-    vec4i: (a: wgsl.v4i, b: number) =>
-      unary4i((e) => NumberOps.divInteger(e, b))(a),
-    vec4u: (a: wgsl.v4u, b: number) =>
-      unary4u((e) => NumberOps.divInteger(e, b))(a),
+    vec4i: (a: wgsl.v4i, b: number) => unary4i((e) => divInteger(e, b))(a),
+    vec4u: (a: wgsl.v4u, b: number) => unary4u((e) => divInteger(e, b))(a),
   } as Record<VecKind, <T extends vBase>(lhs: T, rhs: number) => T>,
 
   dot: {

packages/typegpu/src/std/index.ts

@@ -37,6 +37,7 @@ export {
   reflect,
   sign,
   sin,
+  smoothstep,
   sqrt,
   tanh
 } from './numeric.ts';

packages/typegpu/src/std/numeric.ts

@@ -1,5 +1,6 @@
 import { vecTypeToConstructor } from '../data/vector.ts';
 import { type AnyData, snip, type Snippet } from '../data/dataTypes.ts';
+import { smoothstepScalar } from '../data/numberOps.ts';
 import { f32 } from '../data/numeric.ts';
 import { VectorOps } from '../data/vectorOps.ts';
 import {
@@ -764,3 +765,25 @@ export const tanh = createDualImpl(
   (value) => snip(`tanh(${value.value})`, value.dataType),
   'tanh',
 );
+
+export const smoothstep = createDualImpl(
+  // CPU implementation
+  <T extends AnyFloatVecInstance | number>(edge0: T, edge1: T, x: T): T => {
+    if (typeof x === 'number') {
+      return smoothstepScalar(
+        edge0 as number,
+        edge1 as number,
+        x as number,
+      ) as T;
+    }
+    return VectorOps.smoothstep[x.kind](
+      edge0 as AnyFloatVecInstance,
+      edge1 as AnyFloatVecInstance,
+      x as AnyFloatVecInstance,
+    ) as T;
+  },
+  // GPU implementation
+  (edge0, edge1, x) =>
+    snip(`smoothstep(${edge0.value}, ${edge1.value}, ${x.value})`, x.dataType),
+  'smoothstep',
+);

packages/typegpu/tests/std/numeric/smoothstep.test.ts

@@ -0,0 +1,147 @@
+import { describe, expect, it } from 'vitest';
+import { vec2f, vec3f, vec4f } from '../../../src/data/index.ts';
+import { smoothstep } from '../../../src/std/index.ts';
+import { isCloseTo } from '../../../src/std/boolean.ts';
+
+describe('smoothstep', () => {
+  it('returns 0 when x is less than or equal to edge0', () => {
+    expect(smoothstep(0, 1, -0.5)).toBe(0);
+    expect(smoothstep(0, 1, 0)).toBe(0);
+  });
+
+  it('returns 1 when x is greater than or equal to edge1', () => {
+    expect(smoothstep(0, 1, 1)).toBe(1);
+    expect(smoothstep(0, 1, 1.5)).toBe(1);
+  });
+
+  it('returns smoothly interpolated value between 0 and 1 when x is between edge0 and edge1', () => {
+    expect(smoothstep(0, 1, 0.25)).toBeCloseTo(0.15625); // t = 0.25, t² * (3 - 2t) = 0.15625
+    expect(smoothstep(0, 1, 0.5)).toBeCloseTo(0.5); // t = 0.5, t² * (3 - 2t) = 0.5
+    expect(smoothstep(0, 1, 0.75)).toBeCloseTo(0.84375); // t = 0.75, t² * (3 - 2t) = 0.84375
+  });
+
+  it('works with vec2f', () => {
+    // all components less than edge0
+    const result1 = smoothstep(
+      vec2f(0.3, 0.5),
+      vec2f(0.8, 0.9),
+      vec2f(0.1, 0.2),
+    );
+    expect(isCloseTo(result1, vec2f(0, 0))).toBe(true);
+
+    // all components greater than edge1
+    const result2 = smoothstep(
+      vec2f(0.3, 0.5),
+      vec2f(0.8, 0.9),
+      vec2f(0.9, 1.0),
+    );
+    expect(isCloseTo(result2, vec2f(1, 1))).toBe(true);
+
+    // components between edge0 and edge1
+    const result3 = smoothstep(
+      vec2f(0.3, 0.5),
+      vec2f(0.8, 0.9),
+      vec2f(0.55, 0.7),
+    );
+    expect(isCloseTo(result3, vec2f(0.5, 0.5))).toBe(true);
+
+    // mixed results
+    const result4 = smoothstep(
+      vec2f(0.3, 0.5),
+      vec2f(0.8, 0.9),
+      vec2f(0.2, 0.95),
+    );
+    expect(isCloseTo(result4, vec2f(0, 1))).toBe(true);
+  });
+
+  it('works with vec3f', () => {
+    const result = smoothstep(
+      vec3f(0.1, 0.2, 0.3),
+      vec3f(0.6, 0.7, 0.8),
+      vec3f(0.35, 0.7, 0.55),
+    );
+    expect(isCloseTo(result, vec3f(0.5, 1, 0.5))).toBe(true);
+  });
+
+  it('works with vec4f', () => {
+    const result = smoothstep(
+      vec4f(0.0, 0.1, 0.2, 0.3),
+      vec4f(1.0, 0.9, 0.8, 0.7),
+      vec4f(0.0, 0.5, 0.8, 0.7),
+    );
+    expect(isCloseTo(result, vec4f(0, 0.5, 1, 1))).toBe(true);
+  });
+
+  it('works with vector edges with same values for all components', () => {
+    const result = smoothstep(
+      vec3f(0.3, 0.3, 0.3),
+      vec3f(0.7, 0.7, 0.7),
+      vec3f(0.2, 0.5, 0.8),
+    );
+    // For first component: x < edge0, result = 0
+    // For second component: t = (0.5-0.3)/(0.7-0.3) = 0.5, result = 0.5
+    // For third component: x > edge1, result = 1
+    expect(isCloseTo(result, vec3f(0, 0.5, 1))).toBe(true);
+  });
+
+  it('handles edge case with equal edge values', () => {
+    expect(smoothstep(0.5, 0.5, 0.4)).toBe(0);
+    expect(smoothstep(0.5, 0.5, 0.5)).toBe(0);
+    expect(smoothstep(0.5, 0.5, 0.6)).toBe(0);
+
+    const result = smoothstep(
+      vec2f(0.5, 0.3),
+      vec2f(0.5, 0.3),
+      vec2f(0.4, 0.3),
+    );
+    expect(isCloseTo(result, vec2f(0, 1))).toBe(false);
+  });
+
+  it('handles reversed edge values (edge0 > edge1)', () => {
+    expect(smoothstep(0.8, 0.2, 0.5)).toBeCloseTo(0.5);
+
+    const result = smoothstep(
+      vec2f(0.8, 0.7),
+      vec2f(0.2, 0.3),
+      vec2f(0.5, 0.5),
+    );
+    expect(isCloseTo(result, vec2f(0.5, 0.5))).toBe(true);
+  });
+
+  it('handles negative values correctly', () => {
+    expect(smoothstep(-2, -1, -1.5)).toBeCloseTo(0.5);
+    expect(smoothstep(-10, -5, -20)).toBe(0);
+    expect(smoothstep(-10, -5, 0)).toBe(1);
+
+    const result = smoothstep(
+      vec3f(-1, 0, -10),
+      vec3f(1, 1, -5),
+      vec3f(0, 0.5, -7.5),
+    );
+    expect(isCloseTo(result, vec3f(0.5, 0.5, 0.5))).toBe(true);
+  });
+
+  it('handles very small differences between edges', () => {
+    expect(smoothstep(0.5000, 0.5001, 0.50005)).toBeCloseTo(0.5);
+
+    const result = smoothstep(
+      vec2f(0.1, 0.2),
+      vec2f(0.1001, 0.2001),
+      vec2f(0.1, 0.20005),
+    );
+    expect(isCloseTo(result, vec2f(0, 0.5))).toBe(true);
+  });
+
+  it('handles extreme values', () => {
+    expect(smoothstep(1000, 2000, 1500)).toBeCloseTo(0.5);
+
+    expect(smoothstep(0.00001, 0.00002, 0.000015)).toBeCloseTo(0.5);
+
+    const result = smoothstep(
+      vec3f(0.00001, 1000, 0),
+      vec3f(0.00002, 2000, 1),
+      vec3f(0.000015, 1500, 0.5),
+    );
+    expect(isCloseTo(result, vec3f(0.5, 0.5, 0.5))).toBe(true);
+  });
+});