"""The per-material resources a custom shader declares, read from its source.
A ``ShaderMaterial`` reaches its own resources through group 2 of the custom
shader ABI: binding 0 is the std140 uniform block, and binding 1 upwards are
SEPARATED textures and samplers. Which name lives at which binding is stated by
the shader author, in the shader, so both backends have to read it back out of
the source to bind anything to it.
This module is that reader, and it is deliberately backend-free: the Vulkan
renderer parses the GLSL a material was written in, the web exporter parses the
WGSL naga emitted from it, and the two must agree on the same (name, binding)
pairs or a texture would land on the wrong sampler in the browser. Keeping one
implementation of the model, with one notion of what is supported, is what makes
that agreement checkable rather than hopeful.
Anything the model cannot express -- a combined sampler, a storage image, a cube
or array texture, a second uniform block -- raises :class:`ShaderBindingError`
naming the declaration and its line, so an author hears about it from the tool
that read the file rather than from a pipeline failure three layers down.
The reader is a lexical one. Comments are stripped before anything is matched,
so commenting a declaration out really does remove it from the model, but the
preprocessor is NOT evaluated: a declaration inside ``#if 0``, one produced by a
macro, and one living in an ``#include``d chunk are all invisible to it (an
included texture would then fail pipeline creation, having never been declared
in the layout). Declare a material's textures and samplers in the stage source
itself.
"""
from __future__ import annotations
import re
from dataclasses import dataclass
__all__ = [
"MATERIAL_GROUP",
"MAX_MATERIAL_SAMPLERS",
"MAX_MATERIAL_TEXTURES",
"MaterialBindings",
"ResourceBinding",
"ShaderBindingError",
"material_bindings_for_stages",
"parse_material_bindings",
]
#: The bind group / descriptor set a material's own resources live in.
MATERIAL_GROUP = 2
#: Binding 0 of the material group is the std140 custom-uniform block.
UNIFORM_BINDING = 0
# WebGPU guarantees only 16 sampled textures and 16 samplers per shader stage,
# shared with everything else the stage binds, so a material's own share is
# capped well below that. A shader wanting more is refused by name rather than
# left to fail pipeline creation in a browser.
MAX_MATERIAL_TEXTURES = 8
MAX_MATERIAL_SAMPLERS = 4
# GLSL: layout(... set = 2, binding = N ...) uniform <type> <name>;
_GLSL_DECL = re.compile(
r"layout\s*\(([^)]*)\)\s*uniform\s+(\w+)\s+(\w+)\s*(\[[^\]]*\])?\s*[;{]",
re.MULTILINE,
)
_GLSL_SET = re.compile(r"\bset\s*=\s*(\d+)")
_GLSL_BINDING = re.compile(r"\bbinding\s*=\s*(\d+)")
# WGSL: @group(2) @binding(N) var[<space>] name : type;
_WGSL_DECL = re.compile(
r"@group\s*\(\s*(\d+)\s*\)\s*@binding\s*\(\s*(\d+)\s*\)\s*(var(?:<[^>]*>)?)\s+(\w+)\s*:\s*([^;]+);",
re.MULTILINE,
)
# The one texture type both backends support: a sampled 2D float texture.
_GLSL_TEXTURE_TYPE = "texture2D"
_GLSL_SAMPLER_TYPE = "sampler"
_WGSL_TEXTURE_TYPE = re.compile(r"^texture_2d\s*<\s*f32\s*>$")
_WGSL_SAMPLER_TYPE = re.compile(r"^sampler$")
[docs]
class ShaderBindingError(ValueError):
"""A custom shader declares a per-material resource neither backend can bind."""
[docs]
@dataclass(frozen=True)
class ResourceBinding:
"""One texture or sampler a shader declares in the material group."""
name: str
binding: int
#: 1-based line of the declaration in the source it was parsed from.
line: int
[docs]
@dataclass(frozen=True)
class MaterialBindings:
"""The material group's declared textures and samplers, sorted by binding."""
textures: tuple[ResourceBinding, ...] = ()
samplers: tuple[ResourceBinding, ...] = ()
[docs]
@property
def is_empty(self) -> bool:
"""True when the shader declares no texture and no sampler."""
return not self.textures and not self.samplers
[docs]
def texture_binding(self, name: str) -> int:
"""The binding number the named texture is declared at, or -1 if absent."""
for tex in self.textures:
if tex.name == name:
return tex.binding
return -1
[docs]
def merged_with(self, other: MaterialBindings) -> MaterialBindings:
"""Combine two stages' models, requiring any shared name to agree.
A vertex and a fragment stage may each declare the same texture (a
displacement map read in both), and they must place it at the same
binding: there is one descriptor set behind both stages.
"""
textures = _merge("texture", self.textures, other.textures)
samplers = _merge("sampler", self.samplers, other.samplers)
_reject_binding_clashes(textures + samplers)
return MaterialBindings(textures=textures, samplers=samplers)
def _merge(kind: str, a: tuple[ResourceBinding, ...], b: tuple[ResourceBinding, ...]) -> tuple[ResourceBinding, ...]:
by_name: dict[str, ResourceBinding] = {r.name: r for r in a}
for res in b:
seen = by_name.get(res.name)
if seen is not None and seen.binding != res.binding:
raise ShaderBindingError(
f"{kind} `{res.name}` is declared at binding {seen.binding} in one stage and "
f"binding {res.binding} in the other (line {res.line}); both stages share one "
"descriptor set, so a name must sit at the same binding in each."
)
by_name.setdefault(res.name, res)
return tuple(sorted(by_name.values(), key=lambda r: r.binding))
[docs]
def parse_material_bindings(source: str | None, *, language: str = "glsl") -> MaterialBindings:
"""Read the material group's texture and sampler declarations from one stage.
``language`` is ``"glsl"`` (a desktop shader stage) or ``"wgsl"`` (what the
web exporter transpiled, or a hand-written escape hatch stage). Returns an
empty model for ``None`` / textureless sources. Raises
:class:`ShaderBindingError` on a declaration the ABI cannot express, naming
the offending line.
"""
if not source:
return MaterialBindings()
if language == "glsl":
textures, samplers = _parse_glsl(_strip_comments(source, nested_blocks=False))
elif language == "wgsl":
textures, samplers = _parse_wgsl(_strip_comments(source, nested_blocks=True))
else:
raise ValueError(f"language must be 'glsl' or 'wgsl', got {language!r}")
if len(textures) > MAX_MATERIAL_TEXTURES:
raise ShaderBindingError(
f"a ShaderMaterial may sample at most {MAX_MATERIAL_TEXTURES} textures "
f"(this shader declares {len(textures)}); WebGPU guarantees only 16 sampled "
"textures per shader stage, shared with everything else the stage binds."
)
if len(samplers) > MAX_MATERIAL_SAMPLERS:
raise ShaderBindingError(
f"a ShaderMaterial may declare at most {MAX_MATERIAL_SAMPLERS} samplers "
f"(this shader declares {len(samplers)})."
)
_reject_binding_clashes(textures + samplers)
return MaterialBindings(
textures=tuple(sorted(textures, key=lambda r: r.binding)),
samplers=tuple(sorted(samplers, key=lambda r: r.binding)),
)
[docs]
def material_bindings_for_stages(
vertex_source: str | None,
fragment_source: str | None,
*,
wgsl_vertex: str | None = None,
wgsl_fragment: str | None = None,
) -> MaterialBindings:
"""The material-group model of a whole material, from whichever stages it has.
Hand-written WGSL wins over GLSL for a stage, matching the escape hatch the
web exporter honours, so a material written in WGSL reports the bindings the
browser will actually see. The two stages are merged: a name declared in both
must sit at the same binding.
"""
if wgsl_vertex is not None or wgsl_fragment is not None:
vertex = parse_material_bindings(wgsl_vertex, language="wgsl")
fragment = parse_material_bindings(wgsl_fragment, language="wgsl")
else:
vertex = parse_material_bindings(vertex_source, language="glsl")
fragment = parse_material_bindings(fragment_source, language="glsl")
return vertex.merged_with(fragment)
def _reject_binding_clashes(resources: tuple[ResourceBinding, ...]) -> None:
seen: dict[int, ResourceBinding] = {}
for res in sorted(resources, key=lambda r: (r.binding, r.line)):
clash = seen.get(res.binding)
if clash is not None:
raise ShaderBindingError(
f"`{res.name}` (line {res.line}) and `{clash.name}` (line {clash.line}) both sit "
f"at group {MATERIAL_GROUP} binding {res.binding}; each per-material resource "
"needs a binding of its own."
)
seen[res.binding] = res
def _line_of(source: str, offset: int) -> int:
return source.count("\n", 0, offset) + 1
def _strip_comments(source: str, *, nested_blocks: bool) -> str:
"""Blank out ``//`` and ``/* */`` comments, keeping every character position.
Comment bodies become spaces rather than disappearing, so an offset into the
result still names the line the author wrote it on. Neither language has
string literals to protect. WGSL block comments nest; GLSL's end at the first
``*/``, which ``nested_blocks`` selects between.
"""
if "/" not in source:
return source
out = list(source)
index, end = 0, len(source)
while index < end:
if source[index] != "/" or index + 1 >= end:
index += 1
continue
following = source[index + 1]
if following == "/":
while index < end and source[index] != "\n":
out[index] = " "
index += 1
elif following == "*":
depth = 1
out[index] = out[index + 1] = " "
index += 2
while index < end and depth:
if source[index] == "\n":
index += 1
elif nested_blocks and source.startswith("/*", index):
depth += 1
out[index] = out[index + 1] = " "
index += 2
elif source.startswith("*/", index):
depth -= 1
out[index] = out[index + 1] = " "
index += 2
else:
out[index] = " "
index += 1
else:
index += 1
return "".join(out)
def _parse_glsl(source: str) -> tuple[tuple[ResourceBinding, ...], tuple[ResourceBinding, ...]]:
textures: list[ResourceBinding] = []
samplers: list[ResourceBinding] = []
for match in _GLSL_DECL.finditer(source):
qualifiers, type_name, name, array = match.group(1), match.group(2), match.group(3), match.group(4)
set_match = _GLSL_SET.search(qualifiers)
binding_match = _GLSL_BINDING.search(qualifiers)
if set_match is None or int(set_match.group(1)) != MATERIAL_GROUP or binding_match is None:
continue
binding = int(binding_match.group(1))
line = _line_of(source, match.start())
if binding == UNIFORM_BINDING:
# The uniform block itself: its members are the std140 layout, not a
# binding model. A texture there would displace it.
if type_name in (_GLSL_TEXTURE_TYPE, _GLSL_SAMPLER_TYPE):
raise ShaderBindingError(
f"`{name}` (line {line}) is declared at group {MATERIAL_GROUP} binding "
f"{UNIFORM_BINDING}, which is reserved for the material's uniform block. "
"Move textures and samplers to binding 1 upwards."
)
continue
if array:
raise ShaderBindingError(
f"`{name}` (line {line}) is an array of resources; a ShaderMaterial binds one "
"texture or sampler per binding."
)
if type_name == _GLSL_TEXTURE_TYPE:
textures.append(ResourceBinding(name, binding, line))
elif type_name == _GLSL_SAMPLER_TYPE:
samplers.append(ResourceBinding(name, binding, line))
else:
raise ShaderBindingError(
f"`{name}` (line {line}) is declared as `{type_name}` at group {MATERIAL_GROUP} "
f"binding {binding}. A ShaderMaterial supports separated `texture2D` and "
"`sampler` declarations there; combine them in code, e.g. "
"`texture(sampler2D(tex, samp), uv)`."
)
return tuple(textures), tuple(samplers)
def _parse_wgsl(source: str) -> tuple[tuple[ResourceBinding, ...], tuple[ResourceBinding, ...]]:
textures: list[ResourceBinding] = []
samplers: list[ResourceBinding] = []
for match in _WGSL_DECL.finditer(source):
group, binding = int(match.group(1)), int(match.group(2))
var_kind, name, type_name = match.group(3), match.group(4), match.group(5).strip()
if group != MATERIAL_GROUP:
continue
line = _line_of(source, match.start())
if binding == UNIFORM_BINDING:
if var_kind != "var<uniform>":
raise ShaderBindingError(
f"`{name}` (line {line}) is a `{var_kind}` at group {MATERIAL_GROUP} binding "
f"{UNIFORM_BINDING}, which is reserved for the material's uniform block."
)
continue
if var_kind != "var":
raise ShaderBindingError(
f"`{name}` (line {line}) is a `{var_kind}` at group {MATERIAL_GROUP} binding "
f"{binding}. Only sampled textures and samplers can be bound there; a second "
"uniform or storage buffer has nowhere to go in the material group."
)
if _WGSL_TEXTURE_TYPE.match(type_name):
textures.append(ResourceBinding(name, binding, line))
elif _WGSL_SAMPLER_TYPE.match(type_name):
samplers.append(ResourceBinding(name, binding, line))
else:
raise ShaderBindingError(
f"`{name}` (line {line}) is declared as `{type_name}` at group {MATERIAL_GROUP} "
f"binding {binding}. A ShaderMaterial supports `texture_2d<f32>` and `sampler` "
"there; storage, depth, multisampled, cube and array textures are not bound by "
"the custom-material path."
)
return tuple(textures), tuple(samplers)