Skip to content

Commit 6d2959c

Browse files
iwoplazaiwoplaza
committed
chore: Use public API in test utilities (extractSnippet, ...)
1 parent f4a5f43 commit 6d2959c

6 files changed

Lines changed: 146 additions & 144 deletions

File tree

‎packages/typegpu/src/tgsl/shaderGenerator_members.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,5 @@ export { UnknownData } from '../data/dataTypes.ts';
22

33
// types
44
export type { ResolutionCtx } from '../types.ts';
5+
export type { Snippet } from '../data/snippet.ts';
6+
export type { Origin } from '../data/snippet.ts';

‎packages/typegpu/src/tgsl/wgslGenerator.ts‎

Lines changed: 70 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -910,6 +910,75 @@ ${this.ctx.pre}}`;
910910
return snip(stitch`${this.ctx.resolve(schema).value}(${args})`, schema, 'runtime');
911911
}
912912

913+
public _return(statement: tinyest.Return): string {
914+
const returnNode = statement[1];
915+
916+
if (returnNode !== undefined) {
917+
const expectedReturnType = this.ctx.topFunctionReturnType;
918+
let returnSnippet = expectedReturnType
919+
? this._typedExpression(returnNode, expectedReturnType)
920+
: this._expression(returnNode);
921+
922+
if (returnSnippet.value instanceof RefOperator) {
923+
throw new WgslTypeError(
924+
stitch`Cannot return references, returning '${returnSnippet.value.snippet}'`,
925+
);
926+
}
927+
928+
// Arguments cannot be returned from functions without copying. A simple example why is:
929+
// const identity = (x) => {
930+
// 'use gpu';
931+
// return x;
932+
// };
933+
//
934+
// const foo = (arg: d.v3f) => {
935+
// 'use gpu';
936+
// const marg = identity(arg);
937+
// marg.x = 1; // 'marg's origin would be 'runtime', so we wouldn't be able to track this misuse.
938+
// };
939+
if (
940+
returnSnippet.origin === 'argument' &&
941+
!wgsl.isNaturallyEphemeral(returnSnippet.dataType) &&
942+
// Only restricting this use in non-entry functions, as the function
943+
// is giving up ownership of all references anyway.
944+
this.ctx.topFunctionScope?.functionType === 'normal'
945+
) {
946+
throw new WgslTypeError(
947+
stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`,
948+
);
949+
}
950+
951+
if (
952+
!expectedReturnType &&
953+
!isEphemeralSnippet(returnSnippet) &&
954+
returnSnippet.origin !== 'this-function'
955+
) {
956+
const str = this.ctx.resolve(returnSnippet.value, returnSnippet.dataType).value;
957+
const typeStr = this.ctx.resolve(unptr(returnSnippet.dataType)).value;
958+
throw new WgslTypeError(
959+
`'return ${str};' is invalid, cannot return references.
960+
-----
961+
Try 'return ${typeStr}(${str});' instead.
962+
-----`,
963+
);
964+
}
965+
966+
returnSnippet = tryConvertSnippet(
967+
this.ctx,
968+
returnSnippet,
969+
unptr(returnSnippet.dataType) as wgsl.AnyWgslData,
970+
false,
971+
);
972+
973+
invariant(returnSnippet.dataType !== UnknownData, 'Return type should be known');
974+
975+
this.ctx.reportReturnType(returnSnippet.dataType);
976+
return stitch`${this.ctx.pre}return ${returnSnippet};`;
977+
}
978+
979+
return `${this.ctx.pre}return;`;
980+
}
981+
913982
public _statement(statement: tinyest.Statement): string {
914983
if (typeof statement === 'string') {
915984
const id = this._identifier(statement);
@@ -923,72 +992,7 @@ ${this.ctx.pre}}`;
923992
}
924993

925994
if (statement[0] === NODE.return) {
926-
const returnNode = statement[1];
927-
928-
if (returnNode !== undefined) {
929-
const expectedReturnType = this.ctx.topFunctionReturnType;
930-
let returnSnippet = expectedReturnType
931-
? this._typedExpression(returnNode, expectedReturnType)
932-
: this._expression(returnNode);
933-
934-
if (returnSnippet.value instanceof RefOperator) {
935-
throw new WgslTypeError(
936-
stitch`Cannot return references, returning '${returnSnippet.value.snippet}'`,
937-
);
938-
}
939-
940-
// Arguments cannot be returned from functions without copying. A simple example why is:
941-
// const identity = (x) => {
942-
// 'use gpu';
943-
// return x;
944-
// };
945-
//
946-
// const foo = (arg: d.v3f) => {
947-
// 'use gpu';
948-
// const marg = identity(arg);
949-
// marg.x = 1; // 'marg's origin would be 'runtime', so we wouldn't be able to track this misuse.
950-
// };
951-
if (
952-
returnSnippet.origin === 'argument' &&
953-
!wgsl.isNaturallyEphemeral(returnSnippet.dataType) &&
954-
// Only restricting this use in non-entry functions, as the function
955-
// is giving up ownership of all references anyway.
956-
this.ctx.topFunctionScope?.functionType === 'normal'
957-
) {
958-
throw new WgslTypeError(
959-
stitch`Cannot return references to arguments, returning '${returnSnippet}'. Copy the argument before returning it.`,
960-
);
961-
}
962-
963-
if (
964-
!expectedReturnType &&
965-
!isEphemeralSnippet(returnSnippet) &&
966-
returnSnippet.origin !== 'this-function'
967-
) {
968-
const str = this.ctx.resolve(returnSnippet.value, returnSnippet.dataType).value;
969-
const typeStr = this.ctx.resolve(unptr(returnSnippet.dataType)).value;
970-
throw new WgslTypeError(
971-
`'return ${str};' is invalid, cannot return references.
972-
-----
973-
Try 'return ${typeStr}(${str});' instead.
974-
-----`,
975-
);
976-
}
977-
978-
returnSnippet = tryConvertSnippet(
979-
this.ctx,
980-
returnSnippet,
981-
unptr(returnSnippet.dataType) as wgsl.AnyWgslData,
982-
false,
983-
);
984-
985-
invariant(returnSnippet.dataType !== UnknownData, 'Return type should be known');
986-
987-
this.ctx.reportReturnType(returnSnippet.dataType);
988-
return stitch`${this.ctx.pre}return ${returnSnippet};`;
989-
}
990-
991-
return `${this.ctx.pre}return;`;
995+
return this._return(statement);
992996
}
993997

994998
if (statement[0] === NODE.if) {

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -104,22 +104,22 @@ describe('add', () => {
104104
it('infers types when adding constants', () => {
105105
const int_int = () => {
106106
'use gpu';
107-
1 + 2;
107+
return 1 + 2;
108108
};
109109

110110
const float_float = () => {
111111
'use gpu';
112-
1.1 + 2.3;
112+
return 1.1 + 2.3;
113113
};
114114

115115
const int_float = () => {
116116
'use gpu';
117-
1.1 + 2;
117+
return 1.1 + 2;
118118
};
119119

120120
const float_int = () => {
121121
'use gpu';
122-
1 + 2.3;
122+
return 1 + 2.3;
123123
};
124124

125125
expectDataTypeOf(int_int).toBe(abstractInt);

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

Lines changed: 18 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { describe } from 'vitest';
22
import { it } from 'typegpu-testing-utility';
33
import { expectSnippetOf } from '../utils/parseResolved.ts';
4-
import { snip } from '../../src/data/snippet.ts';
54
import tgpu, { d } from '../../src/index.js';
65

76
describe('Member Access', () => {
@@ -12,41 +11,41 @@ describe('Member Access', () => {
1211
it('should access member properties of literals', () => {
1312
expectSnippetOf(() => {
1413
'use gpu';
15-
Boid().pos;
16-
}).toStrictEqual(snip('Boid().pos', d.vec3f, 'runtime'));
14+
return Boid().pos;
15+
}).toStrictEqual(['Boid().pos', d.vec3f, 'runtime']);
1716

1817
expectSnippetOf(() => {
1918
'use gpu';
20-
Boid().pos.xyz;
21-
}).toStrictEqual(snip('Boid().pos.xyz', d.vec3f, 'runtime'));
19+
return Boid().pos.xyz;
20+
}).toStrictEqual(['Boid().pos.xyz', d.vec3f, 'runtime']);
2221
});
2322

2423
it('should access member properties of externals', () => {
2524
const boid = Boid({ pos: d.vec3f(1, 2, 3) });
2625

2726
expectSnippetOf(() => {
2827
'use gpu';
29-
boid.pos;
30-
}).toStrictEqual(snip(d.vec3f(1, 2, 3), d.vec3f, 'constant'));
28+
return boid.pos;
29+
}).toStrictEqual([d.vec3f(1, 2, 3), d.vec3f, 'constant']);
3130

3231
expectSnippetOf(() => {
3332
'use gpu';
34-
boid.pos.zyx;
35-
}).toStrictEqual(snip(d.vec3f(3, 2, 1), d.vec3f, 'constant'));
33+
return boid.pos.zyx;
34+
}).toStrictEqual([d.vec3f(3, 2, 1), d.vec3f, 'constant']);
3635
});
3736

3837
it('should access member properties of variables', () => {
3938
const boidVar = tgpu.privateVar(Boid);
4039

4140
expectSnippetOf(() => {
4241
'use gpu';
43-
boidVar.$.pos;
44-
}).toStrictEqual(snip('boidVar.pos', d.vec3f, 'private'));
42+
return boidVar.$.pos;
43+
}).toStrictEqual(['boidVar.pos', d.vec3f, 'private']);
4544

4645
expectSnippetOf(() => {
4746
'use gpu';
48-
boidVar.$.pos.xyz;
49-
}).toStrictEqual(snip('boidVar.pos.xyz', d.vec3f, 'runtime')); // < swizzles are new objects
47+
return boidVar.$.pos.xyz;
48+
}).toStrictEqual(['boidVar.pos.xyz', d.vec3f, 'runtime']); // < swizzles are new objects
5049
});
5150

5251
it('derefs access to local variables with proper address space', () => {
@@ -56,8 +55,8 @@ describe('Member Access', () => {
5655
const boid = Boid();
5756
// Taking a reference that is local to this function
5857
const boidRef = boid;
59-
boidRef.pos;
60-
}).toStrictEqual(snip('(*boidRef).pos', d.vec3f, 'this-function'));
58+
return boidRef.pos;
59+
}).toStrictEqual(['(*boidRef).pos', d.vec3f, 'this-function']);
6160
});
6261

6362
it('derefs access to storage with proper address space', ({ root }) => {
@@ -68,14 +67,14 @@ describe('Member Access', () => {
6867
'use gpu';
6968
// Taking a reference to a storage variable
7069
const boidRef = boidReadonly.$;
71-
boidRef.pos;
72-
}).toStrictEqual(snip('(*boidRef).pos', d.vec3f, 'readonly'));
70+
return boidRef.pos;
71+
}).toStrictEqual(['(*boidRef).pos', d.vec3f, 'readonly']);
7372

7473
expectSnippetOf(() => {
7574
'use gpu';
7675
// Taking a reference to a storage variable
7776
const boidRef = boidMutable.$;
78-
boidRef.pos;
79-
}).toStrictEqual(snip('(*boidRef).pos', d.vec3f, 'mutable'));
77+
return boidRef.pos;
78+
}).toStrictEqual(['(*boidRef).pos', d.vec3f, 'mutable']);
8079
});
8180
});

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import { CodegenState } from '../../src/types.ts';
1616
import { it } from 'typegpu-testing-utility';
1717
import { ArrayExpression } from '../../src/tgsl/generationHelpers.ts';
1818
import { extractSnippetFromFn } from '../utils/parseResolved.ts';
19-
import { UnknownData } from '../../src/tgsl/shaderGenerator_members.ts';
2019

2120
const { NodeTypeCatalog: NODE } = tinyest;
2221

@@ -1086,7 +1085,7 @@ describe('wgslGenerator', () => {
10861085
it('creates intermediate representation for array expression', () => {
10871086
const testFn = () => {
10881087
'use gpu';
1089-
[d.u32(1), 8, 8, 2];
1088+
return [d.u32(1), 8, 8, 2];
10901089
};
10911090

10921091
const snippet = extractSnippetFromFn(testFn);

0 commit comments

Comments
 (0)