use crate::ecs::camera::queries::query_active_camera_matrices;
use crate::ecs::mesh::components::Vertex;
use crate::ecs::prefab::resources::mesh_cache_iter;
use crate::ecs::world::{GLOBAL_TRANSFORM, SKIN, VISIBILITY, WATER, World};
use std::collections::HashMap;
use super::{
BUFFER_GROWTH_FACTOR, CullingUniforms, DrawIndexedIndirect, MeshBoundsData,
VOLUMETRIC_CUBE_MESH_NAME, WaterInstance, WaterMaterialData, WaterMeshData, WaterMeshPass,
WaterObjectData, WaterType, WaterUniforms, compute_volumetric_bounds, create_unit_cube_mesh,
extract_frustum_planes,
};
impl WaterMeshPass {
pub(super) fn ensure_mesh_in_buffer(
&mut self,
device: &wgpu::Device,
queue: &wgpu::Queue,
name: &str,
vertices: &[Vertex],
indices: &[u32],
) -> Option<u32> {
if let Some((mesh_id, _)) = self.meshes.get(name) {
return Some(*mesh_id);
}
let vertex_count = vertices.len() as u32;
let index_count = indices.len() as u32;
let required_vertices = self.current_vertex_offset as usize + vertices.len();
if required_vertices > self.vertex_buffer_size {
let new_size = ((required_vertices as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.vertex_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Unified Vertex Buffer (Resized)"),
size: (std::mem::size_of::<Vertex>() * new_size) as u64,
usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.vertex_buffer_size = new_size;
self.current_vertex_offset = 0;
self.current_index_offset = 0;
self.meshes.clear();
self.next_mesh_id = 0;
return None;
}
let required_indices = self.current_index_offset as usize + indices.len();
if required_indices > self.index_buffer_size {
let new_size = ((required_indices as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.index_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Unified Index Buffer (Resized)"),
size: (std::mem::size_of::<u32>() * new_size) as u64,
usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.index_buffer_size = new_size;
self.current_vertex_offset = 0;
self.current_index_offset = 0;
self.meshes.clear();
self.next_mesh_id = 0;
return None;
}
let vertex_offset = self.current_vertex_offset;
let index_offset = self.current_index_offset;
queue.write_buffer(
&self.vertex_buffer,
(vertex_offset as usize * std::mem::size_of::<Vertex>()) as u64,
bytemuck::cast_slice(vertices),
);
queue.write_buffer(
&self.index_buffer,
(index_offset as usize * std::mem::size_of::<u32>()) as u64,
bytemuck::cast_slice(indices),
);
let bounds = Self::compute_mesh_bounds(vertices);
let mesh_id = self.next_mesh_id;
self.next_mesh_id += 1;
if mesh_id as usize >= self.mesh_bounds_buffer_size {
let new_size = ((self.mesh_bounds_buffer_size as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.mesh_bounds_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Bounds Buffer (Resized)"),
size: (std::mem::size_of::<MeshBoundsData>() * new_size) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.mesh_bounds_buffer_size = new_size;
self.culling_bind_group = None;
}
queue.write_buffer(
&self.mesh_bounds_buffer,
(mesh_id as usize * std::mem::size_of::<MeshBoundsData>()) as u64,
bytemuck::cast_slice(&[bounds]),
);
self.meshes.insert(
name.to_string(),
(
mesh_id,
WaterMeshData {
vertex_offset,
index_offset,
index_count,
},
),
);
self.current_vertex_offset += vertex_count;
self.current_index_offset += index_count;
Some(mesh_id)
}
fn compute_mesh_bounds(vertices: &[Vertex]) -> MeshBoundsData {
if vertices.is_empty() {
return MeshBoundsData {
center: [0.0, 0.0, 0.0],
radius: 0.0,
};
}
let mut min = [f32::MAX, f32::MAX, f32::MAX];
let mut max = [f32::MIN, f32::MIN, f32::MIN];
for vertex in vertices {
for index in 0..3 {
min[index] = min[index].min(vertex.position[index]);
max[index] = max[index].max(vertex.position[index]);
}
}
let center = [
(min[0] + max[0]) * 0.5,
(min[1] + max[1]) * 0.5,
(min[2] + max[2]) * 0.5,
];
let mut radius = 0.0f32;
for vertex in vertices {
let dx = vertex.position[0] - center[0];
let dy = vertex.position[1] - center[1];
let dz = vertex.position[2] - center[2];
let dist = (dx * dx + dy * dy + dz * dz).sqrt();
radius = radius.max(dist);
}
MeshBoundsData { center, radius }
}
fn ensure_volumetric_mesh_with_bounds(
&mut self,
device: &wgpu::Device,
queue: &wgpu::Queue,
volume_size: [f32; 3],
) -> Option<u32> {
if !self.meshes.contains_key(VOLUMETRIC_CUBE_MESH_NAME) {
let (vertices, indices) = create_unit_cube_mesh();
self.ensure_mesh_in_buffer(
device,
queue,
VOLUMETRIC_CUBE_MESH_NAME,
&vertices,
&indices,
);
}
let (_, base_mesh_data) = self.meshes.get(VOLUMETRIC_CUBE_MESH_NAME)?;
let base_mesh_data = base_mesh_data.clone();
let bounds = compute_volumetric_bounds(volume_size);
let mesh_id = self.next_mesh_id;
self.next_mesh_id += 1;
if mesh_id as usize >= self.mesh_bounds_buffer_size {
let new_size = ((self.mesh_bounds_buffer_size as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.mesh_bounds_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Bounds Buffer (Resized)"),
size: (std::mem::size_of::<MeshBoundsData>() * new_size) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.mesh_bounds_buffer_size = new_size;
self.culling_bind_group = None;
}
queue.write_buffer(
&self.mesh_bounds_buffer,
(mesh_id as usize * std::mem::size_of::<MeshBoundsData>()) as u64,
bytemuck::cast_slice(&[bounds]),
);
self.meshes.insert(
format!("__volumetric_{}__", mesh_id),
(mesh_id, base_mesh_data),
);
Some(mesh_id)
}
pub(super) fn rebuild_render_bind_group(&mut self, device: &wgpu::Device) {
self.render_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("Water Mesh Render Bind Group"),
layout: &self.render_bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: self.uniform_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: self.material_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: self.instance_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: self.visible_indices_buffer.as_entire_binding(),
},
],
});
}
pub(super) fn rebuild_culling_bind_group(&mut self, device: &wgpu::Device) {
self.culling_bind_group = Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("Water Mesh Culling Bind Group"),
layout: &self.culling_bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: self.instance_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: self.culling_uniform_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: self.mesh_bounds_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: self.object_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 4,
resource: self.indirect_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 5,
resource: self.visible_indices_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 6,
resource: self.material_buffer.as_entire_binding(),
},
],
}));
}
pub(super) fn prepare_pass(
&mut self,
device: &wgpu::Device,
queue: &wgpu::Queue,
world: &World,
) {
self.horizontal_batches.clear();
self.vertical_batches.clear();
self.volumetric_batches.clear();
self.batch_count = 0;
self.object_count = 0;
let Some(camera_matrices) = query_active_camera_matrices(world) else {
return;
};
let view = camera_matrices.view;
let projection = camera_matrices.projection;
let view_projection = projection * view;
let mut material_to_index: HashMap<WaterMaterialData, u32> = HashMap::new();
let mut materials_data: Vec<WaterMaterialData> = Vec::new();
let mut instances: Vec<WaterInstance> = Vec::new();
let mut objects: Vec<WaterObjectData> = Vec::new();
struct MeshBatchInfo {
mesh_name: String,
water_type: WaterType,
instances: Vec<(u32, WaterInstance)>,
volume_size: Option<[f32; 3]>,
}
let mut mesh_batches: Vec<MeshBatchInfo> = Vec::new();
let mut needed_meshes: HashMap<String, (Vec<Vertex>, Vec<u32>)> = HashMap::new();
for entity in world
.core
.query_entities(WATER | GLOBAL_TRANSFORM | VISIBILITY)
{
if world.core.entity_has_components(entity, SKIN) {
continue;
}
let Some(visibility) = world.core.get_visibility(entity) else {
continue;
};
if !visibility.visible {
continue;
}
let Some(water) = world.core.get_water(entity) else {
continue;
};
if !water.enabled {
continue;
}
let Some(global_transform) = world.core.get_global_transform(entity) else {
continue;
};
let mesh_name: String;
if water.is_volumetric {
mesh_name = VOLUMETRIC_CUBE_MESH_NAME.to_string();
if !needed_meshes.contains_key(VOLUMETRIC_CUBE_MESH_NAME) {
let (vertices, indices) = create_unit_cube_mesh();
needed_meshes
.insert(VOLUMETRIC_CUBE_MESH_NAME.to_string(), (vertices, indices));
}
} else {
let Some(render_mesh) = world.core.get_render_mesh(entity) else {
continue;
};
mesh_name = render_mesh.name.clone();
if !needed_meshes.contains_key(&mesh_name) {
for (name, mesh) in mesh_cache_iter(&world.resources.mesh_cache) {
if name == &mesh_name {
needed_meshes.insert(
name.clone(),
(mesh.vertices.clone(), mesh.indices.clone()),
);
break;
}
}
}
if !needed_meshes.contains_key(&mesh_name) {
continue;
}
}
let material_data = WaterMaterialData {
base_color: water.base_color,
water_color: water.water_color,
wave_height: water.wave_height,
choppy: water.choppy,
speed: water.speed,
freq: water.frequency,
specular_strength: water.specular_strength,
fresnel_power: water.fresnel_power,
volume_shape: water.volume_shape as u32,
volume_flow_type: water.volume_flow_type as u32,
volume_size: water.volume_size,
is_volumetric: if water.is_volumetric { 1 } else { 0 },
flow_direction: water.flow_direction,
flow_strength: water.flow_strength,
_flow_padding: 0.0,
};
let material_index = *material_to_index.entry(material_data).or_insert_with(|| {
let index = materials_data.len() as u32;
materials_data.push(material_data);
index
});
let model: [[f32; 4]; 4] = global_transform.0.into();
let instance = WaterInstance {
model,
material_index,
_padding: [0; 3],
};
let instance_index = instances.len() as u32;
instances.push(instance);
let water_type = if water.is_volumetric {
WaterType::Volumetric
} else if water.is_vertical {
WaterType::Vertical
} else {
WaterType::Horizontal
};
if water_type == WaterType::Volumetric {
mesh_batches.push(MeshBatchInfo {
mesh_name: mesh_name.clone(),
water_type,
instances: vec![(instance_index, instance)],
volume_size: Some(water.volume_size),
});
} else {
let existing_batch = mesh_batches
.iter_mut()
.find(|b| b.mesh_name == mesh_name && b.water_type == water_type);
if let Some(batch) = existing_batch {
batch.instances.push((instance_index, instance));
} else {
mesh_batches.push(MeshBatchInfo {
mesh_name: mesh_name.clone(),
water_type,
instances: vec![(instance_index, instance)],
volume_size: None,
});
}
}
}
if instances.is_empty() {
return;
}
for (name, (vertices, indices)) in &needed_meshes {
if name != VOLUMETRIC_CUBE_MESH_NAME {
self.ensure_mesh_in_buffer(device, queue, name, vertices, indices);
}
}
let mut volumetric_mesh_ids: HashMap<usize, u32> = HashMap::new();
for (batch_idx, batch_info) in mesh_batches.iter().enumerate() {
if batch_info.water_type == WaterType::Volumetric
&& let Some(volume_size) = batch_info.volume_size
&& let Some(mesh_id) =
self.ensure_volumetric_mesh_with_bounds(device, queue, volume_size)
{
volumetric_mesh_ids.insert(batch_idx, mesh_id);
}
}
let mut bind_groups_need_rebuild = false;
if instances.len() > self.instance_buffer_size {
let new_size = ((instances.len() as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.instance_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Instance Buffer (Resized)"),
size: (std::mem::size_of::<WaterInstance>() * new_size) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.instance_buffer_size = new_size;
bind_groups_need_rebuild = true;
}
if materials_data.len() > self.material_buffer_size {
let new_size = ((materials_data.len() as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.material_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Material Buffer (Resized)"),
size: (std::mem::size_of::<WaterMaterialData>() * new_size) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.material_buffer_size = new_size;
bind_groups_need_rebuild = true;
}
for (batch_idx, batch_info) in mesh_batches.iter().enumerate() {
let mesh_id = if batch_info.water_type == WaterType::Volumetric {
let Some(&mesh_id) = volumetric_mesh_ids.get(&batch_idx) else {
continue;
};
mesh_id
} else {
let Some((mesh_id, _)) = self.meshes.get(&batch_info.mesh_name) else {
continue;
};
*mesh_id
};
let batch_id = self.batch_count as u32;
for (instance_index, _) in &batch_info.instances {
objects.push(WaterObjectData {
instance_index: *instance_index,
mesh_id,
batch_id,
_padding: 0,
});
}
match batch_info.water_type {
WaterType::Horizontal => self.horizontal_batches.push(self.batch_count),
WaterType::Vertical => self.vertical_batches.push(self.batch_count),
WaterType::Volumetric => {
let position = if let Some((idx, _)) = batch_info.instances.first() {
let inst = &instances[*idx as usize];
[inst.model[3][0], inst.model[3][1], inst.model[3][2]]
} else {
[0.0, 0.0, 0.0]
};
self.volumetric_batches.push((self.batch_count, position));
}
}
self.batch_count += 1;
}
self.object_count = objects.len() as u32;
if objects.len() > self.object_buffer_size {
let new_size = ((objects.len() as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.object_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Object Buffer (Resized)"),
size: (std::mem::size_of::<WaterObjectData>() * new_size) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.object_buffer_size = new_size;
bind_groups_need_rebuild = true;
}
if instances.len() > self.visible_indices_buffer_size {
let new_size = ((instances.len() as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.visible_indices_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Visible Indices Buffer (Resized)"),
size: (std::mem::size_of::<u32>() * new_size) as u64,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.visible_indices_buffer_size = new_size;
bind_groups_need_rebuild = true;
}
if self.batch_count > self.indirect_buffer_size {
let new_size = ((self.batch_count as f32) * BUFFER_GROWTH_FACTOR) as usize;
self.indirect_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Indirect Buffer (Resized)"),
size: (std::mem::size_of::<DrawIndexedIndirect>() * new_size) as u64,
usage: wgpu::BufferUsages::INDIRECT
| wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.indirect_reset_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Water Mesh Indirect Reset Buffer (Resized)"),
size: (std::mem::size_of::<DrawIndexedIndirect>() * new_size) as u64,
usage: wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
self.indirect_buffer_size = new_size;
bind_groups_need_rebuild = true;
}
if bind_groups_need_rebuild {
self.rebuild_render_bind_group(device);
self.culling_bind_group = None;
}
let view_inverse = nalgebra_glm::inverse(&view);
let camera_position = [
view_inverse[(0, 3)],
view_inverse[(1, 3)],
view_inverse[(2, 3)],
1.0,
];
self.camera_position = [camera_position[0], camera_position[1], camera_position[2]];
self.volumetric_batches.sort_by(|a, b| {
let dist_a = (a.1[0] - self.camera_position[0]).powi(2)
+ (a.1[1] - self.camera_position[1]).powi(2)
+ (a.1[2] - self.camera_position[2]).powi(2);
let dist_b = (b.1[0] - self.camera_position[0]).powi(2)
+ (b.1[1] - self.camera_position[1]).powi(2)
+ (b.1[2] - self.camera_position[2]).powi(2);
dist_b
.partial_cmp(&dist_a)
.unwrap_or(std::cmp::Ordering::Equal)
});
let time = world.resources.window.timing.uptime_milliseconds as f32 / 1000.0;
let uniforms = WaterUniforms {
view: view.into(),
projection: projection.into(),
view_projection: view_projection.into(),
camera_position,
time,
_padding: [0.0; 3],
};
queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));
let frustum_planes = extract_frustum_planes(&view_projection);
let culling_uniforms = CullingUniforms {
frustum_planes,
object_count: self.object_count,
_padding: [0; 3],
};
queue.write_buffer(
&self.culling_uniform_buffer,
0,
bytemuck::cast_slice(&[culling_uniforms]),
);
queue.write_buffer(
&self.material_buffer,
0,
bytemuck::cast_slice(&materials_data),
);
queue.write_buffer(&self.instance_buffer, 0, bytemuck::cast_slice(&instances));
queue.write_buffer(&self.object_buffer, 0, bytemuck::cast_slice(&objects));
let mut indirect_commands = Vec::with_capacity(self.batch_count);
let mut indirect_reset_commands = Vec::with_capacity(self.batch_count);
let mut visible_indices = Vec::with_capacity(instances.len());
let mut first_instance_offset = 0u32;
for (batch_idx, batch_info) in mesh_batches.iter().enumerate() {
let mesh_data = if batch_info.water_type == WaterType::Volumetric {
let Some(&mesh_id) = volumetric_mesh_ids.get(&batch_idx) else {
continue;
};
let mesh_name = format!("__volumetric_{}__", mesh_id);
let Some((_, data)) = self.meshes.get(&mesh_name) else {
continue;
};
data.clone()
} else {
let Some((_, data)) = self.meshes.get(&batch_info.mesh_name) else {
continue;
};
data.clone()
};
let instance_count = batch_info.instances.len() as u32;
indirect_commands.push(DrawIndexedIndirect {
index_count: mesh_data.index_count,
instance_count,
first_index: mesh_data.index_offset,
base_vertex: mesh_data.vertex_offset as i32,
first_instance: first_instance_offset,
});
indirect_reset_commands.push(DrawIndexedIndirect {
index_count: mesh_data.index_count,
instance_count: 0,
first_index: mesh_data.index_offset,
base_vertex: mesh_data.vertex_offset as i32,
first_instance: first_instance_offset,
});
for (instance_index, _) in &batch_info.instances {
visible_indices.push(*instance_index);
}
first_instance_offset += instance_count;
}
queue.write_buffer(
&self.indirect_buffer,
0,
bytemuck::cast_slice(&indirect_commands),
);
queue.write_buffer(
&self.indirect_reset_buffer,
0,
bytemuck::cast_slice(&indirect_reset_commands),
);
queue.write_buffer(
&self.visible_indices_buffer,
0,
bytemuck::cast_slice(&visible_indices),
);
self.indirect_reset_count = indirect_commands.len();
if self.culling_bind_group.is_none() {
self.rebuild_culling_bind_group(device);
}
}
}