Skip to content

Commit 1baffd1

Browse files
cieplypolarcieplypolar
authored
impr: block scopes in the nameRegistry (#2177)
1 parent d5ebf40 commit 1baffd1

5 files changed

Lines changed: 385 additions & 141 deletions

File tree

‎packages/typegpu/src/nameRegistry.ts‎

Lines changed: 64 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import { invariant } from './errors.ts';
2+
13
const bannedTokens = new Set([
24
// keywords
35
'alias',
@@ -385,6 +387,8 @@ export interface NameRegistry {
385387

386388
pushFunctionScope(): void;
387389
popFunctionScope(): void;
390+
pushBlockScope(): void;
391+
popBlockScope(): void;
388392
}
389393

390394
function sanitizePrimer(primer: string | undefined) {
@@ -429,26 +433,35 @@ export function isValidProp(ident: string): boolean {
429433
const prefix = ident.split('_')[0] as string;
430434
return !bannedTokens.has(prefix);
431435
}
436+
type FunctionScopeLayer = {
437+
type: 'functionScope';
438+
};
439+
440+
type BlockScopeLayer = {
441+
type: 'blockScope';
442+
usedBlockScopeNames: Set<string>;
443+
};
444+
445+
type ScopeLayer = FunctionScopeLayer | BlockScopeLayer;
432446

433447
abstract class NameRegistryImpl implements NameRegistry {
434448
abstract getUniqueVariant(base: string): string;
435449

436450
readonly #usedNames: Set<string>;
437-
readonly #usedFunctionScopeNamesStack: Set<string>[];
451+
readonly #scopeStack: ScopeLayer[];
438452

439453
constructor() {
440454
this.#usedNames = new Set<string>([
441455
...bannedTokens,
442456
...builtins,
443457
]);
444-
this.#usedFunctionScopeNamesStack = [];
458+
this.#scopeStack = [];
445459
}
446460

447-
get usedFunctionScopeNames(): Set<string> | undefined {
448-
return this
449-
.#usedFunctionScopeNamesStack[
450-
this.#usedFunctionScopeNamesStack.length - 1
451-
];
461+
get #usedBlockScopeNames(): Set<string> | undefined {
462+
return (this.#scopeStack[this.#scopeStack.length - 1] as
463+
| BlockScopeLayer
464+
| undefined)?.usedBlockScopeNames;
452465
}
453466

454467
makeUnique(primer: string | undefined, global: boolean): string {
@@ -458,31 +471,68 @@ abstract class NameRegistryImpl implements NameRegistry {
458471
if (global) {
459472
this.#usedNames.add(name);
460473
} else {
461-
this.usedFunctionScopeNames?.add(name);
474+
this.#usedBlockScopeNames?.add(name);
462475
}
463476

464477
return name;
465478
}
466479

480+
#isUsedInBlocksBefore(name: string): boolean {
481+
const functionScopeIndex = this.#scopeStack.findLastIndex((scope) =>
482+
scope.type === 'functionScope'
483+
);
484+
return this.#scopeStack.slice(functionScopeIndex + 1).some((scope) =>
485+
(scope as BlockScopeLayer).usedBlockScopeNames.has(name)
486+
);
487+
}
488+
467489
makeValid(primer: string): string {
468-
if (isValidIdentifier(primer) && !this.#usedNames.has(primer)) {
469-
this.usedFunctionScopeNames?.add(primer);
490+
if (
491+
isValidIdentifier(primer) && !this.#usedNames.has(primer) &&
492+
!this.#isUsedInBlocksBefore(primer)
493+
) {
494+
this.#usedBlockScopeNames?.add(primer);
470495
return primer;
471496
}
472497
return this.makeUnique(primer, false);
473498
}
474499

475500
isUsed(name: string): boolean {
476-
return this.#usedNames.has(name) ||
477-
!!this.usedFunctionScopeNames?.has(name);
501+
return this.#usedNames.has(name) || this.#isUsedInBlocksBefore(name);
478502
}
479503

480504
pushFunctionScope(): void {
481-
this.#usedFunctionScopeNamesStack.push(new Set<string>());
505+
this.#scopeStack.push({ type: 'functionScope' });
506+
this.#scopeStack.push({
507+
type: 'blockScope',
508+
usedBlockScopeNames: new Set(),
509+
});
482510
}
483511

484512
popFunctionScope(): void {
485-
this.#usedFunctionScopeNamesStack.pop();
513+
const functionScopeIndex = this.#scopeStack.findLastIndex((scope) =>
514+
scope.type === 'functionScope'
515+
);
516+
517+
if (functionScopeIndex === -1) {
518+
throw new Error('Tried to pop function scope when no scope was present.');
519+
}
520+
521+
this.#scopeStack.splice(functionScopeIndex);
522+
}
523+
524+
pushBlockScope(): void {
525+
this.#scopeStack.push({
526+
type: 'blockScope',
527+
usedBlockScopeNames: new Set(),
528+
});
529+
}
530+
popBlockScope(): void {
531+
invariant(
532+
this.#scopeStack[this.#scopeStack.length - 1]?.type === 'blockScope',
533+
'Tried to pop block scope, but it is not present',
534+
);
535+
this.#scopeStack.pop();
486536
}
487537
}
488538

‎packages/typegpu/src/resolutionCtx.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,10 +420,12 @@ export class ResolutionCtxImpl implements ResolutionCtx {
420420
}
421421

422422
pushBlockScope() {
423+
this.#namespaceInternal.nameRegistry.pushBlockScope();
423424
this._itemStateStack.pushBlockScope();
424425
}
425426

426427
popBlockScope() {
428+
this.#namespaceInternal.nameRegistry.popBlockScope();
427429
this._itemStateStack.pop('blockScope');
428430
}
429431

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

Lines changed: 113 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -1142,18 +1142,23 @@ ${this.ctx.pre}else ${alternate}`;
11421142
if (statement[0] === NODE.for) {
11431143
const [_, init, condition, update, body] = statement;
11441144

1145-
const [initStatement, conditionExpr, updateStatement] = this.ctx
1146-
.withResetIndentLevel(() => [
1147-
init ? this.statement(init) : undefined,
1148-
condition ? this.typedExpression(condition, bool) : undefined,
1149-
update ? this.statement(update) : undefined,
1150-
]);
1151-
1152-
const initStr = initStatement ? initStatement.slice(0, -1) : '';
1153-
const updateStr = updateStatement ? updateStatement.slice(0, -1) : '';
1154-
1155-
const bodyStr = this.block(blockifySingleStatement(body));
1156-
return stitch`${this.ctx.pre}for (${initStr}; ${conditionExpr}; ${updateStr}) ${bodyStr}`;
1145+
try {
1146+
this.ctx.pushBlockScope();
1147+
const [initStatement, conditionExpr, updateStatement] = this.ctx
1148+
.withResetIndentLevel(() => [
1149+
init ? this.statement(init) : undefined,
1150+
condition ? this.typedExpression(condition, bool) : undefined,
1151+
update ? this.statement(update) : undefined,
1152+
]);
1153+
1154+
const initStr = initStatement ? initStatement.slice(0, -1) : '';
1155+
const updateStr = updateStatement ? updateStatement.slice(0, -1) : '';
1156+
1157+
const bodyStr = this.block(blockifySingleStatement(body));
1158+
return stitch`${this.ctx.pre}for (${initStr}; ${conditionExpr}; ${updateStr}) ${bodyStr}`;
1159+
} finally {
1160+
this.ctx.popBlockScope();
1161+
}
11571162
}
11581163

11591164
if (statement[0] === NODE.while) {
@@ -1174,120 +1179,122 @@ ${this.ctx.pre}else ${alternate}`;
11741179
'`for ... of ...` loops only support iterables stored in variables',
11751180
);
11761181
}
1182+
try {
1183+
this.ctx.pushBlockScope();
11771184

1178-
// Our index name will be some element from infinite sequence (i, ii, iii, ...).
1179-
// If user defines `i` and `ii` before `for ... of ...` loop, then our index name will be `iii`.
1180-
// If user defines `i` inside `for ... of ...` then it will be scoped to a new block,
1181-
// so we can safely use `i`.
1182-
let index = 'i'; // it will be valid name, no need to call this.ctx.makeNameValid
1183-
while (this.ctx.getById(index) !== null) {
1184-
index += 'i';
1185-
}
1185+
const index = this.ctx.makeNameValid('i');
11861186

1187-
const elementSnippet = accessIndex(
1188-
iterableSnippet,
1189-
snip(index, u32, 'runtime'),
1190-
);
1191-
if (!elementSnippet) {
1192-
throw new WgslTypeError(
1193-
'`for ... of ...` loops only support array or vector iterables',
1187+
const elementSnippet = accessIndex(
1188+
iterableSnippet,
1189+
snip(index, u32, 'runtime'),
11941190
);
1195-
}
1191+
if (!elementSnippet) {
1192+
throw new WgslTypeError(
1193+
'`for ... of ...` loops only support array or vector iterables',
1194+
);
1195+
}
11961196

1197-
const iterableDataType = iterableSnippet.dataType;
1198-
let elementCountSnippet: Snippet;
1199-
let elementType = elementSnippet.dataType;
1197+
const iterableDataType = iterableSnippet.dataType;
1198+
let elementCountSnippet: Snippet;
1199+
let elementType = elementSnippet.dataType;
12001200

1201-
if (elementType === UnknownData) {
1202-
throw new WgslTypeError(
1203-
stitch`The elements in iterable ${iterableSnippet} are of unknown type`,
1204-
);
1205-
}
1201+
if (elementType === UnknownData) {
1202+
throw new WgslTypeError(
1203+
stitch`The elements in iterable ${iterableSnippet} are of unknown type`,
1204+
);
1205+
}
12061206

1207-
if (wgsl.isWgslArray(iterableDataType)) {
1208-
elementCountSnippet = iterableDataType.elementCount > 0
1209-
? snip(
1210-
`${iterableDataType.elementCount}`,
1207+
if (wgsl.isWgslArray(iterableDataType)) {
1208+
elementCountSnippet = iterableDataType.elementCount > 0
1209+
? snip(
1210+
`${iterableDataType.elementCount}`,
1211+
u32,
1212+
'constant',
1213+
)
1214+
: arrayLength[$gpuCallable].call(this.ctx, [iterableSnippet]);
1215+
} else if (wgsl.isVec(iterableDataType)) {
1216+
elementCountSnippet = snip(
1217+
`${Number(iterableDataType.type.match(/\d/))}`,
12111218
u32,
12121219
'constant',
1213-
)
1214-
: arrayLength[$gpuCallable].call(this.ctx, [iterableSnippet]);
1215-
} else if (wgsl.isVec(iterableDataType)) {
1216-
elementCountSnippet = snip(
1217-
`${Number(iterableDataType.type.match(/\d/))}`,
1218-
u32,
1219-
'constant',
1220-
);
1221-
} else {
1222-
throw new WgslTypeError(
1223-
'`for ... of ...` loops only support array or vector iterables',
1224-
);
1225-
}
1220+
);
1221+
} else {
1222+
throw new WgslTypeError(
1223+
'`for ... of ...` loops only support array or vector iterables',
1224+
);
1225+
}
12261226

1227-
if (loopVar[0] !== NODE.const) {
1228-
throw new WgslTypeError(
1229-
'Only `for (const ... of ... )` loops are supported',
1230-
);
1231-
}
1227+
if (loopVar[0] !== NODE.const) {
1228+
throw new WgslTypeError(
1229+
'Only `for (const ... of ... )` loops are supported',
1230+
);
1231+
}
12321232

1233-
// If it's ephemeral, it's a value that cannot change. If it's a reference, we take
1234-
// an implicit pointer to it
1235-
let loopVarKind = 'let';
1236-
const loopVarName = this.ctx.makeNameValid(loopVar[1]);
1233+
// If it's ephemeral, it's a value that cannot change. If it's a reference, we take
1234+
// an implicit pointer to it
1235+
let loopVarKind = 'let';
1236+
const loopVarName = this.ctx.makeNameValid(loopVar[1]);
1237+
1238+
if (!isEphemeralSnippet(elementSnippet)) {
1239+
if (elementSnippet.origin === 'constant-tgpu-const-ref') {
1240+
loopVarKind = 'const';
1241+
} else if (elementSnippet.origin === 'runtime-tgpu-const-ref') {
1242+
loopVarKind = 'let';
1243+
} else {
1244+
loopVarKind = 'let';
1245+
if (!wgsl.isPtr(elementType)) {
1246+
const ptrType = createPtrFromOrigin(
1247+
elementSnippet.origin,
1248+
concretize(
1249+
elementType as wgsl.AnyWgslData,
1250+
) as wgsl.StorableData,
1251+
);
1252+
invariant(
1253+
ptrType !== undefined,
1254+
`Creating pointer type from origin ${elementSnippet.origin}`,
1255+
);
1256+
elementType = ptrType;
1257+
}
12371258

1238-
if (!isEphemeralSnippet(elementSnippet)) {
1239-
if (elementSnippet.origin === 'constant-tgpu-const-ref') {
1240-
loopVarKind = 'const';
1241-
} else if (elementSnippet.origin === 'runtime-tgpu-const-ref') {
1242-
loopVarKind = 'let';
1243-
} else {
1244-
loopVarKind = 'let';
1245-
if (!wgsl.isPtr(elementType)) {
1246-
const ptrType = createPtrFromOrigin(
1247-
elementSnippet.origin,
1248-
concretize(elementType as wgsl.AnyWgslData) as wgsl.StorableData,
1249-
);
1250-
invariant(
1251-
ptrType !== undefined,
1252-
`Creating pointer type from origin ${elementSnippet.origin}`,
1253-
);
1254-
elementType = ptrType;
1259+
elementType = implicitFrom(elementType as wgsl.Ptr);
12551260
}
1256-
1257-
elementType = implicitFrom(elementType as wgsl.Ptr);
12581261
}
1259-
}
12601262

1261-
const loopVarSnippet = snip(
1262-
loopVarName,
1263-
elementType,
1264-
elementSnippet.origin,
1265-
);
1266-
this.ctx.defineVariable(loopVarName, loopVarSnippet);
1263+
const loopVarSnippet = snip(
1264+
loopVarName,
1265+
elementType,
1266+
elementSnippet.origin,
1267+
);
12671268

1268-
const forStr = stitch`${this.ctx.pre}for (var ${index} = 0u; ${index} < ${
1269-
tryConvertSnippet(this.ctx, elementCountSnippet, u32, false)
1270-
}; ${index}++) {`;
1269+
this.ctx.defineVariable(loopVar[1], loopVarSnippet);
12711270

1272-
this.ctx.indent();
1271+
const forStr =
1272+
stitch`${this.ctx.pre}for (var ${index} = 0u; ${index} < ${
1273+
tryConvertSnippet(this.ctx, elementCountSnippet, u32, false)
1274+
}; ${index}++) {`;
12731275

1274-
const loopVarDeclStr =
1275-
stitch`${this.ctx.pre}${loopVarKind} ${loopVarName} = ${
1276-
tryConvertSnippet(
1277-
this.ctx,
1278-
elementSnippet,
1279-
elementType,
1280-
false,
1281-
)
1282-
};`;
1276+
this.ctx.indent();
12831277

1284-
const bodyStr = `${this.ctx.pre}${
1285-
this.block(blockifySingleStatement(body))
1286-
}`;
1278+
const loopVarDeclStr =
1279+
stitch`${this.ctx.pre}${loopVarKind} ${loopVarName} = ${
1280+
tryConvertSnippet(
1281+
this.ctx,
1282+
elementSnippet,
1283+
elementType,
1284+
false,
1285+
)
1286+
};`;
12871287

1288-
this.ctx.dedent();
1288+
const bodyStr = `${this.ctx.pre}${
1289+
this.block(blockifySingleStatement(body))
1290+
}`;
1291+
1292+
this.ctx.dedent();
12891293

1290-
return stitch`${forStr}\n${loopVarDeclStr}\n${bodyStr}\n${this.ctx.pre}}`;
1294+
return stitch`${forStr}\n${loopVarDeclStr}\n${bodyStr}\n${this.ctx.pre}}`;
1295+
} finally {
1296+
this.ctx.popBlockScope();
1297+
}
12911298
}
12921299

12931300
if (statement[0] === NODE.continue) {

0 commit comments

Comments
 (0)