chore: Remove 'createDualImpl' (#1961)

Author: iwoplaza

packages/typegpu/src/core/function/dualImpl.ts

@@ -1,42 +1,12 @@
 import type { DualFn } from '../../data/dualFn.ts';
-import {
-  type MapValueToSnippet,
-  snip,
-  type Snippet,
-} from '../../data/snippet.ts';
-import { inCodegenMode } from '../../execMode.ts';
-import { type FnArgsConversionHint, isKnownAtComptime } from '../../types.ts';
+import { type MapValueToSnippet, snip } from '../../data/snippet.ts';
+import { getResolutionCtx, inCodegenMode } from '../../execMode.ts';
+import { isKnownAtComptime, NormalState } from '../../types.ts';
 import { setName } from '../../shared/meta.ts';
 import { $internal } from '../../shared/symbols.ts';
 import { tryConvertSnippet } from '../../tgsl/conversion.ts';
 import type { AnyData } from '../../data/dataTypes.ts';
 
-export function createDualImpl<T extends (...args: never[]) => unknown>(
-  jsImpl: T,
-  gpuImpl: (...args: MapValueToSnippet<Parameters<T>>) => Snippet,
-  name: string,
-  argConversionHint: FnArgsConversionHint = 'keep',
-): DualFn<T> {
-  const impl = ((...args: Parameters<T>) => {
-    if (inCodegenMode()) {
-      return gpuImpl(...(args as MapValueToSnippet<Parameters<T>>)) as Snippet;
-    }
-    return jsImpl(...args);
-  }) as T;
-
-  setName(impl, name);
-  impl.toString = () => name;
-  Object.defineProperty(impl, $internal, {
-    value: {
-      jsImpl,
-      gpuImpl,
-      argConversionHint,
-    },
-  });
-
-  return impl as DualFn<T>;
-}
-
 type MapValueToDataType<T> = { [K in keyof T]: AnyData };
 
 interface DualImplOptions<T extends (...args: never[]) => unknown> {
@@ -68,6 +38,8 @@ export function dualImpl<T extends (...args: never[]) => unknown>(
   options: DualImplOptions<T>,
 ): DualFn<T> {
   const gpuImpl = (...args: MapValueToSnippet<Parameters<T>>) => {
+    // biome-ignore lint/style/noNonNullAssertion: it's there
+    const ctx = getResolutionCtx()!;
     const { argTypes, returnType } = typeof options.signature === 'function'
       ? options.signature(
         ...args.map((s) => {
@@ -94,6 +66,7 @@ export function dualImpl<T extends (...args: never[]) => unknown>(
       converted.every((s) => isKnownAtComptime(s)) &&
       typeof options.normalImpl === 'function'
     ) {
+      ctx.pushMode(new NormalState());
       try {
         return snip(
           options.normalImpl(...converted.map((s) => s.value) as never[]),
@@ -108,6 +81,8 @@ export function dualImpl<T extends (...args: never[]) => unknown>(
         if (!(e instanceof MissingCpuImplError)) {
           throw e;
         }
+      } finally {
+        ctx.popMode('normal');
       }
     }
 
@@ -135,9 +110,11 @@ export function dualImpl<T extends (...args: never[]) => unknown>(
     value: {
       jsImpl: options.normalImpl,
       gpuImpl,
-      strictSignature: typeof options.signature !== 'function'
-        ? options.signature
-        : undefined,
+      get strictSignature() {
+        return typeof options.signature !== 'function'
+          ? options.signature
+          : undefined;
+      },
       argConversionHint: 'keep',
     },
   });

packages/typegpu/src/core/valueProxyUtils.ts

@@ -1,7 +1,7 @@
 import type { Snippet } from '../data/snippet.ts';
 import { getGPUValue } from '../getGPUValue.ts';
 import { $internal, $ownSnippet, $resolve } from '../shared/symbols.ts';
-import { accessProp } from '../tgsl/generationHelpers.ts';
+import { accessProp } from '../tgsl/accessProp.ts';
 import {
   getOwnSnippet,
   type SelfResolvable,

packages/typegpu/src/data/array.ts

@@ -1,9 +1,7 @@
-import { createDualImpl } from '../core/function/dualImpl.ts';
+import { comptime } from '../core/function/comptime.ts';
 import { $internal } from '../shared/symbols.ts';
-import { UnknownData } from './dataTypes.ts';
-import { sizeOf } from './sizeOf.ts';
-import { snip, type Snippet } from './snippet.ts';
 import { schemaCallWrapper } from './schemaCallWrapper.ts';
+import { sizeOf } from './sizeOf.ts';
 import type { AnyWgslData, WgslArray } from './wgslTypes.ts';
 
 // ----------
@@ -37,39 +35,14 @@ interface WgslArrayConstructor {
  * @param elementType The type of elements in the array.
  * @param elementCount The number of elements in the array.
  */
-export const arrayOf = createDualImpl(
-  // JS implementation
+export const arrayOf = comptime(
   ((elementType, elementCount) => {
     if (elementCount === undefined) {
-      return (count: number) => cpu_arrayOf(elementType, count);
+      return comptime((count: number) => cpu_arrayOf(elementType, count));
     }
     return cpu_arrayOf(elementType, elementCount);
   }) as WgslArrayConstructor,
-  // CODEGEN implementation
-  (elementType, elementCount) => {
-    if (elementCount?.value === undefined) {
-      const partial = (count: Snippet) =>
-        arrayOf[$internal].gpuImpl(elementType, count);
-      // Marking so the WGSL generator lets this function through
-      partial[$internal] = true;
-
-      return snip(partial, UnknownData, /* origin */ 'constant');
-    }
-
-    if (typeof elementCount.value !== 'number') {
-      throw new Error(
-        `Cannot create array schema with count unknown at compile-time: '${elementCount.value}'`,
-      );
-    }
-
-    return snip(
-      cpu_arrayOf(elementType.value as AnyWgslData, elementCount.value),
-      elementType.value as AnyWgslData,
-      /* origin */ 'runtime',
-    );
-  },
-  'arrayOf',
-);
+).$name('arrayOf');
 
 // --------------
 // Implementation

packages/typegpu/src/data/disarray.ts

@@ -1,10 +1,7 @@
-import { createDualImpl } from '../core/function/dualImpl.ts';
+import { comptime } from '../core/function/comptime.ts';
 import { $internal } from '../shared/symbols.ts';
-import { UnknownData } from './dataTypes.ts';
-import { snip, type Snippet } from './snippet.ts';
 import type { AnyData, Disarray } from './dataTypes.ts';
 import { schemaCallWrapper } from './schemaCallWrapper.ts';
-import type { AnyWgslData } from './wgslTypes.ts';
 
 // ----------
 // Public API
@@ -43,39 +40,14 @@ interface DisarrayConstructor {
  * @param elementType The type of elements in the array.
  * @param elementCount The number of elements in the array.
  */
-export const disarrayOf = createDualImpl(
-  // JS implementation
+export const disarrayOf = comptime(
   ((elementType, elementCount) => {
     if (elementCount === undefined) {
       return (count: number) => cpu_disarrayOf(elementType, count);
     }
     return cpu_disarrayOf(elementType, elementCount);
   }) as DisarrayConstructor,
-  // CODEGEN implementation
-  (elementType, elementCount) => {
-    if (elementCount === undefined || elementCount.value === undefined) {
-      const partial = (count: Snippet) =>
-        disarrayOf[$internal].gpuImpl(elementType, count);
-      // Marking so the WGSL generator lets this function through
-      partial[$internal] = true;
-
-      return snip(partial, UnknownData, /* origin */ 'constant');
-    }
-
-    if (typeof elementCount.value !== 'number') {
-      throw new Error(
-        `Cannot create disarray schema with count unknown at compile-time: '${elementCount.value}'`,
-      );
-    }
-
-    return snip(
-      cpu_disarrayOf(elementType.value as AnyWgslData, elementCount.value),
-      elementType.value as AnyWgslData,
-      /* origin */ 'runtime',
-    );
-  },
-  'disarrayOf',
-);
+).$name('disarrayOf');
 
 export function cpu_disarrayOf<TElement extends AnyData>(
   elementType: TElement,

packages/typegpu/src/data/index.ts

@@ -2,10 +2,7 @@
  * @module typegpu/data
  */
 
-import {
-  type InfixOperator,
-  infixOperators,
-} from '../tgsl/generationHelpers.ts';
+import { type InfixOperator, infixOperators } from '../tgsl/accessProp.ts';
 import { $internal } from '../shared/symbols.ts';
 import { MatBase } from './matrix.ts';
 import { VecBase } from './vectorImpl.ts';