Skip to content

fix: External traversal for raw WGSL implemented functions #2648

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aleksanderkatan wants to merge 8 commits into
base: fix/forbid-multiple-uses
Choose a base branch
from
Draft

fix: External traversal for raw WGSL implemented functions #2648

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a7abca9
Initial rewrite
aleksanderkatan Jun 24, 2026
b0b5472
Add an end of chain test
aleksanderkatan Jun 24, 2026
ca3e9e1
Simplify end of chain tracking
aleksanderkatan Jun 24, 2026
48727bb
Update snapshots
aleksanderkatan Jun 24, 2026
e346c8e
Update declare tests
aleksanderkatan Jun 24, 2026
4974305
Remove warn test
aleksanderkatan Jun 24, 2026
65fcedf
Simplify regex
aleksanderkatan Jun 24, 2026
38a0643
Update packages/typegpu/src/core/resolve/externals.ts
aleksanderkatan Jun 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 28 additions & 28 deletions apps/typegpu-docs/tests/individual-example-tests/log-test.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ describe('console log example', () => {
expect(shaderCodes).toMatchInlineSnapshot(`
"@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -48,8 +50,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -90,6 +90,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -99,8 +101,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -150,6 +150,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -159,8 +161,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -203,6 +203,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -212,8 +214,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn log1serializer() {
Expand Down Expand Up @@ -834,6 +834,8 @@ describe('console log example', () => {
num: u32,
}

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -843,8 +845,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -995,6 +995,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1004,8 +1006,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn log1serializer() {
Expand Down Expand Up @@ -1052,6 +1052,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1061,8 +1063,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -1103,6 +1103,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1112,8 +1114,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -1158,6 +1158,8 @@ describe('console log example', () => {

@group(0) @binding(1) var<uniform> logCountUniform: u32;

var<private> dataBlockIndex: u32;

@group(0) @binding(2) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1167,8 +1169,6 @@ describe('console log example', () => {

@group(0) @binding(3) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -1212,6 +1212,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1221,8 +1223,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -1299,6 +1299,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1308,8 +1310,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn log1serializer() {
Expand Down Expand Up @@ -1453,6 +1453,8 @@ describe('console log example', () => {
return mainVertex_Output(vec4f(positions[vertexIndex], 0f, 1f));
}

var<private> dataBlockIndex: u32;

@group(0) @binding(0) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1462,8 +1464,6 @@ describe('console log example', () => {

@group(0) @binding(1) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -1506,6 +1506,8 @@ describe('console log example', () => {
return mainVertex_Output(vec4f(positions[vertexIndex], 0f, 1f));
}

var<private> dataBlockIndex: u32;

@group(0) @binding(0) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1515,8 +1517,6 @@ describe('console log example', () => {

@group(0) @binding(1) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down Expand Up @@ -1552,6 +1552,8 @@ describe('console log example', () => {

@group(0) @binding(0) var<uniform> sizeUniform: vec3u;

var<private> dataBlockIndex: u32;

@group(0) @binding(1) var<storage, read_write> indexBuffer: atomic<u32>;

struct SerializedLogData {
Expand All @@ -1561,8 +1563,6 @@ describe('console log example', () => {

@group(0) @binding(2) var<storage, read_write> dataBuffer: array<SerializedLogData, 40>;

var<private> dataBlockIndex: u32;

var<private> dataByteIndex: u32;

fn nextByteIndex() -> u32 {
Expand Down
94 changes: 47 additions & 47 deletions packages/typegpu/src/core/resolve/externals.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,18 +66,15 @@ export function addReturnTypeToExternals(
}
}

function identifierRegex(name: string) {
return new RegExp(
`(?<![\\w\\$_.])${name.replaceAll('.', '\\.').replaceAll('$', '\\$')}(?![\\w\\$_])`,
'g',
);
}
const anyIdent = /([$_\p{XID_Start}][$\p{XID_Continue}]*)/u; // WGSL ident, modified to include $
const anyPropChain = new RegExp(`(${anyIdent.source})(\\.${anyIdent.source})*`, 'ug');
const boundedPropChain = new RegExp(`(?<![\\w\\$_.])${anyPropChain.source}(?![\\w\\$_])`, 'ug');

/**
* Replaces all occurrences of external names in WGSL code with their resolved values.
* It adds all necessary definitions to the resolution context.
* @param ctx - The resolution context.
* @param externalMap - The external map.
* @param externalMap - The external map. Assumes that keys don't contain dots.
* @param wgsl - The WGSL code.
*
* @returns The WGSL code with all external names replaced with their resolved values.
Expand All @@ -87,52 +84,55 @@ export function replaceExternalsInWgsl(
externalMap: ExternalMap,
wgsl: string,
): string {
return Object.entries(externalMap).reduce((acc, [externalName, external]) => {
const externalRegex = identifierRegex(externalName);
if (wgsl && externalName !== 'Out' && externalName !== 'in' && !externalRegex.test(wgsl)) {
console.warn(`The external '${externalName}' wasn't used in the resolved template.`);
// continue anyway, we still might need to resolve the external
const keys = Object.keys(externalMap);
if (keys.length === 0) {
return wgsl;
}

// Avoid resolving the same item multiple times during one call.
const cache: Map<unknown, string> = new Map();

return wgsl.replaceAll(boundedPropChain, (match) => {
const chain = match.split('.');

if (!Object.hasOwn(externalMap, chain.at(0) as string)) {
// this prop access does not start with an external
return match;
}

if (isResolvable(external)) {
if (isNamable(external) && getName(external) === undefined) {
setName(external, externalName.split('.').at(-1) as string);
let currentItem: unknown = externalMap;
let name: string | undefined = undefined;
let suffix = '';

for (const [i, elem] of chain.entries()) {
currentItem = (currentItem as ExternalMap)[elem];
name = elem;
if (isResolvable(currentItem)) {
suffix = chain
.slice(i + 1)
.map((s) => `.${s}`)
.join('');
break;
}

if (typeof currentItem !== 'object' || currentItem === null || i === chain.length - 1) {
console.warn(
`During resolution, the external '${chain.slice(0, i + 1).join('.')}' has been omitted. Only TGPU resources, 'use gpu' functions, primitives, and plain JS objects can be used as externals.`,
);
return match;
}
return acc.replaceAll(externalRegex, ctx.resolve(external).value);
}

if (external !== null && typeof external === 'object') {
const foundProperties = [
...wgsl.matchAll(
new RegExp(
`${externalName
.replaceAll('.', '\\.')
.replaceAll('$', '\\$')}\\.(?<prop>.*?)(?![\\w\\$_])`,
'g',
),
),
].map((found) => found[1]);
const uniqueProperties = [...new Set(foundProperties)];

return uniqueProperties.reduce(
(innerAcc: string, prop) =>
prop && prop in external
? replaceExternalsInWgsl(
ctx,
{
[`${externalName}.${prop}`]: external[prop as keyof typeof external],
},
innerAcc,
)
: innerAcc,
acc,
);
if (isNamable(currentItem) && getName(currentItem) === undefined && name !== undefined) {
setName(currentItem, name);
}

console.warn(
`During resolution, the external '${externalName}' has been omitted. Only TGPU resources, 'use gpu' functions, primitives, and plain JS objects can be used as externals.`,
);
let resolved = cache.get(currentItem);
if (resolved === undefined) {
resolved = ctx.resolve(currentItem).value;
cache.set(currentItem, resolved);
}

return acc;
}, wgsl);
return resolved + suffix;
});
}
Loading
Loading