fix: convert array elements when casting between array types

packages/typegpu/src/tgsl/arrayExpression.ts

@@ -0,0 +1,61 @@
+import { isStoredInMemory } from '../data/snippet.ts';
+import type { ResolvedSnippet, Snippet } from '../data/snippet.ts';
+import {
+  type AnyWgslData,
+  type BaseData,
+  isNaturallyEphemeral,
+  type WgslArray,
+} from '../data/wgslTypes.ts';
+import { WgslTypeError } from '../errors.ts';
+import { $comptimeValueOf, $internal, $resolve } from '../shared/symbols.ts';
+import {
+  asComptime,
+  type ComptimeValue,
+  type ResolutionCtx,
+  type SelfResolvable,
+  type WithComptimeValue,
+} from '../types.ts';
+import { stringifySnippet } from './stringifySnippet.ts';
+
+export class ArrayExpression implements SelfResolvable, WithComptimeValue {
+  readonly [$internal] = true;
+  readonly type: WgslArray<AnyWgslData>;
+  readonly elements: Snippet[];
+
+  constructor(type: WgslArray<AnyWgslData>, elements: Snippet[]) {
+    this.type = type;
+    this.elements = elements;
+  }
+
+  toString(): string {
+    return 'ArrayExpression';
+  }
+
+  [$comptimeValueOf](): ComptimeValue {
+    // The array is known at comptime only if every element is known at comptime
+    const values: unknown[] = [];
+    for (const elem of this.elements) {
+      const known = asComptime(elem);
+      if (!known) {
+        return { comptime: false };
+      }
+      values.push(known.value);
+    }
+    return { comptime: true, value: values };
+  }
+
+  [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
+    for (const elem of this.elements) {
+      // We check if there are no references among the elements
+      if (isStoredInMemory(elem) && !isNaturallyEphemeral(elem.dataType)) {
+        const snippetStr = stringifySnippet(elem);
+        const snippetType = ctx.resolve(elem.dataType as BaseData).value;
+        throw new WgslTypeError(
+          `'${snippetStr}' reference cannot be used in an array constructor.\n-----\nTry '${snippetType}(${snippetStr})' or 'arrayOf(${snippetType}, count)([...])' to copy the value instead.\n-----`,
+        );
+      }
+    }
+
+    return ctx.gen.typeInstantiation(this.type, this.elements);
+  }
+}

packages/typegpu/src/tgsl/conversion.ts

@@ -13,9 +13,11 @@ import {
   isMat,
   isPtr,
   isVec,
+  isWgslArray,
   isWgslStruct,
   type Ptr,
   type U32,
+  type WgslArray,
   type WgslStruct,
 } from '../data/wgslTypes.ts';
 import { invariant, WgslTypeError } from '../errors.ts';
@@ -25,6 +27,7 @@ import { assertExhaustive } from '../shared/utilityTypes.ts';
 import { logger } from '../tgpuLogger.ts';
 import type { ResolutionCtx } from '../types.ts';
 import { accessStructProp } from './accessStructProp.ts';
+import { ArrayExpression } from './arrayExpression.ts';
 import { stringifySnippet } from './stringifySnippet.ts';
 
 type ConversionAction = 'ref' | 'deref' | 'cast' | 'none';
@@ -38,6 +41,19 @@ const INFINITE_RANK: ConversionRankInfo = {
   action: 'none',
 };
 
+function needsArrayCast(src: WgslArray, dest: WgslArray): boolean {
+  // An element count of 0 means runtime-sized, which can't be constructed element by element.
+  if (src === dest || src.elementCount === 0 || dest.elementCount === 0) {
+    return false;
+  }
+  const elementRank = getAutoConversionRank(src.elementType, dest.elementType);
+  return (
+    src.elementCount !== dest.elementCount ||
+    elementRank.rank === Number.POSITIVE_INFINITY ||
+    elementRank.action !== 'none'
+  );
+}
+
 function getAutoConversionRank(src: BaseData, dest: BaseData): ConversionRankInfo {
   const trueSrc = undecorate(src);
   const trueDst = undecorate(dest);
@@ -46,6 +62,12 @@ function getAutoConversionRank(src: BaseData, dest: BaseData): ConversionRankInf
     if (trueSrc.type === 'struct' && trueSrc !== trueDst) {
       return { rank: 1, action: 'cast', targetType: dest };
     }
+    if (isWgslArray(trueSrc) && isWgslArray(trueDst) && needsArrayCast(trueSrc, trueDst)) {
+      // Includes the cost of converting an element, so that `[[1], [1.5]]` becomes f32 in any row order.
+      const elementRank = getConversionRank(trueSrc.elementType, trueDst.elementType, true).rank;
+      const rank = 1 + (Number.isFinite(elementRank) ? elementRank : 0);
+      return { rank, action: 'cast', targetType: dest };
+    }
     return { rank: 0, action: 'none' };
   }
 
@@ -263,8 +285,12 @@ function applyActionToSnippet(
         snippet.origin,
         snippet.possibleSideEffects,
       );
-    case 'deref':
-      return derefSnippet(snippet);
+    case 'deref': {
+      const dereferenced = derefSnippet(snippet);
+      return isWgslArray(targetType) && isWgslArray(dereferenced.dataType)
+        ? tryConvertSnippet(ctx, dereferenced, targetType)
+        : dereferenced;
+    }
     case 'cast': {
       if (isWgslStruct(snippet.dataType) && isWgslStruct(targetType)) {
         const typeName = getName(snippet.dataType) ?? '<unnamed>';
@@ -300,6 +326,59 @@ function applyActionToSnippet(
           false,
         );
       }
+      if (isWgslArray(snippet.dataType) && isWgslArray(targetType)) {
+        const typeName = String(snippet.dataType);
+        const targetName = String(targetType);
+
+        const convertElement = (element: Snippet) => {
+          const converted = convertToCommonType(ctx, [element], [targetType.elementType]);
+          if (!converted || !converted[0]) {
+            throw new Error(
+              `Cannot auto-convert array '${typeName}' to '${targetName}' because its elements are not convertible.`,
+            );
+          }
+          return converted[0];
+        };
+
+        // Array to array casting
+        if (snippet.dataType.elementCount !== targetType.elementCount) {
+          throw new Error(
+            `Cannot auto-convert array '${typeName}' to '${targetName}' because their element counts differ.`,
+          );
+        }
+
+        // An array literal stays one, so that its elements can still be known at comptime.
+        if (snippet.value instanceof ArrayExpression) {
+          return snip(
+            new ArrayExpression(
+              targetType as WgslArray<AnyWgslData>,
+              snippet.value.elements.map(convertElement),
+            ),
+            targetType,
+            snippet.origin,
+            snippet.possibleSideEffects,
+          );
+        }
+
+        if (snippet.possibleSideEffects) {
+          throw new Error(
+            `Cannot resolve array cast from '${typeName}' to '${targetName}'. Store the value to a variable first, then cast it.`,
+          );
+        }
+
+        const source = ctx.resolveSnippet(snippet).value;
+        const elementType = snippet.dataType.elementType;
+        const elementSnips = Array.from({ length: targetType.elementCount }, (_, index) =>
+          convertElement(snip(`${source}[${index}]`, elementType, 'runtime', false)),
+        );
+        const targetSnippet = ctx.resolve(targetType).value;
+        return snip(
+          `${targetSnippet}(${elementSnips.map((elem) => ctx.resolveSnippet(elem).value).join(', ')})`,
+          targetType,
+          'runtime',
+          false,
+        );
+      }
       // Casting means calling the schema with the snippet as an argument.
       return schemaCallWrapperGPU(ctx, targetType, snippet);
     }

packages/typegpu/src/tgsl/generationHelpers.ts

@@ -1,31 +1,21 @@
 import { UnknownData } from '../data/dataTypes.ts';
 import { abstractFloat, abstractInt, bool, f32, i32 } from '../data/numeric.ts';
 import { isRef } from '../data/ref.ts';
-import { isSnippet, isStoredInMemory, snip, withDataType } from '../data/snippet.ts';
-import type { ResolvedSnippet, Snippet } from '../data/snippet.ts';
+import { isSnippet, snip, withDataType } from '../data/snippet.ts';
+import type { Snippet } from '../data/snippet.ts';
 import {
   type AnyWgslData,
   type BaseData,
   type F32,
   type I32,
   isMatInstance,
-  isNaturallyEphemeral,
   isVecInstance,
-  type WgslArray,
   WORKAROUND_getSchema,
 } from '../data/wgslTypes.ts';
-import {
-  asComptime,
-  type ComptimeValue,
-  getOwnSnippet,
-  type ResolutionCtx,
-  type SelfResolvable,
-  type WithComptimeValue,
-} from '../types.ts';
-import { WgslTypeError } from '../errors.ts';
-import { $comptimeValueOf, $internal, $resolve } from '../shared/symbols.ts';
+import { getOwnSnippet } from '../types.ts';
 import { logger } from '../tgpuLogger.ts';
-import { stringifySnippet } from './stringifySnippet.ts';
+
+export { ArrayExpression } from './arrayExpression.ts';
 
 export function numericLiteralToSnippet(value: number): Snippet {
   if (value >= 2 ** 63 || value < -(2 ** 63)) {
@@ -106,45 +96,3 @@ export function coerceToSnippet(value: unknown): Snippet {
  * Defers resolution. Stores array elements as snippets so the
  * generator can access them when needed.
  */
-export class ArrayExpression implements SelfResolvable, WithComptimeValue {
-  readonly [$internal] = true;
-  readonly type: WgslArray<AnyWgslData>;
-  readonly elements: Snippet[];
-
-  constructor(type: WgslArray<AnyWgslData>, elements: Snippet[]) {
-    this.type = type;
-    this.elements = elements;
-  }
-
-  toString(): string {
-    return 'ArrayExpression';
-  }
-
-  [$comptimeValueOf](): ComptimeValue {
-    // The array is known at comptime only if every element is known at comptime
-    const values: unknown[] = [];
-    for (const elem of this.elements) {
-      const known = asComptime(elem);
-      if (!known) {
-        return { comptime: false };
-      }
-      values.push(known.value);
-    }
-    return { comptime: true, value: values };
-  }
-
-  [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
-    for (const elem of this.elements) {
-      // We check if there are no references among the elements
-      if (isStoredInMemory(elem) && !isNaturallyEphemeral(elem.dataType)) {
-        const snippetStr = stringifySnippet(elem);
-        const snippetType = ctx.resolve(concretize(elem.dataType as BaseData)).value;
-        throw new WgslTypeError(
-          `'${snippetStr}' reference cannot be used in an array constructor.\n-----\nTry '${snippetType}(${snippetStr})' or 'arrayOf(${snippetType}, count)([...])' to copy the value instead.\n-----`,
-        );
-      }
-    }
-
-    return ctx.gen.typeInstantiation(this.type, this.elements);
-  }
-}