impr: console and Math method reference checks (#2438)

Author: aleksanderkatan

packages/typegpu/src/data/dataTypes.ts

@@ -1,4 +1,4 @@
-import { setName, type TgpuNamable } from '../shared/meta.ts';
+import { type TgpuNamable } from '../shared/meta.ts';
 import { isMarkedInternal } from '../shared/symbols.ts';
 import type {
   Infer,
@@ -23,7 +23,6 @@ import type {
   $reprPatch,
   $validVertexSchema,
 } from '../shared/symbols.ts';
-import { $internal } from '../shared/symbols.ts';
 import type { Prettify } from '../shared/utilityTypes.ts';
 import { vertexFormats } from '../shared/vertexFormat.ts';
 import type { WgslExternalTexture, WgslStorageTexture, WgslTexture } from './texture.ts';
@@ -261,13 +260,3 @@ export class MatrixColumnsAccess {
     this.matrix = matrix;
   }
 }
-
-export class ConsoleLog {
-  [$internal] = true;
-  readonly op: string;
-
-  constructor(op: string) {
-    this.op = op;
-    setName(this, 'consoleLog');
-  }
-}

packages/typegpu/src/resolutionCtx.ts

@@ -31,7 +31,7 @@ import {
   type TgpuLayoutEntry,
 } from './tgpuBindGroupLayout.ts';
 import { LogGeneratorImpl, LogGeneratorNullImpl } from './tgsl/consoleLog/logGenerator.ts';
-import type { LogGenerator, LogResources } from './tgsl/consoleLog/types.ts';
+import type { LogGenerator, LogResources, SupportedLogOp } from './tgsl/consoleLog/types.ts';
 import { getBestConversion } from './tgsl/conversion.ts';
 import { coerceToSnippet, concretize, numericLiteralToSnippet } from './tgsl/generationHelpers.ts';
 import type { ShaderGenerator } from './tgsl/shaderGenerator.ts';
@@ -458,7 +458,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     this._itemStateStack.clearBlockExternals();
   }
 
-  generateLog(op: string, args: Snippet[]): Snippet {
+  generateLog(op: SupportedLogOp, args: Snippet[]): Snippet {
     return this.#logGenerator.generateLog(this, op, args);
   }
 

packages/typegpu/src/tgsl/consoleLog/deserializers.ts

@@ -174,7 +174,7 @@ export function logDataFromGPU(resources: LogResources) {
         if (results.length === 0) {
           results.push('');
         }
-        console[op](
+        op.bind(console)(
           `%c${options.messagePrefix}%c ${results[0]}`,
           'background: #936ff5; color: white;',
           'color: inherit; background: none',

packages/typegpu/src/tgsl/consoleLog/logGenerator.ts

@@ -26,8 +26,7 @@ import {
   type LogMeta,
   type LogResources,
   type SerializedLogCallData,
-  type SupportedLogOps,
-  supportedLogOps,
+  type SupportedLogOp,
 } from './types.ts';
 
 const defaultOptions: Required<LogGeneratorOptions> = {
@@ -78,14 +77,9 @@ export class LogGeneratorImpl implements LogGenerator {
    * @param args Argument snippets. Snippets of UnknownType will be treated as string literals.
    * @returns A snippet containing the call to the logging function.
    */
-  generateLog(ctx: GenerationCtx, op: string, args: Snippet[]): Snippet {
+  generateLog(ctx: GenerationCtx, op: SupportedLogOp, args: Snippet[]): Snippet {
     if (shaderStageSlot.$ === 'vertex') {
-      console.warn(`'console.${op}' is not supported in vertex shaders.`);
-      return fallbackSnippet;
-    }
-
-    if (!supportedLogOps.includes(op as SupportedLogOps)) {
-      console.warn(`Unsupported log method '${op}'.`);
+      console.warn(`'console' operations are not supported in vertex shaders.`);
       return fallbackSnippet;
     }
 
@@ -122,7 +116,7 @@ export class LogGeneratorImpl implements LogGenerator {
     );
 
     this.#logIdToMeta.set(id, {
-      op: op as SupportedLogOps,
+      op,
       argTypes: concreteArgsWithStrings.map((e) =>
         e?.dataType === UnknownData ? (e?.value as string) : (e?.dataType as AnyWgslData),
       ),

packages/typegpu/src/tgsl/consoleLog/types.ts

@@ -2,6 +2,9 @@ import type { TgpuMutable } from '../../core/buffer/bufferShorthand.ts';
 import type { Snippet } from '../../data/snippet.ts';
 import type { AnyWgslData, Atomic, U32, WgslArray, WgslStruct } from '../../data/wgslTypes.ts';
 import type { GenerationCtx } from '../generationHelpers.ts';
+import type { supportedLogOps } from '../jsPolyfills.ts';
+
+export type SupportedLogOp = ReturnType<typeof supportedLogOps>[number];
 
 /**
  * Options for configuring GPU log generation.
@@ -32,7 +35,7 @@ export type SerializedLogCallData = WgslStruct<{
 }>;
 
 export interface LogMeta {
-  op: SupportedLogOps;
+  op: SupportedLogOp;
   argTypes: (string | AnyWgslData)[];
 }
 
@@ -52,10 +55,6 @@ export interface LogResources {
 }
 
 export interface LogGenerator {
-  generateLog(ctx: GenerationCtx, op: string, args: Snippet[]): Snippet;
+  generateLog(ctx: GenerationCtx, op: SupportedLogOp, args: Snippet[]): Snippet;
   get logResources(): LogResources | undefined;
 }
-
-export const supportedLogOps = ['log', 'debug', 'info', 'warn', 'error', 'clear'] as const;
-
-export type SupportedLogOps = (typeof supportedLogOps)[number];

packages/typegpu/src/tgsl/generationHelpers.ts

@@ -29,6 +29,7 @@ import type { ShelllessRepository } from './shellless.ts';
 import { stitch } from '../../src/core/resolve/stitch.ts';
 import { WgslTypeError } from '../../src/errors.ts';
 import { $internal, $resolve } from '../../src/shared/symbols.ts';
+import type { SupportedLogOp } from './consoleLog/types.ts';
 
 export function numericLiteralToSnippet(value: number): Snippet {
   if (value >= 2 ** 63 || value < -(2 ** 63)) {
@@ -85,7 +86,7 @@ export type GenerationCtx = ResolutionCtx & {
   dedent(): string;
   pushBlockScope(): void;
   popBlockScope(): void;
-  generateLog(op: string, args: Snippet[]): Snippet;
+  generateLog(op: SupportedLogOp, args: Snippet[]): Snippet;
   getById(id: string): Snippet | null;
   defineVariable(id: string, snippet: Snippet): void;
   setBlockExternals(externals: Record<string, Snippet>): void;

packages/typegpu/src/tgsl/jsPolyfills.ts

@@ -0,0 +1,78 @@
+import type { AnyFn } from '../core/function/fnTypes.ts';
+import { f32 } from '../data/numeric.ts';
+import {
+  abs,
+  acos,
+  acosh,
+  asin,
+  asinh,
+  atan,
+  atan2,
+  atanh,
+  ceil,
+  cos,
+  cosh,
+  countLeadingZeros,
+  exp,
+  floor,
+  log,
+  log2,
+  max,
+  min,
+  pow,
+  sign,
+  sin,
+  sinh,
+  sqrt,
+  tan,
+  tanh,
+  trunc,
+} from '../std/numeric.ts';
+import type { DualFn } from '../types.ts';
+
+export const mathToStd = new Map<AnyFn, DualFn<AnyFn>>([
+  // -- one to one Math to WGSL correlation --
+  [Math.abs, abs],
+  [Math.acos, acos],
+  [Math.acosh, acosh],
+  [Math.asin, asin],
+  [Math.asinh, asinh],
+  [Math.atan, atan],
+  [Math.atan2, atan2],
+  [Math.atanh, atanh],
+  [Math.ceil, ceil],
+  [Math.cos, cos],
+  [Math.cosh, cosh],
+  [Math.exp, exp],
+  [Math.floor, floor],
+  [Math.fround, f32 as DualFn<AnyFn>],
+  [Math.clz32, countLeadingZeros],
+  [Math.trunc, trunc],
+  [Math.log, log],
+  [Math.log2, log2],
+  [Math.pow, pow],
+  [Math.sign, sign],
+  [Math.sin, sin],
+  [Math.sinh, sinh],
+  [Math.sqrt, sqrt],
+  [Math.tan, tan],
+  [Math.tanh, tanh],
+  // -- varying in Math and two arg in WGSL, but we support varying in std --
+  [Math.max, max],
+  [Math.min, min],
+  // -- possible if we extend std --
+  // [Math.cbrt, ???],
+  // [Math.log10, ???],
+  // [Math.log1p, ???],
+  // [Math.f16round, ???],
+  // [Math.hypot, ???],
+  // [Math.expm1, ???],
+  // -- skipped --
+  // [Math.random, ???],
+  // [Math.imul, ???],
+  // [Math.round, ???], // round(2.5) is 3 in JS and 2 in WGSL
+]);
+
+// deferring the array creation lets us handle console mocks
+export const supportedLogOps = () =>
+  [console.log, console.debug, console.info, console.warn, console.error, console.clear] as const;