fix: Function call member access type inference (#1062)

* Fix function member access type inference

* Add WGSL void type and error handling

* Move $repr symbol from type import to value import

* Refactor internal symbol and dual implementation

* Set Internal Implementation to jsImpl

* Review fixes

---------

Co-authored-by: Iwo Plaza <iwoplaza@gmail.com>

Author: reczkok

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

@@ -221,11 +221,7 @@ function createFn<
     }),
   );
 
-  Object.defineProperty(call, $internal, {
-    value: {
-      implementation,
-    },
-  });
+  call[$internal].implementation = implementation;
 
   const fn = Object.assign(call, fnBase as This) as unknown as TgpuFn<
     Args,
@@ -310,11 +306,7 @@ function createBoundFunction<
     },
   });
 
-  Object.defineProperty(fn, $internal, {
-    value: {
-      implementation: innerFn[$internal].implementation,
-    },
-  });
+  fn[$internal].implementation = innerFn[$internal].implementation;
 
   return fn;
 }

packages/typegpu/src/core/resolve/resolveData.ts

@@ -257,5 +257,9 @@ export function resolveData(ctx: ResolutionCtx, data: AnyData): string {
     throw new Error('Abstract types have no concrete representation in WGSL');
   }
 
+  if (data.type === 'void') {
+    throw new Error('Void has no representation in WGSL');
+  }
+
   assertExhaustive(data, 'resolveData');
 }

packages/typegpu/src/data/dataTypes.ts

@@ -6,11 +6,19 @@ import type {
   InferPartialRecord,
   InferRecord,
 } from '../shared/repr.js';
+import { $internal } from '../shared/symbols.js';
 import type { Prettify } from '../shared/utilityTypes.js';
 import { vertexFormats } from '../shared/vertexFormat.js';
 import type { PackedData } from './vertexFormatData.js';
 import * as wgsl from './wgslTypes.js';
 
+export type TgpuDualFn<TImpl extends (...args: unknown[]) => unknown> =
+  TImpl & {
+    [$internal]: {
+      implementation: TImpl | string;
+    };
+  };
+
 /**
  * Array schema constructed via `d.disarrayOf` function.
  *
@@ -139,5 +147,5 @@ export function isData(value: unknown): value is AnyData {
 export type AnyData = wgsl.AnyWgslData | AnyLooseData;
 export type AnyConcreteData = Exclude<
   AnyData,
-  wgsl.AbstractInt | wgsl.AbstractFloat
+  wgsl.AbstractInt | wgsl.AbstractFloat | wgsl.Void
 >;

packages/typegpu/src/data/matrix.ts

@@ -3,6 +3,7 @@ import { $repr } from '../shared/repr.js';
 import type { SelfResolvable } from '../types';
 import { vec2f, vec3f, vec4f } from './vector';
 import type {
+  AnyWgslData,
   Mat2x2f,
   Mat3x3f,
   Mat4x4f,
@@ -53,7 +54,7 @@ function createMatSchema<
     [$repr]: undefined as unknown as ValueType,
     type: options.type,
     label: options.type,
-  };
+  } as unknown as AnyWgslData;
 
   const construct = createDualImpl(
     // CPU implementation

packages/typegpu/src/data/numeric.ts

@@ -60,7 +60,7 @@ const u32Cast = createDualImpl(
  */
 export const u32: U32 = Object.assign(u32Cast, {
   type: 'u32',
-}) as U32;
+}) as unknown as U32;
 
 const i32Cast = createDualImpl(
   // CPU implementation
@@ -99,7 +99,7 @@ const i32Cast = createDualImpl(
  */
 export const i32: I32 = Object.assign(i32Cast, {
   type: 'i32',
-}) as I32;
+}) as unknown as I32;
 
 const f32Cast = createDualImpl(
   // CPU implementation
@@ -127,7 +127,7 @@ const f32Cast = createDualImpl(
  */
 export const f32: F32 = Object.assign(f32Cast, {
   type: 'f32',
-}) as F32;
+}) as unknown as F32;
 
 const f16Cast = createDualImpl(
   // CPU implementation
@@ -158,4 +158,4 @@ const f16Cast = createDualImpl(
  */
 export const f16: F16 = Object.assign(f16Cast, {
   type: 'f16',
-}) as F16;
+}) as unknown as F16;

packages/typegpu/src/data/wgslTypes.ts

@@ -1,6 +1,5 @@
 import type { TgpuNamable } from '../namable.js';
 import type {
-  $repr,
   Infer,
   InferGPU,
   InferGPURecord,
@@ -10,6 +9,7 @@ import type {
   MemIdentity,
   MemIdentityRecord,
 } from '../shared/repr.js';
+import { $repr } from '../shared/repr.js';
 import type { Prettify } from '../shared/utilityTypes.js';
 
 type DecoratedLocation<T extends BaseData> = Decorated<T, Location<number>[]>;
@@ -50,6 +50,12 @@ export interface AbstractFloat {
   readonly [$repr]: number;
 }
 
+export interface Void {
+  readonly type: 'void';
+  readonly [$repr]: undefined;
+}
+export const Void: Void = { type: 'void', [$repr]: undefined };
+
 interface Swizzle2<T2, T3, T4> {
   readonly xx: T2;
   readonly xy: T2;
@@ -1158,6 +1164,7 @@ export const wgslTypeLiterals = [
   'decorated',
   'abstractInt',
   'abstractFloat',
+  'void',
 ] as const;
 
 export type WgslTypeLiteral = (typeof wgslTypeLiterals)[number];
@@ -1222,7 +1229,8 @@ export type AnyWgslData =
   | Atomic
   | Decorated
   | AbstractInt
-  | AbstractFloat;
+  | AbstractFloat
+  | Void;
 
 // #endregion
 

packages/typegpu/src/shared/generators.ts

@@ -1,5 +1,7 @@
+import type { TgpuDualFn } from '../data/dataTypes';
 import { inGPUMode } from '../gpuMode';
 import type { Resource } from '../types';
+import { $internal } from './symbols';
 
 /**
  * Yields values in the sequence 0,1,2..∞ except for the ones in the `excluded` set.
@@ -24,8 +26,8 @@ type MapValueToResource<T> = { [K in keyof T]: Resource };
 export function createDualImpl<T extends (...args: any[]) => any>(
   jsImpl: T,
   gpuImpl: (...args: MapValueToResource<Parameters<T>>) => Resource,
-): T {
-  return ((...args: Parameters<T>) => {
+): TgpuDualFn<T> {
+  const impl = ((...args: Parameters<T>) => {
     if (inGPUMode()) {
       return gpuImpl(
         ...(args as unknown as MapValueToResource<Parameters<T>>),
@@ -34,4 +36,12 @@ export function createDualImpl<T extends (...args: any[]) => any>(
     // biome-ignore lint/suspicious/noExplicitAny: <it's very convenient>
     return jsImpl(...(args as any));
   }) as T;
+
+  Object.defineProperty(impl, $internal, {
+    value: {
+      implementation: jsImpl,
+    },
+  });
+
+  return impl as TgpuDualFn<T>;
 }

packages/typegpu/src/smol/generationHelpers.ts

@@ -1,4 +1,5 @@
 import { isDerived, isSlot } from '../core/slot/slotTypes';
+import type { AnyData } from '../data/dataTypes';
 import { mat2x2f, mat3x3f, mat4x4f } from '../data/matrix';
 import {
   abstractFloat,
@@ -27,8 +28,6 @@ import {
   vec4u,
 } from '../data/vector';
 import {
-  type AnyWgslData,
-  type BaseData,
   type WgslStruct,
   isDecorated,
   isWgslArray,
@@ -67,7 +66,7 @@ type SwizzleLength = 1 | 2 | 3 | 4;
 
 const swizzleLenToType: Record<
   SwizzleableType,
-  Record<SwizzleLength, AnyWgslData>
+  Record<SwizzleLength, AnyData>
 > = {
   f: {
     1: f32,
@@ -146,7 +145,7 @@ const indexableTypeToResult = {
 export function getTypeForPropAccess(
   targetType: Wgsl,
   propName: string,
-): BaseData | UnknownData {
+): AnyData | UnknownData {
   if (
     typeof targetType === 'string' ||
     typeof targetType === 'number' ||
@@ -164,23 +163,25 @@ export function getTypeForPropAccess(
     }
     const unwrapped = ctx.unwrap(targetType);
 
-    return getTypeFromWgsl(unwrapped as Wgsl) as BaseData;
+    return getTypeFromWgsl(unwrapped);
   }
 
-  let target = targetType as BaseData;
+  let target = targetType as AnyData;
 
   if (hasInternalDataType(target)) {
-    target = target[$internal].dataType;
+    target = target[$internal].dataType as AnyData;
   }
   while (isDecorated(target)) {
-    target = target.inner;
+    target = target.inner as AnyData;
   }
 
   const targetTypeStr =
     'kind' in target ? (target.kind as string) : target.type;
 
   if (targetTypeStr === 'struct') {
-    return (target as WgslStruct).propTypes[propName] ?? UnknownData;
+    return (
+      ((target as WgslStruct).propTypes[propName] as AnyData) ?? UnknownData
+    );
   }
 
   const propLength = propName.length;
@@ -204,11 +205,11 @@ export function getTypeForPropAccess(
   return isWgslData(target) ? target : UnknownData;
 }
 
-export function getTypeForIndexAccess(resource: Wgsl): BaseData | UnknownData {
+export function getTypeForIndexAccess(resource: Wgsl): AnyData | UnknownData {
   if (isWgslData(resource)) {
     // array
     if (isWgslArray(resource)) {
-      return resource.elementType;
+      return resource.elementType as AnyData;
     }
 
     // vector or matrix
@@ -222,7 +223,7 @@ export function getTypeForIndexAccess(resource: Wgsl): BaseData | UnknownData {
   return UnknownData;
 }
 
-export function getTypeFromWgsl(resource: Wgsl): BaseData | UnknownData {
+export function getTypeFromWgsl(resource: Wgsl): AnyData | UnknownData {
   if (isDerived(resource) || isSlot(resource)) {
     return getTypeFromWgsl(resource.value as Wgsl);
   }

packages/typegpu/src/smol/wgslGenerator.ts

@@ -1,12 +1,13 @@
 import type * as smol from 'tinyest';
 import * as d from '../data';
+import type { AnyData } from '../data/dataTypes';
 import { abstractInt } from '../data/numeric.js';
 import * as wgsl from '../data/wgslTypes.js';
 import {
   type ResolutionCtx,
   type Resource,
   UnknownData,
-  type Wgsl,
+  isMarkedInternal,
   isWgsl,
 } from '../types.js';
 import {
@@ -46,8 +47,8 @@ type Operator =
   | smol.UnaryOperator;
 
 function operatorToType<
-  TL extends wgsl.AnyWgslData | UnknownData,
-  TR extends wgsl.AnyWgslData | UnknownData,
+  TL extends AnyData | UnknownData,
+  TR extends AnyData | UnknownData,
 >(lhs: TL, op: Operator, rhs?: TR): TL | TR | wgsl.Bool {
   if (!rhs) {
     if (op === '!' || op === '~') {
@@ -190,7 +191,9 @@ export function generateExpression(
     if (typeof target.value === 'string') {
       return {
         value: `${target.value}.${property}`,
-        dataType: getTypeForPropAccess(target.dataType as Wgsl, property),
+        dataType: d.isData(target.dataType)
+          ? getTypeForPropAccess(target.dataType, property)
+          : UnknownData,
       };
     }
 
@@ -214,10 +217,24 @@ export function generateExpression(
     // biome-ignore lint/suspicious/noExplicitAny: <sorry TypeScript>
     const propValue = (target.value as any)[property];
 
+    if (target.dataType.type !== 'unknown') {
+      if (wgsl.isMat(target.dataType) && property === 'columns') {
+        return {
+          value: target.value,
+          dataType: target.dataType,
+        };
+      }
+
+      return {
+        value: propValue,
+        dataType: getTypeForPropAccess(target.dataType, property),
+      };
+    }
+
     if (isWgsl(target.value)) {
       return {
         value: propValue,
-        dataType: getTypeForPropAccess(target.value as d.AnyWgslData, property),
+        dataType: getTypeForPropAccess(target.value, property),
       };
     }
 
@@ -245,7 +262,9 @@ export function generateExpression(
 
     return {
       value: `${targetStr}[${propertyStr}]`,
-      dataType: getTypeForIndexAccess(targetExpr.dataType as d.AnyWgslData),
+      dataType: d.isData(targetExpr.dataType)
+        ? getTypeForIndexAccess(targetExpr.dataType)
+        : UnknownData,
     };
   }
 
@@ -291,10 +310,21 @@ export function generateExpression(
       };
     }
 
+    if (!isMarkedInternal(idValue)) {
+      throw new Error(
+        `Function ${String(idValue)} has not been created using TypeGPU APIs. Did you mean to wrap the function with tgpu.fn(args, return)(...) ?`,
+      );
+    }
+
     // Assuming that `id` is callable
-    return (idValue as unknown as (...args: unknown[]) => unknown)(
+    const fnRes = (idValue as unknown as (...args: unknown[]) => unknown)(
       ...resolvedResources,
     ) as Resource;
+
+    return {
+      value: resolveRes(ctx, fnRes),
+      dataType: fnRes.dataType,
+    };
   }
 
   if ('o' in expression) {
@@ -369,7 +399,7 @@ export function generateExpression(
 
     return {
       value: `${arrayType}( ${arrayValues.join(', ')} )`,
-      dataType: d.arrayOf(type as d.AnyWgslData, values.length),
+      dataType: d.arrayOf(type, values.length) as d.AnyWgslData,
     };
   }
 
@@ -449,6 +479,10 @@ ${alternate}`;
       throw new Error('Cannot create variable without an initial value.');
     }
 
+    if (d.isLooseData(eq.dataType)) {
+      throw new Error('Cannot create variable with loose data type.');
+    }
+
     registerBlockVariable(ctx, rawId, eq.dataType);
     const id = resolveRes(ctx, generateIdentifier(ctx, rawId));