Skip to content

Commit 2164494

Browse files
aleksanderkatanaleksanderkatan
committed
work in progress
1 parent 439cc62 commit 2164494

8 files changed

Lines changed: 108 additions & 205 deletions

File tree

‎apps/typegpu-docs/tests/individual-example-tests/function-visualizer.test.ts‎

Lines changed: 1 addition & 184 deletions
Original file line numberDiff line numberDiff line change
@@ -21,189 +21,6 @@ describe('function visualizer example', () => {
2121
device,
2222
);
2323

24-
expect(shaderCodes).toMatchInlineSnapshot(`
25-
"@group(0) @binding(0) var<uniform> sizeUniform: vec3u;
26-
27-
struct Properties {
28-
transformation: mat4x4f,
29-
inverseTransformation: mat4x4f,
30-
interpolationPoints: u32,
31-
lineWidth: f32,
32-
}
33-
34-
@group(0) @binding(1) var<uniform> propertiesUniform: Properties;
35-
36-
fn interpolatedFunction(x: f32) -> f32 {
37-
return x;
38-
}
39-
40-
@group(1) @binding(0) var<storage, read_write> lineVertices: array<vec2f>;
41-
42-
fn computePointsFn(x: u32, _arg_1: u32, _arg_2: u32) {
43-
let properties = (&propertiesUniform);
44-
let start = ((*properties).transformation * vec4f(-1, 0, 0, 1)).x;
45-
let end = ((*properties).transformation * vec4f(1, 0, 0, 1)).x;
46-
let pointX = (start + (((end - start) / (f32((*properties).interpolationPoints) - 1f)) * f32(x)));
47-
let pointY = interpolatedFunction(pointX);
48-
var result = ((*properties).inverseTransformation * vec4f(pointX, pointY, 0f, 1f));
49-
lineVertices[x] = result.xy;
50-
}
51-
52-
@compute @workgroup_size(256, 1, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) {
53-
if (any(id >= sizeUniform)) {
54-
return;
55-
}
56-
computePointsFn(id.x, id.y, id.z);
57-
}
58-
59-
@group(0) @binding(0) var<uniform> sizeUniform: vec3u;
60-
61-
struct Properties {
62-
transformation: mat4x4f,
63-
inverseTransformation: mat4x4f,
64-
interpolationPoints: u32,
65-
lineWidth: f32,
66-
}
67-
68-
@group(0) @binding(1) var<uniform> propertiesUniform: Properties;
69-
70-
fn interpolatedFunction(x: f32) -> f32 {
71-
return cos(x*5)/3-x;
72-
}
73-
74-
@group(1) @binding(0) var<storage, read_write> lineVertices: array<vec2f>;
75-
76-
fn computePointsFn(x: u32, _arg_1: u32, _arg_2: u32) {
77-
let properties = (&propertiesUniform);
78-
let start = ((*properties).transformation * vec4f(-1, 0, 0, 1)).x;
79-
let end = ((*properties).transformation * vec4f(1, 0, 0, 1)).x;
80-
let pointX = (start + (((end - start) / (f32((*properties).interpolationPoints) - 1f)) * f32(x)));
81-
let pointY = interpolatedFunction(pointX);
82-
var result = ((*properties).inverseTransformation * vec4f(pointX, pointY, 0f, 1f));
83-
lineVertices[x] = result.xy;
84-
}
85-
86-
@compute @workgroup_size(256, 1, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) {
87-
if (any(id >= sizeUniform)) {
88-
return;
89-
}
90-
computePointsFn(id.x, id.y, id.z);
91-
}
92-
93-
@group(0) @binding(0) var<uniform> sizeUniform: vec3u;
94-
95-
struct Properties {
96-
transformation: mat4x4f,
97-
inverseTransformation: mat4x4f,
98-
interpolationPoints: u32,
99-
lineWidth: f32,
100-
}
101-
102-
@group(0) @binding(1) var<uniform> propertiesUniform: Properties;
103-
104-
fn interpolatedFunction(x: f32) -> f32 {
105-
return x*sin(log(abs(x)));
106-
}
107-
108-
@group(1) @binding(0) var<storage, read_write> lineVertices: array<vec2f>;
109-
110-
fn computePointsFn(x: u32, _arg_1: u32, _arg_2: u32) {
111-
let properties = (&propertiesUniform);
112-
let start = ((*properties).transformation * vec4f(-1, 0, 0, 1)).x;
113-
let end = ((*properties).transformation * vec4f(1, 0, 0, 1)).x;
114-
let pointX = (start + (((end - start) / (f32((*properties).interpolationPoints) - 1f)) * f32(x)));
115-
let pointY = interpolatedFunction(pointX);
116-
var result = ((*properties).inverseTransformation * vec4f(pointX, pointY, 0f, 1f));
117-
lineVertices[x] = result.xy;
118-
}
119-
120-
@compute @workgroup_size(256, 1, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) {
121-
if (any(id >= sizeUniform)) {
122-
return;
123-
}
124-
computePointsFn(id.x, id.y, id.z);
125-
}
126-
127-
struct Properties {
128-
transformation: mat4x4f,
129-
inverseTransformation: mat4x4f,
130-
interpolationPoints: u32,
131-
lineWidth: f32,
132-
}
133-
134-
@group(0) @binding(0) var<uniform> propertiesUniform: Properties;
135-
136-
struct backgroundVertex_Output {
137-
@builtin(position) pos: vec4f,
138-
}
139-
140-
@vertex fn backgroundVertex(@builtin(vertex_index) vid: u32, @builtin(instance_index) iid: u32) -> backgroundVertex_Output {
141-
let properties = (&propertiesUniform);
142-
var leftBot = ((*properties).transformation * vec4f(-1, -1, 0, 1));
143-
var rightTop = ((*properties).transformation * vec4f(1, 1, 0, 1));
144-
let aspectRatio = ((rightTop.x - leftBot.x) / (rightTop.y - leftBot.y));
145-
var transformedPoints = array<vec2f, 4>(vec2f(leftBot.x, 0f), vec2f(rightTop.x, 0f), vec2f(0f, leftBot.y), vec2f(0f, rightTop.y));
146-
var currentPoint = ((*properties).inverseTransformation * vec4f(transformedPoints[((2u * iid) + u32((f32(vid) / 2f)))].xy, 0f, 1f));
147-
return backgroundVertex_Output(vec4f((currentPoint.x + (((f32(iid) * select(-1f, 1f, ((vid % 2u) == 0u))) * 5e-3f) / aspectRatio)), (currentPoint.y + ((f32((1u - iid)) * select(-1f, 1f, ((vid % 2u) == 0u))) * 5e-3f)), currentPoint.zw));
148-
}
149-
150-
@fragment fn backgroundFragment() -> @location(0) vec4f {
151-
return vec4f(0.8999999761581421, 0.8999999761581421, 0.8999999761581421, 1);
152-
}
153-
154-
struct Properties {
155-
transformation: mat4x4f,
156-
inverseTransformation: mat4x4f,
157-
interpolationPoints: u32,
158-
lineWidth: f32,
159-
}
160-
161-
@group(0) @binding(0) var<uniform> propertiesUniform: Properties;
162-
163-
@group(1) @binding(0) var<storage, read> lineVertices_1: array<vec2f>;
164-
165-
fn orthonormalForLine(p1: vec2f, p2: vec2f) -> vec2f {
166-
var line = (p2 - p1);
167-
var ortho = vec2f(-(line.y), line.x);
168-
return normalize(ortho);
169-
}
170-
171-
fn orthonormalForVertex(index: f32) -> vec2f {
172-
if (((index == 0f) || (index == 255f))) {
173-
return vec2f(0, 1);
174-
}
175-
let lineVertices = (&lineVertices_1);
176-
let previous = (&(*lineVertices)[u32((index - 1f))]);
177-
let current = (&(*lineVertices)[u32(index)]);
178-
let next = (&(*lineVertices)[u32((index + 1f))]);
179-
var n1 = orthonormalForLine((*previous), (*current));
180-
var n2 = orthonormalForLine((*current), (*next));
181-
var avg = ((n1 + n2) / 2f);
182-
return normalize(avg);
183-
}
184-
185-
struct vertex_Output {
186-
@builtin(position) pos: vec4f,
187-
}
188-
189-
@vertex fn vertex(@builtin(vertex_index) vid: u32) -> vertex_Output {
190-
let properties = (&propertiesUniform);
191-
let lineVertices = (&lineVertices_1);
192-
let currentVertex = (f32(vid) / 2f);
193-
var orthonormal = orthonormalForVertex(currentVertex);
194-
var offset = ((orthonormal * (*properties).lineWidth) * select(-1f, 1f, ((vid % 2u) == 0u)));
195-
var leftBot = ((*properties).transformation * vec4f(-1, -1, 0, 1));
196-
var rightTop = ((*properties).transformation * vec4f(1, 1, 0, 1));
197-
let canvasRatio = ((rightTop.x - leftBot.x) / (rightTop.y - leftBot.y));
198-
var adjustedOffset = vec2f((offset.x / canvasRatio), offset.y);
199-
return vertex_Output(vec4f(((*lineVertices)[u32(currentVertex)] + adjustedOffset), 0f, 1f));
200-
}
201-
202-
@group(1) @binding(1) var<uniform> color: vec4f;
203-
204-
@fragment fn fragment() -> @location(0) vec4f {
205-
return color;
206-
}"
207-
`);
24+
expect(shaderCodes).toMatchInlineSnapshot(`""`);
20825
});
20926
});

‎packages/typegpu/src/core/function/fnCore.ts‎

Lines changed: 31 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,19 @@ import { applyExternals, type ExternalMap, replaceExternalsInWgsl } from '../res
1010
import { extractArgs } from './extractArgs.ts';
1111
import type { Implementation, SeparatedEntryArgs } from './fnTypes.ts';
1212

13+
// TODO: deslopify
14+
export function extractDeclaredNames(codeString: string) {
15+
// The Regex pattern
16+
const regex =
17+
/(?:\b(?:var|let|const)\s+(?<varName>[a-zA-Z_$][\w$]*))|(?<=\(|,)\s*(?<argName>[a-zA-Z_$][\w$]*)\s*(?=:|,|\))/g;
18+
19+
// Extract the named groups from all matches
20+
return Array.from(
21+
codeString.matchAll(regex),
22+
(match) => match.groups?.varName || (match.groups?.argName as string),
23+
);
24+
}
25+
1326
export interface FnCore {
1427
applyExternals: (newExternals: ExternalMap) => void;
1528
resolve(
@@ -81,17 +94,26 @@ export function createFnCore(
8194
throw new Error('Explicit return type is required for string implementation');
8295
}
8396

84-
const validArgNames = entryInput
85-
? Object.fromEntries(
86-
entryInput.positionalArgs.map((a) => [a.schemaKey, ctx.makeNameValid(a.schemaKey)]),
87-
)
88-
: undefined;
97+
const usedNames = new Set(extractDeclaredNames(implementation));
98+
const [validArgNames, replacedImpl] = ctx.withReservedNames(new Set(usedNames), () => {
99+
console.log('TOKENS\n', usedNames);
89100

90-
if (validArgNames && Object.keys(validArgNames).length > 0) {
91-
applyExternals(externalMap, { in: validArgNames });
92-
}
101+
const validArgNames = entryInput
102+
? Object.fromEntries(
103+
entryInput.positionalArgs.map((a) => [a.schemaKey, ctx.makeNameValid(a.schemaKey)]),
104+
)
105+
: undefined;
93106

94-
const replacedImpl = replaceExternalsInWgsl(ctx, externalMap, implementation);
107+
if (validArgNames && Object.keys(validArgNames).length > 0) {
108+
applyExternals(externalMap, { in: validArgNames });
109+
}
110+
111+
console.log('STILL TOKENS\n', usedNames);
112+
return [
113+
validArgNames,
114+
replaceExternalsInWgsl(ctx, externalMap, implementation, usedNames),
115+
];
116+
});
95117

96118
let header = '';
97119
let body = '';

‎packages/typegpu/src/core/resolve/externals.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,9 @@ export function replaceExternalsInWgsl(
8282
ctx: ResolutionCtx,
8383
externalMap: ExternalMap,
8484
wgsl: string,
85+
forbiddenNames?: Set<string>,
8586
): string {
87+
console.log('ENTERING R.E. WITH', forbiddenNames);
8688
return Object.entries(externalMap).reduce((acc, [externalName, external]) => {
8789
const externalRegex = identifierRegex(externalName);
8890
if (wgsl && externalName !== 'Out' && externalName !== 'in' && !externalRegex.test(wgsl)) {
@@ -91,7 +93,14 @@ export function replaceExternalsInWgsl(
9193
}
9294

9395
if (isWgsl(external) || isLooseData(external) || hasTinyestMetadata(external)) {
94-
return acc.replaceAll(externalRegex, ctx.resolve(external).value);
96+
const externalName = ctx.resolve(external).value;
97+
if (forbiddenNames?.has(externalName)) {
98+
console.log('FORBIDDEN', forbiddenNames);
99+
throw new Error(
100+
`Name clash on external and variable '${externalName}' and external/argument '${typeof external === 'string' ? external : getName(external)}' that was automatically renamed to '${externalName}'. Please either rename the variable, or give the external a different name using '.$name()'.`,
101+
);
102+
}
103+
return acc.replaceAll(externalRegex, externalName);
95104
}
96105

97106
if (external !== null && typeof external === 'object') {
@@ -116,6 +125,7 @@ export function replaceExternalsInWgsl(
116125
[`${externalName}.${prop}`]: external[prop as keyof typeof external],
117126
},
118127
innerAcc,
128+
forbiddenNames,
119129
)
120130
: innerAcc,
121131
acc,

‎packages/typegpu/src/nameRegistry.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -385,7 +385,7 @@ export interface NameRegistry {
385385
*/
386386
makeValid(primer: string): string;
387387

388-
pushFunctionScope(): void;
388+
pushFunctionScope(initial?: Set<string>): void;
389389
popFunctionScope(): void;
390390
pushBlockScope(): void;
391391
popBlockScope(): void;
@@ -498,11 +498,11 @@ abstract class NameRegistryImpl implements NameRegistry {
498498
return this.#usedNames.has(name) || this.#isUsedInBlocksBefore(name);
499499
}
500500

501-
pushFunctionScope(): void {
501+
pushFunctionScope(initial?: Set<string>): void {
502502
this.#scopeStack.push({ type: 'functionScope' });
503503
this.#scopeStack.push({
504504
type: 'blockScope',
505-
usedBlockScopeNames: new Set(),
505+
usedBlockScopeNames: initial ?? new Set(),
506506
});
507507
}
508508

‎packages/typegpu/src/resolutionCtx.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,13 @@ export class ResolutionCtxImpl implements ResolutionCtx {
390390
return this.#namespaceInternal.nameRegistry.makeValid(name);
391391
}
392392

393+
withReservedNames<T>(names: Set<string>, cb: () => T): T {
394+
this.#namespaceInternal.nameRegistry.pushFunctionScope(names);
395+
const result = cb();
396+
this.#namespaceInternal.nameRegistry.popFunctionScope();
397+
return result;
398+
}
399+
393400
get pre(): string {
394401
return this._indentController.pre;
395402
}

‎packages/typegpu/src/types.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,20 @@ export interface ResolutionCtx {
329329
* @param name the temporary name to assign to the item (if missing, just returns `callback()`)
330330
*/
331331
withRenamed<T>(item: object, name: string | undefined, callback: () => T): T;
332+
/**
333+
* Temporarily reserves the set of names.
334+
* Useful for resolving externals of wgsl-implemented functions,
335+
* since we have to prevent externals receiving names used by the function.
336+
* Example where this would be an issue:
337+
* ```ts
338+
* const jsConst = tgpu.const(d.u32, 1).$name('myConst');
339+
* const fn = tgpu.fn([])`() {
340+
* const myConst = 0;
341+
* const otherConst = extConst;
342+
* }`.$uses({ extConst: jsConst });
343+
* ```
344+
*/
345+
withReservedNames<T>(names: Set<string>, cb: () => T): T;
332346

333347
getUniqueName(resource: object): string;
334348
makeNameValid(name: string): string;

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -143,9 +143,9 @@ describe('autogenerating wgsl headers for tgpu entry functions with raw string W
143143
@builtin(position) outPos: vec4f,
144144
}
145145
146-
@vertex fn mainVertex(@builtin(vertex_index) vi_0: u32, @builtin(instance_index) ii_0: u32) -> mainVertex_Output {
147-
var vi = vi_0;
148-
let ii = ii_0;
146+
@vertex fn mainVertex(@builtin(vertex_index) vi_1: u32, @builtin(instance_index) ii_1: u32) -> mainVertex_Output {
147+
var vi = vi_1;
148+
let ii = ii_1;
149149
150150
return mainVertex_Output(vec4f(vi, ii, 0, 1));
151151
}"
@@ -198,10 +198,10 @@ describe('autogenerating wgsl headers for tgpu entry functions with raw string W
198198
@builtin(position) outPos: vec4f,
199199
}
200200
201-
@vertex fn mainVertex(@builtin(vertex_index) vi_0: u32, @builtin(instance_index) vi_0_0: u32) -> mainVertex_Output {
201+
@vertex fn mainVertex(@builtin(vertex_index) vi_1: u32, @builtin(instance_index) vi_0: u32) -> mainVertex_Output {
202202
var vi = 0;
203-
var a = vi_0;
204-
var b = vi_0_0;
203+
var a = vi_1;
204+
var b = vi_0;
205205
206206
return mainVertex_Output(vec4f(a, b, 0, 1));
207207
}"

0 commit comments

Comments
 (0)