feat: std.getShaderStage

[iwoplaza]

I additionally merged our ShaderStage and TgpuShaderStage into one type, then exported it from typegpu instead of typegpu/~internal.

[github-actions]

pkg.pr.new

packages
Ready to be installed by your favorite package manager ⬇️

https://pkg.pr.new/software-mansion/TypeGPU/eslint-plugin-typegpu@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/tgpu-gen@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/tinyest-for-wgsl@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/typegpu@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/cli@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/color@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/gl@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/noise@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/radiance-cascades@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/react@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/sdf@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/@typegpu/three@b48fecd9808f31ea94cbf2e576e25284124281b3

https://pkg.pr.new/software-mansion/TypeGPU/unplugin-typegpu@b48fecd9808f31ea94cbf2e576e25284124281b3

benchmark
view benchmark

commit
view commit

[Copilot]

Pull request overview

This PR adds a new standard-library probe (std.getShaderStage) to allow shader code to branch on the entry-point stage during resolution, while also consolidating the internal stage type (TgpuShaderStage) into a single exported ShaderStage type.

Changes:

• Added std.getShaderStage() (comptime probe) backed by an internal shaderStageSlot.
• Unified TgpuShaderStage and ShaderStage into a single exported ShaderStage type and updated internal usages.
• Added test coverage and updated docs to describe the new probe.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
packages/typegpu/tests/std/environment.test.ts	Adds coverage for std.getShaderStage() across top-level, resolution, JS execution, and simulate contexts.
packages/typegpu/src/types.ts	Replaces TgpuShaderStage with exported ShaderStage and updates related type unions.
packages/typegpu/src/tgsl/shaderGenerator.ts	Updates function definition typing to use ShaderStage.
packages/typegpu/src/tgpuBindGroupLayout.ts	Updates bind group visibility typing and defaults to use ShaderStage[].
packages/typegpu/src/std/index.ts	Re-exports getShaderStage from std.
packages/typegpu/src/std/environment.ts	Implements getShaderStage() as a comptime probe reading from shaderStageSlot.
packages/typegpu/src/resolutionCtx.ts	Updates function-scope typing to use ShaderStage.
packages/typegpu/src/internal.ts	Updates internal type exports to use ShaderStage.
packages/typegpu/src/indexNamedExports.ts	Exports ShaderStage from the public named exports.
packages/typegpu/src/core/slot/internalSlots.ts	Changes shaderStageSlot default to null and types it as `ShaderStage
packages/typegpu/src/core/function/fnCore.ts	Updates function core typing to use ShaderStage.
packages/typegpu-gl/src/glslGenerator.ts	Updates GLSL generator’s function type tracking to use ShaderStage.
apps/typegpu-docs/src/content/docs/apis/utils.mdx	Documents std.getShaderStage() and adjusts the “resolution environment” section text.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

Resolution Time Benchmark

---
config:
  themeVariables:
    xyChart:
      plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
  title "Random Branching (🔴 PR | 🔵 main | 🟢 release)"
  x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
  y-axis "time (ms)"
  line [0.95, 1.90, 4.04, 6.46, 7.77, 12.07, 23.21, 25.06]
  line [0.96, 1.96, 4.08, 6.21, 7.81, 10.89, 23.08, 24.78]
  line [0.97, 1.92, 4.59, 7.16, 8.24, 11.90, 22.65, 27.99]

Loading

---
config:
  themeVariables:
    xyChart:
      plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
  title "Linear Recursion (🔴 PR | 🔵 main | 🟢 release)"
  x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
  y-axis "time (ms)"
  line [0.31, 0.53, 0.70, 0.82, 1.07, 1.24, 1.43, 1.60]
  line [0.28, 0.46, 0.71, 0.87, 1.17, 1.26, 1.49, 1.66]
  line [0.31, 0.55, 0.71, 0.87, 1.25, 1.32, 1.51, 1.69]

Loading

---
config:
  themeVariables:
    xyChart:
      plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
  title "Full Tree (🔴 PR | 🔵 main | 🟢 release)"
  x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
  y-axis "time (ms)"
  line [0.84, 2.32, 4.06, 7.25, 12.54, 26.07, 55.53, 116.14]
  line [0.84, 2.08, 4.42, 6.76, 12.92, 26.26, 56.43, 113.78]
  line [1.09, 2.13, 4.33, 6.81, 13.39, 27.52, 57.56, 119.91]

Loading

[github-actions]

Bundle size comparison (import * as ... in PR vs import * as ... in target):

🟢 Decreased (max -0.28%)	➖ Unchanged	🔴 Increased (max 2.79%)	❔ Unknown
6	130	187	1

import * as ... in PR vs import * as ... in target (did bundle size increase?):

Test	tsdown
std_isBeingTranspiled.ts	15.71 kB ($${\color{red}+2.8\%}$$)
std_getTargetShaderLanguage.ts	15.77 kB ($${\color{red}+2.8\%}$$)
std_getShaderStage.ts	15.76 kB

import { ... } in PR vs import * as ... in PR (is the library tree-Shakeable?):

Test	tsdown
tgpu_init.ts	270.29 kB ($${\color{green}-3.3\%}$$)
tgpu_initFromDevice.ts	269.74 kB ($${\color{green}-3.5\%}$$)
tgpu_resolve.ts	170.54 kB ($${\color{green}-39.0\%}$$)
tgpu_resolveWithContext.ts	170.48 kB ($${\color{green}-39.0\%}$$)
tgpu_bindGroupLayout.ts	73.82 kB ($${\color{green}-73.6\%}$$)
tgpu_mutableAccessor.ts	68.55 kB ($${\color{green}-75.5\%}$$)
tgpu_accessor.ts	68.54 kB ($${\color{green}-75.5\%}$$)
tgpu_privateVar.ts	67.23 kB ($${\color{green}-75.9\%}$$)
tgpu_workgroupVar.ts	67.23 kB ($${\color{green}-75.9\%}$$)
tgpu_const.ts	66.65 kB ($${\color{green}-76.2\%}$$)
tgpu_lazy.ts	66.45 kB ($${\color{green}-76.2\%}$$)
tgpu_fragmentFn.ts	38.92 kB ($${\color{green}-86.1\%}$$)
tgpu_fn.ts	38.87 kB ($${\color{green}-86.1\%}$$)
tgpu_vertexFn.ts	38.74 kB ($${\color{green}-86.1\%}$$)
tgpu_computeFn.ts	38.44 kB ($${\color{green}-86.2\%}$$)
tgpu_vertexLayout.ts	27.57 kB ($${\color{green}-90.1\%}$$)
tgpu_comptime.ts	15.18 kB ($${\color{green}-94.6\%}$$)
tgpu_unroll.ts	1.75 kB ($${\color{green}-99.4\%}$$)
tgpu_slot.ts	1.70 kB ($${\color{green}-99.4\%}$$)

If you wish to run a comparison for other, slower bundlers, run the 'Tree-shake test' from the GitHub Actions menu.

[pullfrog]

ℹ️ No critical issues — minor suggestions inline.

Reviewed changes

• std.getShaderStage — new comptime probe returning 'vertex' / 'fragment' / 'compute' when the calling 'use gpu' function resolves (transitively) inside a matching entry fn, undefined otherwise. Built on the shaderStageSlot the entry-fn and auto-IO resolvers already bind.
• Type merge — TgpuShaderStage merged with the internal ShaderStage (previously ... | undefined) into a single public ShaderStage = 'compute' | 'vertex' | 'fragment', exported from indexNamedExports; the slot's unbound sentinel moved to null. git grep confirms no TgpuShaderStage references remain, and TgpuShaderStage was never publicly exported, so there's no downstream break.
• Tests — five new getShaderStage cases, including a committed snapshot that pins per-stage duplication of a shared function used across vertex/fragment/compute resolution.
• Docs — getShaderStage added to the resolution-environment section of utils.mdx.

Tests (environment.test.ts, bindGroupLayout.test.ts, glsl suite) and both package typechecks pass. The rename and sentinel change are behavior-preserving: logGenerator.ts's shaderStageSlot.$ === 'vertex' check treats null exactly as it did undefined.

  ^{｜ Fix all ➔ ｜ Fix 👍s ➔ ｜ View workflow run ｜ Using DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏}

[pullfrog]

ℹ️ The delta since the prior review is clean — one docs nit from the earlier review remains open, nothing new introduced.

Reviewed changes (delta between prior review 7c27874c and current head 49361fbb)

• ShaderStage export home moved — the type is no longer re-exported from typegpu/~internal (internal.ts); it now ships only through the public typegpu entry (indexNamedExports.ts). GlslGenerator was updated to import it from 'typegpu' instead, and git grep confirms no ~internal consumer of ShaderStage remains.
• Docs prose — the comma splice in the resolution-environment section was fixed by splitting the sentence.

Verified on 49361fbb: @typegpu/gl and core package typechecks pass, and environment.test.ts (17 tests) passes.

  ^{｜ Fix it ➔ ｜ View workflow run ｜ Using DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏}

[aleksanderkatan]

Nice! 🦈🏚️

[pullfrog]

✅ No new issues found.

Reviewed changes

• Docs table fix — the env-table header row in the resolution-environment section now keeps its trailing |, consistent with every other row, resolving the previously open nit at utils.mdx:133. The comma-splice thread was already resolved in an earlier push.

The delta of the PR's own files since the prior review is limited to this docs correction — the core std.getShaderStage() implementation, the TgpuShaderStage→ShaderStage type merge, and the test coverage are unchanged from the previously reviewed, clean pass. The remaining noise in the incremental diff is base-branch (feat/fn-call-through-gen) shader-generator refactor work, not part of this PR. Both prior Pullfrog threads are now resolved, and nothing new was introduced.

  ^{｜ View workflow run ｜ Using DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏}