"""Built-in water pass (design D16, RM-E1).
A dedicated forward pass (its own GLSL shader pair, not a ShaderMaterial) drawn
in the TRANSPARENT slot after the scene colour/depth copy (A7/A8), so it can
refract what is behind it. The vertex stage generates a tessellated grid with no
vertex buffer (like the grid / tilemap passes) and displaces it by a sum of
Gerstner waves driven by the FrameGlobals wind + time; the fragment stage
refracts + depth-fades against the scene copy, mixes a Fresnel reflection from
the environment cubemap, and adds shore foam.
Set 0 is the shared forward descriptor set (FrameGlobals b13, environment cube
b4, scene colour/depth b14/b15); set 1 is this pass's own camera UBO + per-frame
surface SSBO; set 2 is the engine bindless texture array, from which a
``PlanarReflection3D`` mirror feed is sampled by index (RM-E2). Zero-cost when
unused: nothing renders and no surface triggers the scene-copy split when the
scene holds no visible ``WaterSurface3D``, and a surface with no reflection is
byte-identical to RM-E1 (the planar branch is gated on the per-surface index).
"""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.descriptors import (
allocate_descriptor_set,
create_descriptor_set_layout,
create_pool_for_types,
write_ssbo_descriptor,
write_ubo_descriptor,
)
from ..gpu.memory import create_buffer, upload_numpy
from ..gpu.pipeline import PipelineSpec, build_pipeline, create_shader_module
from ..materials.shader_compiler import compile_shader
from .water_pack import WATER_SURFACE_DTYPE, pack_water_surface
__all__ = ["WaterPass", "WATER_SURFACE_DTYPE"]
log = logging.getLogger(__name__)
MAX_SURFACES = 64
_CAMERA_BYTES = 160 # mat4 view_proj + mat4 inv_view_proj + vec4 cam_pos + vec4 cam_planes
# std430 stride: mat4 (64) + 7 vec4 (112) = 176 bytes; mirrors the GLSL struct.
_SURFACE_BYTES = WATER_SURFACE_DTYPE.itemsize
[docs]
class WaterPass:
"""Renders every submitted ``WaterSurface3D`` as a Gerstner water grid."""
def __init__(self, engine: Any):
self._engine = engine
self._pipeline: Any = None
self._pipeline_layout: Any = None
self._vert_module: Any = None
self._frag_module: Any = None
self._set_layout: Any = None
self._pool: Any = None
self._set: Any = None
self._cam_buf: Any = None
self._cam_mem: Any = None
self._surf_buf: Any = None
self._surf_mem: Any = None
self._forward_ssbo_layout: Any = None
self._ready = False
# Per-frame submissions: (model_matrix, WaterMaterial, subdivisions, size, reflection_tex).
self._submissions: list[tuple[np.ndarray, Any, int, tuple[float, float], int]] = []
# -------------------------------------------------------------- setup
[docs]
def setup(self, ssbo_layout: Any, render_pass: Any = None) -> None:
"""Create GPU resources. ``ssbo_layout`` is the shared forward set-0 layout."""
e = self._engine
device = e.ctx.device
phys = e.ctx.physical_device
self._forward_ssbo_layout = ssbo_layout
# Own descriptor set (set 1): binding 0 camera UBO, binding 1 surface SSBO.
stage = vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT
self._set_layout = create_descriptor_set_layout(
device,
[
(0, vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, stage, 1),
(1, vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, stage, 1),
],
)
self._pool = create_pool_for_types(
device,
{
vk.VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER: 1,
vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER: 1,
},
max_sets=1,
)
self._set = allocate_descriptor_set(device, self._pool, self._set_layout)
host = vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
self._cam_buf, self._cam_mem = create_buffer(
device, phys, _CAMERA_BYTES, vk.VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, host
)
surf_size = MAX_SURFACES * _SURFACE_BYTES
self._surf_buf, self._surf_mem = create_buffer(
device, phys, surf_size, vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, host
)
upload_numpy(device, self._cam_mem, np.zeros(_CAMERA_BYTES, dtype=np.uint8))
upload_numpy(device, self._surf_mem, np.zeros(surf_size, dtype=np.uint8))
write_ubo_descriptor(device, self._set, 0, self._cam_buf, _CAMERA_BYTES)
write_ssbo_descriptor(device, self._set, 1, self._surf_buf, surf_size)
shader_dir = e.shader_dir
self._vert_module = create_shader_module(device, compile_shader(shader_dir / "water.vert"))
self._frag_module = create_shader_module(device, compile_shader(shader_dir / "water.frag"))
self._create_pipeline(device, render_pass or e.render_pass, e.extent)
self._ready = True
log.debug("Water pass initialized")
def _create_pipeline(self, device: Any, render_pass: Any, extent: tuple[int, int]) -> None:
"""Alpha-blended, depth-tested (no write) pipeline; set 0 forward, set 1 own."""
spec = PipelineSpec(
name="water",
topology=vk.VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
vertex_stride=0, # grid generated in the shader from gl_VertexIndex
cull_mode=vk.VK_CULL_MODE_NONE,
depth_test=True,
depth_write=False, # transparent surface: test against opaque depth, no write
depth_compare=vk.VK_COMPARE_OP_LESS_OR_EQUAL,
blend="alpha",
dst_alpha_factor=vk.VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
# Set 2 is the engine bindless texture array: a PlanarReflection3D
# capture (RM-E2) is sampled from it by its bindless index, exactly
# like the tilemap / uber texture path. Bound but unread when no
# surface carries a reflection (byte-identical to RM-E1 water).
set_layouts=(self._forward_ssbo_layout, self._set_layout, self._engine.texture_descriptor_layout),
push_size=16, # int surface_index (+ pad)
push_stages=vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
# Thin G-buffer (D3): match the HDR pass attachment count; frag writes
# only location 0 so the gbuffer normal attachment is left untouched.
attachment_count=(2 if self._engine.gbuffer_active else 1),
)
self._pipeline, self._pipeline_layout = build_pipeline(
device,
spec,
render_pass,
extent,
vert_module=self._vert_module,
frag_module=self._frag_module,
)
[docs]
def rebuild_pipeline(self, render_pass: Any) -> None:
"""Recreate the pipeline against a different render pass (HDR / G-buffer toggle)."""
if not self._ready:
return
device = self._engine.ctx.device
if self._pipeline:
vk.vkDestroyPipeline(device, self._pipeline, None)
if self._pipeline_layout:
vk.vkDestroyPipelineLayout(device, self._pipeline_layout, None)
self._create_pipeline(device, render_pass, self._engine.extent)
# -------------------------------------------------------------- per-frame
[docs]
def begin_frame(self) -> None:
self._submissions.clear()
[docs]
def submit(
self,
model_matrix: np.ndarray,
material: Any,
subdivisions: int,
size: tuple[float, float],
reflection_tex: int = -1,
) -> None:
"""Queue a water surface for this frame.
``reflection_tex`` is the bindless index of a :class:`PlanarReflection3D`
capture feeding a true planar mirror (RM-E2), or ``-1`` for none.
"""
self._submissions.append(
(model_matrix, material, int(subdivisions), (float(size[0]), float(size[1])), int(reflection_tex))
)
[docs]
@property
def has_surfaces(self) -> bool:
return bool(self._submissions)
@staticmethod
def _pack_surface(
model: np.ndarray, mat: Any, subdivisions: int, size: tuple[float, float], reflection_tex: int = -1
) -> np.ndarray:
# Delegates to the shared, Vulkan-free packer so desktop + web produce the
# byte-identical std430 record (design D16, RM-E1w).
return pack_water_surface(model, mat, subdivisions, size, reflection_tex)
[docs]
def render(
self,
cmd: Any,
view_proj: np.ndarray,
cam_pos: np.ndarray,
near: float,
far: float,
has_skybox: bool,
extent: tuple[int, int],
ssbo_set: Any,
submissions: list | None = None,
) -> None:
"""Record the water draws. ``ssbo_set`` is the shared forward set 0."""
subs = self._submissions if submissions is None else submissions
if not self._ready or not subs:
return
device = self._engine.ctx.device
# Camera UBO: view_proj + inv_view_proj (both column-major) + cam_pos + planes.
vp_t = np.ascontiguousarray(view_proj.T, dtype=np.float32)
inv_vp_t = np.ascontiguousarray(np.linalg.inv(view_proj).T, dtype=np.float32)
cam = np.zeros(_CAMERA_BYTES // 4, dtype=np.float32)
cam[0:16] = vp_t.ravel()
cam[16:32] = inv_vp_t.ravel()
cam[32:35] = cam_pos[:3]
cam[36:40] = (float(near), float(far), 1.0 if has_skybox else 0.0, 0.0)
upload_numpy(device, self._cam_mem, cam.view(np.uint8))
# Surface SSBO.
count = min(len(subs), MAX_SURFACES)
packed = np.zeros(count, dtype=WATER_SURFACE_DTYPE)
for i in range(count):
model, mat, sub, size, refl_tex = subs[i]
packed[i] = self._pack_surface(model, mat, sub, size, refl_tex)
upload_numpy(device, self._surf_mem, packed.view(np.uint8))
vk_vp = vk.VkViewport(x=0.0, y=0.0, width=float(extent[0]), height=float(extent[1]), minDepth=0.0, maxDepth=1.0)
vk.vkCmdSetViewport(cmd, 0, 1, [vk_vp])
scissor = vk.VkRect2D(offset=vk.VkOffset2D(x=0, y=0), extent=vk.VkExtent2D(width=extent[0], height=extent[1]))
vk.vkCmdSetScissor(cmd, 0, 1, [scissor])
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline_layout, 0, 1, [ssbo_set], 0, None
)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline_layout, 1, 1, [self._set], 0, None
)
# Set 2: engine bindless texture array (PlanarReflection3D feed, RM-E2).
tex_ds = self._engine.texture_descriptor_set
if tex_ds:
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_GRAPHICS, self._pipeline_layout, 2, 1, [tex_ds], 0, None
)
ffi = vk.ffi
for i in range(count):
_model, _mat, sub, _size, _refl = subs[i]
n = max(int(sub), 1)
pc = np.array([i, 0, 0, 0], dtype=np.int32)
cbuf = ffi.new("char[]", pc.tobytes())
vk._vulkan.lib.vkCmdPushConstants(
cmd,
self._pipeline_layout,
vk.VK_SHADER_STAGE_VERTEX_BIT | vk.VK_SHADER_STAGE_FRAGMENT_BIT,
0,
16,
cbuf,
)
vk.vkCmdDraw(cmd, n * n * 6, 1, 0, 0)
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
for obj, fn in [
(self._pipeline, vk.vkDestroyPipeline),
(self._pipeline_layout, vk.vkDestroyPipelineLayout),
(self._vert_module, vk.vkDestroyShaderModule),
(self._frag_module, vk.vkDestroyShaderModule),
(self._set_layout, vk.vkDestroyDescriptorSetLayout),
(self._pool, vk.vkDestroyDescriptorPool),
]:
if obj:
fn(device, obj, None)
for buf, mem in [(self._cam_buf, self._cam_mem), (self._surf_buf, self._surf_mem)]:
if buf:
vk.vkDestroyBuffer(device, buf, None)
if mem:
vk.vkFreeMemory(device, mem, None)
self._ready = False