chore: Use public API in test utilities (extractSnippet, ...)

packages/typegpu/src/tgsl/shaderGenerator_members.ts

@@ -2,3 +2,5 @@ export { UnknownData } from '../data/dataTypes.ts';
 
 // types
 export type { ResolutionCtx } from '../types.ts';
+export type { Snippet } from '../data/snippet.ts';
+export type { Origin } from '../data/snippet.ts';

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -910,6 +910,75 @@ ${this.ctx.pre}}`;
     return snip(stitch`${this.ctx.resolve(schema).value}(${args})`, schema, 'runtime');
   }
 
+  public _return(statement: tinyest.Return): string {
+    const returnNode = statement[1];
+
+    if (returnNode !== undefined) {
+      const expectedReturnType = this.ctx.topFunctionReturnType;
+      let returnSnippet = expectedReturnType
+        ? this._typedExpression(returnNode, expectedReturnType)
+        : this._expression(returnNode);
+
+      if (returnSnippet.value instanceof RefOperator) {
+        throw new WgslTypeError(
+          stitch`Cannot return references, returning '${returnSnippet.value.snippet}'`,
+        );
+      }
+
+      // Arguments cannot be returned from functions without copying. A simple example why is:
+      // const identity = (x) => {
+      //   'use gpu';
+      //   return x;
+      // };
+      //
+      // const foo = (arg: d.v3f) => {
+      //   'use gpu';
+      //   const marg = identity(arg);
+      //   marg.x = 1; // 'marg's origin would be 'runtime', so we wouldn't be able to track this misuse.
+      // };
+      if (
+        returnSnippet.origin === 'argument' &&
+        !wgsl.isNaturallyEphemeral(returnSnippet.dataType) &&
+        // Only restricting this use in non-entry functions, as the function
+        // is giving up ownership of all references anyway.
+        this.ctx.topFunctionScope?.functionType === 'normal'
+      ) {
+        throw new WgslTypeError(
+          stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`,
+        );
+      }
+
+      if (
+        !expectedReturnType &&
+        !isEphemeralSnippet(returnSnippet) &&
+        returnSnippet.origin !== 'this-function'
+      ) {
+        const str = this.ctx.resolve(returnSnippet.value, returnSnippet.dataType).value;
+        const typeStr = this.ctx.resolve(unptr(returnSnippet.dataType)).value;
+        throw new WgslTypeError(
+          `'return ${str};' is invalid, cannot return references.
+-----
+Try 'return ${typeStr}(${str});' instead.
+-----`,
+        );
+      }
+
+      returnSnippet = tryConvertSnippet(
+        this.ctx,
+        returnSnippet,
+        unptr(returnSnippet.dataType) as wgsl.AnyWgslData,
+        false,
+      );
+
+      invariant(returnSnippet.dataType !== UnknownData, 'Return type should be known');
+
+      this.ctx.reportReturnType(returnSnippet.dataType);
+      return stitch`${this.ctx.pre}return ${returnSnippet};`;
+    }
+
+    return `${this.ctx.pre}return;`;
+  }
+
   public _statement(statement: tinyest.Statement): string {
     if (typeof statement === 'string') {
       const id = this._identifier(statement);
@@ -923,72 +992,7 @@ ${this.ctx.pre}}`;
     }
 
     if (statement[0] === NODE.return) {
-      const returnNode = statement[1];
-
-      if (returnNode !== undefined) {
-        const expectedReturnType = this.ctx.topFunctionReturnType;
-        let returnSnippet = expectedReturnType
-          ? this._typedExpression(returnNode, expectedReturnType)
-          : this._expression(returnNode);
-
-        if (returnSnippet.value instanceof RefOperator) {
-          throw new WgslTypeError(
-            stitch`Cannot return references, returning '${returnSnippet.value.snippet}'`,
-          );
-        }
-
-        // Arguments cannot be returned from functions without copying. A simple example why is:
-        // const identity = (x) => {
-        //   'use gpu';
-        //   return x;
-        // };
-        //
-        // const foo = (arg: d.v3f) => {
-        //   'use gpu';
-        //   const marg = identity(arg);
-        //   marg.x = 1; // 'marg's origin would be 'runtime', so we wouldn't be able to track this misuse.
-        // };
-        if (
-          returnSnippet.origin === 'argument' &&
-          !wgsl.isNaturallyEphemeral(returnSnippet.dataType) &&
-          // Only restricting this use in non-entry functions, as the function
-          // is giving up ownership of all references anyway.
-          this.ctx.topFunctionScope?.functionType === 'normal'
-        ) {
-          throw new WgslTypeError(
-            stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`,
-          );
-        }
-
-        if (
-          !expectedReturnType &&
-          !isEphemeralSnippet(returnSnippet) &&
-          returnSnippet.origin !== 'this-function'
-        ) {
-          const str = this.ctx.resolve(returnSnippet.value, returnSnippet.dataType).value;
-          const typeStr = this.ctx.resolve(unptr(returnSnippet.dataType)).value;
-          throw new WgslTypeError(
-            `'return ${str};' is invalid, cannot return references.
------
-Try 'return ${typeStr}(${str});' instead.
------`,
-          );
-        }
-
-        returnSnippet = tryConvertSnippet(
-          this.ctx,
-          returnSnippet,
-          unptr(returnSnippet.dataType) as wgsl.AnyWgslData,
-          false,
-        );
-
-        invariant(returnSnippet.dataType !== UnknownData, 'Return type should be known');
-
-        this.ctx.reportReturnType(returnSnippet.dataType);
-        return stitch`${this.ctx.pre}return ${returnSnippet};`;
-      }
-
-      return `${this.ctx.pre}return;`;
+      return this._return(statement);
     }
 
     if (statement[0] === NODE.if) {

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

@@ -104,22 +104,22 @@ describe('add', () => {
     it('infers types when adding constants', () => {
       const int_int = () => {
         'use gpu';
-        1 + 2;
+        return 1 + 2;
       };
 
       const float_float = () => {
         'use gpu';
-        1.1 + 2.3;
+        return 1.1 + 2.3;
       };
 
       const int_float = () => {
         'use gpu';
-        1.1 + 2;
+        return 1.1 + 2;
       };
 
       const float_int = () => {
         'use gpu';
-        1 + 2.3;
+        return 1 + 2.3;
       };
 
       expectDataTypeOf(int_int).toBe(abstractInt);

packages/typegpu/tests/tgsl/memberAccess.test.ts

@@ -1,7 +1,6 @@
 import { describe } from 'vitest';
 import { it } from 'typegpu-testing-utility';
 import { expectSnippetOf } from '../utils/parseResolved.ts';
-import { snip } from '../../src/data/snippet.ts';
 import tgpu, { d } from '../../src/index.js';
 
 describe('Member Access', () => {
@@ -12,41 +11,41 @@ describe('Member Access', () => {
   it('should access member properties of literals', () => {
     expectSnippetOf(() => {
       'use gpu';
-      Boid().pos;
-    }).toStrictEqual(snip('Boid().pos', d.vec3f, 'runtime'));
+      return Boid().pos;
+    }).toStrictEqual(['Boid().pos', d.vec3f, 'runtime']);
 
     expectSnippetOf(() => {
       'use gpu';
-      Boid().pos.xyz;
-    }).toStrictEqual(snip('Boid().pos.xyz', d.vec3f, 'runtime'));
+      return Boid().pos.xyz;
+    }).toStrictEqual(['Boid().pos.xyz', d.vec3f, 'runtime']);
   });
 
   it('should access member properties of externals', () => {
     const boid = Boid({ pos: d.vec3f(1, 2, 3) });
 
     expectSnippetOf(() => {
       'use gpu';
-      boid.pos;
-    }).toStrictEqual(snip(d.vec3f(1, 2, 3), d.vec3f, 'constant'));
+      return boid.pos;
+    }).toStrictEqual([d.vec3f(1, 2, 3), d.vec3f, 'constant']);
 
     expectSnippetOf(() => {
       'use gpu';
-      boid.pos.zyx;
-    }).toStrictEqual(snip(d.vec3f(3, 2, 1), d.vec3f, 'constant'));
+      return boid.pos.zyx;
+    }).toStrictEqual([d.vec3f(3, 2, 1), d.vec3f, 'constant']);
   });
 
   it('should access member properties of variables', () => {
     const boidVar = tgpu.privateVar(Boid);
 
     expectSnippetOf(() => {
       'use gpu';
-      boidVar.$.pos;
-    }).toStrictEqual(snip('boidVar.pos', d.vec3f, 'private'));
+      return boidVar.$.pos;
+    }).toStrictEqual(['boidVar.pos', d.vec3f, 'private']);
 
     expectSnippetOf(() => {
       'use gpu';
-      boidVar.$.pos.xyz;
-    }).toStrictEqual(snip('boidVar.pos.xyz', d.vec3f, 'runtime')); // < swizzles are new objects
+      return boidVar.$.pos.xyz;
+    }).toStrictEqual(['boidVar.pos.xyz', d.vec3f, 'runtime']); // < swizzles are new objects
   });
 
   it('derefs access to local variables with proper address space', () => {
@@ -56,8 +55,8 @@ describe('Member Access', () => {
       const boid = Boid();
       // Taking a reference that is local to this function
       const boidRef = boid;
-      boidRef.pos;
-    }).toStrictEqual(snip('(*boidRef).pos', d.vec3f, 'this-function'));
+      return boidRef.pos;
+    }).toStrictEqual(['(*boidRef).pos', d.vec3f, 'this-function']);
   });
 
   it('derefs access to storage with proper address space', ({ root }) => {
@@ -68,14 +67,14 @@ describe('Member Access', () => {
       'use gpu';
       // Taking a reference to a storage variable
       const boidRef = boidReadonly.$;
-      boidRef.pos;
-    }).toStrictEqual(snip('(*boidRef).pos', d.vec3f, 'readonly'));
+      return boidRef.pos;
+    }).toStrictEqual(['(*boidRef).pos', d.vec3f, 'readonly']);
 
     expectSnippetOf(() => {
       'use gpu';
       // Taking a reference to a storage variable
       const boidRef = boidMutable.$;
-      boidRef.pos;
-    }).toStrictEqual(snip('(*boidRef).pos', d.vec3f, 'mutable'));
+      return boidRef.pos;
+    }).toStrictEqual(['(*boidRef).pos', d.vec3f, 'mutable']);
   });
 });

packages/typegpu/tests/tgsl/wgslGenerator.test.ts

@@ -16,7 +16,6 @@ import { CodegenState } from '../../src/types.ts';
 import { it } from 'typegpu-testing-utility';
 import { ArrayExpression } from '../../src/tgsl/generationHelpers.ts';
 import { extractSnippetFromFn } from '../utils/parseResolved.ts';
-import { UnknownData } from '../../src/tgsl/shaderGenerator_members.ts';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -1086,7 +1085,7 @@ describe('wgslGenerator', () => {
   it('creates intermediate representation for array expression', () => {
     const testFn = () => {
       'use gpu';
-      [d.u32(1), 8, 8, 2];
+      return [d.u32(1), 8, 8, 2];
     };
 
     const snippet = extractSnippetFromFn(testFn);