Skip to content

fix: Unique names in raw WGSL code #2499

@iwoplaza

Description

@iwoplaza

We provide ways to integrate raw WGSL code into TypeScript shaders, and the other way around. We don't however properly handle name clashes.

const OFFSET = tgpu.const(d.f32, 10).$name('x');

const hello = tgpu.rawCodeSnippet('buffer[i]', d.f32).$uses({ buffer });
//                                        ^ reference to a local definition
//                                 ^^^^^^ reference to a global definition

// `OFFSET` is resolved before `foo`, therefore by
// the time we start generating `foo`, the identifier
// `x` is already taken, which clashes with the local definition
const foo = tgpu.fn([], d.f32)`() {
  var x = 0f;
  for (var i = 0; i < 10; i++) {
    x += hello;
  }
  return x + OFFSET;
}`.$uses({ OFFSET, hello });

tgpu.resolve([OFFSET, foo]);

Let's keep a new localRenames map

  1. Push a function layer in the ItemStateStack (or an equivalent layer that will allow up to store unique local names)
  2. Collect all identifiers used in the function body, excluding external keys (in this case: x and i)
  3. Filter for those that are already taken (in this case: x)
  4. Generate unique identifiers for those that are taken, save them to the map (x => x_1), find and replace in the WGSL template.
  5. We reserve the non-taken identifiers, so no new global definitions will be generated with those names as long as the function layer is on the stack (in this case just i, as x_1 has already been reserved as part of generating a unique name).
    After this process, the WGSL template will look like this:
() {
  var x_1 = 0f;
  for (var i = 0; i < 10; i++) {
    x_1 += hello;
  }
  return x_1;
}

Now we can resolve all externals, and replace them in the template:

  1. When we resolve tgpu.rawCodeSnippet, we perform basically the same process we did before, where we collect all identifiers used in this now smaller template (excluding external keys), which in this case is just i
  2. We look up if any local identifiers have been renamed (using the localRenames map), and find and replace (in this case, we don't rename anything, as i wasn't renamed in the parent function)
const hello = tgpu.rawCodeSnippet('buffer[i]', d.f32).$uses({ buffer });
//                                        ^ reference to a local definition
//                                 ^^^^^^ reference to a global definition

We can determine whether something is an external definition (external name), or something referencing a local definition (local name) based on if something is passed in $uses(...).

tgpu.resolve with template

We can treat tgpu.resolve with templates the same way we treat WGSL-implemented functions, we just don't need to find and replace anything in the template as we already know the template is the first thing being resolved, so we can just globally reserve every non-external identifier that
is used in the template, and proceed with regular resolving of externals.

const FACTOR = tgpu.const(d.u32, 10).$name('factor');

const shader = tgpu.resolve({
  template: `
fn foo() {
  const factor = 1f;
  const hello = FACTOR;
}
`,
  externals: {
    FACTOR,
  },
});

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workingstabilityTasks that focus on improving stability, that includes tests, refining edge-cases, refactoring.

    Type

    No fields configured for Bug.

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions