Skip to content

Commit 40573f1

Browse files
iwoplazaiwoplaza
authored
feat: Inspect vector type in shader function (#1895)
1 parent 29da05f commit 40573f1

8 files changed

Lines changed: 197 additions & 85 deletions

File tree

‎packages/tinyest-for-wgsl/src/parsers.ts‎

Lines changed: 18 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -24,73 +24,6 @@ function isDeclared(ctx: Context, name: string) {
2424
return ctx.stack.some((scope) => scope.declaredNames.includes(name));
2525
}
2626

27-
const BINARY_OP_MAP = {
28-
'==': '==',
29-
'!=': '!=',
30-
'===': '==',
31-
'!==': '!=',
32-
'<': '<',
33-
'<=': '<=',
34-
'>': '>',
35-
'>=': '>=',
36-
'<<': '<<',
37-
'>>': '>>',
38-
get '>>>'(): never {
39-
throw new Error('The `>>>` operator is unsupported in TGSL.');
40-
},
41-
'+': '+',
42-
'-': '-',
43-
'*': '*',
44-
'/': '/',
45-
'%': '%',
46-
'|': '|',
47-
'^': '^',
48-
'&': '&',
49-
get in(): never {
50-
throw new Error('The `in` operator is unsupported in TGSL.');
51-
},
52-
get instanceof(): never {
53-
throw new Error('The `instanceof` operator is unsupported in TGSL.');
54-
},
55-
'**': '**',
56-
get '|>'(): never {
57-
throw new Error('The `|>` operator is unsupported in TGSL.');
58-
},
59-
} as const;
60-
61-
const LOGICAL_OP_MAP = {
62-
'||': '||',
63-
'&&': '&&',
64-
get '??'(): never {
65-
throw new Error('The `??` operator is unsupported in TGSL.');
66-
},
67-
} as const;
68-
69-
const ASSIGNMENT_OP_MAP = {
70-
'=': '=',
71-
'+=': '+=',
72-
'-=': '-=',
73-
'*=': '*=',
74-
'/=': '/=',
75-
'%=': '%=',
76-
'<<=': '<<=',
77-
'>>=': '>>=',
78-
get '>>>='(): never {
79-
throw new Error('The `>>>=` operator is unsupported in TGSL.');
80-
},
81-
'|=': '|=',
82-
'^=': '^=',
83-
'&=': '&=',
84-
get '**='(): never {
85-
throw new Error('The `**=` operator is unsupported in TGSL.');
86-
},
87-
'||=': '||=',
88-
'&&=': '&&=',
89-
get '??='(): never {
90-
throw new Error('The `??=` operator is unsupported in TGSL.');
91-
},
92-
} as const;
93-
9427
const Transpilers: Partial<
9528
{
9629
[Type in JsNode['type']]: (
@@ -149,24 +82,36 @@ const Transpilers: Partial<
14982
},
15083

15184
BinaryExpression(ctx, node) {
152-
const wgslOp = BINARY_OP_MAP[node.operator];
15385
const left = transpile(ctx, node.left) as tinyest.Expression;
15486
const right = transpile(ctx, node.right) as tinyest.Expression;
155-
return [NODE.binaryExpr, left, wgslOp, right];
87+
return [
88+
NODE.binaryExpr,
89+
left,
90+
node.operator as tinyest.BinaryOperator,
91+
right,
92+
];
15693
},
15794

15895
LogicalExpression(ctx, node) {
159-
const wgslOp = LOGICAL_OP_MAP[node.operator];
16096
const left = transpile(ctx, node.left) as tinyest.Expression;
16197
const right = transpile(ctx, node.right) as tinyest.Expression;
162-
return [NODE.logicalExpr, left, wgslOp, right];
98+
return [
99+
NODE.logicalExpr,
100+
left,
101+
node.operator as tinyest.LogicalOperator,
102+
right,
103+
];
163104
},
164105

165106
AssignmentExpression(ctx, node) {
166-
const wgslOp = ASSIGNMENT_OP_MAP[node.operator as acorn.AssignmentOperator];
167107
const left = transpile(ctx, node.left) as tinyest.Expression;
168108
const right = transpile(ctx, node.right) as tinyest.Expression;
169-
return [NODE.assignmentExpr, left, wgslOp, right];
109+
return [
110+
NODE.assignmentExpr,
111+
left,
112+
node.operator as tinyest.AssignmentOperator,
113+
right,
114+
];
170115
},
171116

172117
UnaryExpression(ctx, node) {

‎packages/tinyest/src/nodes.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,12 +118,15 @@ export type Statement =
118118
export type BinaryOperator =
119119
| '=='
120120
| '!='
121+
| '==='
122+
| '!=='
121123
| '<'
122124
| '<='
123125
| '>'
124126
| '>='
125127
| '<<'
126128
| '>>'
129+
| '>>>'
127130
| '+'
128131
| '-'
129132
| '*'
@@ -132,6 +135,8 @@ export type BinaryOperator =
132135
| '|'
133136
| '^'
134137
| '&'
138+
| 'in'
139+
| 'instanceof'
135140
| '**';
136141

137142
export type BinaryExpression = readonly [
@@ -155,7 +160,9 @@ export type AssignmentOperator =
155160
| '&='
156161
| '**='
157162
| '||='
158-
| '&&=';
163+
| '&&='
164+
| '>>>='
165+
| '??=';
159166

160167
export type AssignmentExpression = readonly [
161168
type: NodeTypeCatalog['assignmentExpr'],
@@ -164,7 +171,7 @@ export type AssignmentExpression = readonly [
164171
rhs: Expression,
165172
];
166173

167-
export type LogicalOperator = '&&' | '||';
174+
export type LogicalOperator = '&&' | '||' | '??';
168175

169176
export type LogicalExpression = readonly [
170177
type: NodeTypeCatalog['logicalExpr'],

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,13 @@ export function accessProp(
177177
return accessProp(derefed, propName);
178178
}
179179

180+
if (isVec(target.dataType)) {
181+
// Example: d.vec3f().kind === 'vec3f'
182+
if (propName === 'kind') {
183+
return snip(target.dataType.type, UnknownData, 'constant');
184+
}
185+
}
186+
180187
const propLength = propName.length;
181188
if (isVec(target.dataType) && propLength >= 1 && propLength <= 4) {
182189
const swizzleTypeChar = target.dataType.type.includes('bool')

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

Lines changed: 73 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ import { safeStringify } from '../shared/stringify.ts';
2626
import { $internal } from '../shared/symbols.ts';
2727
import { pow } from '../std/numeric.ts';
2828
import { add, div, mul, neg, sub } from '../std/operators.ts';
29-
import type { FnArgsConversionHint } from '../types.ts';
29+
import { type FnArgsConversionHint, isKnownAtComptime } from '../types.ts';
3030
import {
3131
convertStructValues,
3232
convertToCommonType,
@@ -50,6 +50,8 @@ const { NodeTypeCatalog: NODE } = tinyest;
5050
const parenthesizedOps = [
5151
'==',
5252
'!=',
53+
'===',
54+
'!==',
5355
'<',
5456
'<=',
5557
'>',
@@ -68,7 +70,64 @@ const parenthesizedOps = [
6870
'||',
6971
];
7072

71-
const binaryLogicalOps = ['&&', '||', '==', '!=', '<', '<=', '>', '>='];
73+
const binaryLogicalOps = [
74+
'&&',
75+
'||',
76+
'==',
77+
'!=',
78+
'===',
79+
'!==',
80+
'<',
81+
'<=',
82+
'>',
83+
'>=',
84+
];
85+
86+
const OP_MAP = {
87+
//
88+
// binary
89+
//
90+
'===': '==',
91+
'!==': '!=',
92+
get '>>>'(): never {
93+
throw new Error('The `>>>` operator is unsupported in TypeGPU functions.');
94+
},
95+
get in(): never {
96+
throw new Error('The `in` operator is unsupported in TypeGPU functions.');
97+
},
98+
get instanceof(): never {
99+
throw new Error(
100+
'The `instanceof` operator is unsupported in TypeGPU functions.',
101+
);
102+
},
103+
get '|>'(): never {
104+
throw new Error('The `|>` operator is unsupported in TypeGPU functions.');
105+
},
106+
//
107+
// logical
108+
//
109+
get '??'(): never {
110+
throw new Error('The `??` operator is unsupported in TypeGPU functions.');
111+
},
112+
//
113+
// assignment
114+
//
115+
get '>>>='(): never {
116+
throw new Error('The `>>>=` operator is unsupported in TypeGPU functions.');
117+
},
118+
get '**='(): never {
119+
throw new Error('The `**=` operator is unsupported in TypeGPU functions.');
120+
},
121+
get '??='(): never {
122+
throw new Error('The `??=` operator is unsupported in TypeGPU functions.');
123+
},
124+
get '&&='(): never {
125+
throw new Error('The `&&=` operator is unsupported in TypeGPU functions.');
126+
},
127+
get '||='(): never {
128+
throw new Error('The `||=` operator is unsupported in TypeGPU functions.');
129+
},
130+
} as Record<string, string>;
72131

73132
type Operator =
74133
| tinyest.BinaryOperator
@@ -261,6 +320,16 @@ ${this.ctx.pre}}`;
261320
);
262321
}
263322

323+
if (op === '==') {
324+
throw new Error('Please use the === operator instead of ==');
325+
}
326+
327+
if (
328+
op === '===' && isKnownAtComptime(lhsExpr) && isKnownAtComptime(rhsExpr)
329+
) {
330+
return snip(lhsExpr.value === rhsExpr.value, bool, 'constant');
331+
}
332+
264333
if (lhsExpr.dataType.type === 'unknown') {
265334
throw new WgslTypeError(`Left-hand side of '${op}' is of unknown type`);
266335
}
@@ -328,8 +397,8 @@ ${this.ctx.pre}}`;
328397

329398
return snip(
330399
parenthesizedOps.includes(op)
331-
? `(${lhsStr} ${op} ${rhsStr})`
332-
: `${lhsStr} ${op} ${rhsStr}`,
400+
? `(${lhsStr} ${OP_MAP[op] ?? op} ${rhsStr})`
401+
: `${lhsStr} ${OP_MAP[op] ?? op} ${rhsStr}`,
333402
type,
334403
// Result of an operation, so not a reference to anything
335404
/* origin */ 'runtime',

‎packages/typegpu/tests/constant.test.ts‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,4 +121,49 @@ describe('tgpu.const', () => {
121121
}"
122122
`);
123123
});
124+
125+
it('cannot be passed directly to shellless functions', () => {
126+
const fn1 = (v: d.v3f) => {
127+
'use gpu';
128+
return v.x * v.y * v.z;
129+
};
130+
131+
const foo = tgpu.const(d.vec3f, d.vec3f(1, 2, 3));
132+
const fn2 = () => {
133+
'use gpu';
134+
return fn1(foo.$);
135+
};
136+
137+
expect(() => tgpu.resolve([fn2])).toThrowErrorMatchingInlineSnapshot(`
138+
[Error: Resolution of the following tree failed:
139+
- <root>
140+
- fn*:fn2
141+
- fn*:fn2(): Cannot pass constant references as function arguments. Explicitly copy them by wrapping them in a schema: 'vec3f(...)']
142+
`);
143+
});
144+
145+
it('cannot be mutated', () => {
146+
const boid = tgpu.const(Boid, {
147+
pos: d.vec3f(1, 2, 3),
148+
vel: d.vec3u(4, 5, 6),
149+
});
150+
151+
const fn = () => {
152+
'use gpu';
153+
// @ts-expect-error: Cannot assign to read-only property
154+
boid.$.pos = d.vec3f(0, 0, 0);
155+
};
156+
157+
expect(() => tgpu.resolve([fn])).toThrowErrorMatchingInlineSnapshot(`
158+
[Error: Resolution of the following tree failed:
159+
- <root>
160+
- fn*:fn
161+
- fn*:fn(): 'boid.pos = vec3f()' is invalid, because boid.pos is a constant.]
162+
`);
163+
164+
// Since we freeze the object, we cannot mutate when running the function in JS either
165+
expect(() => fn()).toThrowErrorMatchingInlineSnapshot(
166+
`[TypeError: Cannot assign to read only property 'pos' of object '#<Object>']`,
167+
);
168+
});
124169
});

‎packages/typegpu/tests/examples/individual/tgsl-parsing-test.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,10 @@ describe('tgsl parsing test example', () => {
3636
3737
fn logicalExpressionTests_1() -> bool {
3838
var s = true;
39-
s = (s && (true == true));
40-
s = (s && (true == true));
41-
s = (s && (true == true));
42-
s = (s && (false == false));
39+
s = (s && true);
40+
s = (s && true);
41+
s = (s && true);
42+
s = (s && true);
4343
s = (s && true);
4444
s = (s && !false);
4545
s = (s && true);

0 commit comments

Comments
 (0)