632 lines
22 KiB
Rust
632 lines
22 KiB
Rust
use std::rc::Rc;
|
||
|
||
use exr::prelude::{ReadChannels, ReadLayers};
|
||
use glam::{Vec2, Vec3};
|
||
use wgpu::util::DeviceExt;
|
||
|
||
use super::{with_device, with_queue, DrawCall, Pipeline};
|
||
use crate::{
|
||
loaders::mesh::{InstanceRaw, Mesh, Vertex},
|
||
paths,
|
||
texture::HeightmapTexture,
|
||
};
|
||
|
||
pub struct SnowConfig
|
||
{
|
||
pub depth_map_path: String,
|
||
pub heightmap_path: String,
|
||
pub terrain_size: Vec2,
|
||
pub resolution: (u32, u32),
|
||
}
|
||
|
||
impl SnowConfig
|
||
{
|
||
pub fn default() -> Self
|
||
{
|
||
Self {
|
||
depth_map_path: paths::textures::snow_depth(),
|
||
heightmap_path: paths::textures::terrain_heightmap(),
|
||
terrain_size: Vec2::new(1000.0, 1000.0),
|
||
resolution: (1000, 1000),
|
||
}
|
||
}
|
||
}
|
||
|
||
pub struct ClipmapConfig
|
||
{
|
||
pub num_levels: u32,
|
||
pub grid_size: u32,
|
||
pub base_cell_size: f32,
|
||
}
|
||
|
||
impl Default for ClipmapConfig
|
||
{
|
||
fn default() -> Self
|
||
{
|
||
Self {
|
||
num_levels: 3,
|
||
grid_size: 192,
|
||
base_cell_size: 0.25,
|
||
}
|
||
}
|
||
}
|
||
|
||
struct ClipmapLevel
|
||
{
|
||
mesh: Rc<Mesh>,
|
||
instance_buffer: wgpu::Buffer,
|
||
}
|
||
|
||
pub struct SnowLayer
|
||
{
|
||
pub depth_texture: wgpu::Texture,
|
||
pub depth_texture_view: wgpu::TextureView,
|
||
pub depth_bind_group: wgpu::BindGroup,
|
||
pub width: u32,
|
||
pub height: u32,
|
||
pub deform_bind_group: wgpu::BindGroup,
|
||
pub deform_pipeline: wgpu::ComputePipeline,
|
||
pub deform_params_buffer: wgpu::Buffer,
|
||
pub terrain_size: Vec2,
|
||
|
||
clipmap_config: ClipmapConfig,
|
||
levels: Vec<ClipmapLevel>,
|
||
displacement_bind_group: wgpu::BindGroup,
|
||
heightmap_texture: HeightmapTexture,
|
||
depth_sampler: wgpu::Sampler,
|
||
}
|
||
|
||
fn create_instance_buffer() -> wgpu::Buffer
|
||
{
|
||
with_device(|device| {
|
||
device.create_buffer(&wgpu::BufferDescriptor {
|
||
label: Some("Snow Clipmap Instance Buffer"),
|
||
size: std::mem::size_of::<InstanceRaw>() as u64,
|
||
usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
|
||
mapped_at_creation: false,
|
||
})
|
||
})
|
||
}
|
||
|
||
impl SnowLayer
|
||
{
|
||
pub fn load(config: &SnowConfig) -> anyhow::Result<Self>
|
||
{
|
||
println!("\n=== Loading Snow Layer ===");
|
||
println!("Depth map path: {}", config.depth_map_path);
|
||
println!("Heightmap path: {}", config.heightmap_path);
|
||
println!("Terrain size: {:?}", config.terrain_size);
|
||
|
||
let (depth_data, width, height) = Self::load_depth_map(&config.depth_map_path)?;
|
||
|
||
println!("Using EXR dimensions: {}×{}", width, height);
|
||
|
||
let (depth_texture, depth_texture_view, depth_bind_group) =
|
||
Self::create_depth_texture(&depth_data, width, height);
|
||
|
||
let (deform_pipeline, deform_bind_group, deform_params_buffer) =
|
||
Self::create_deform_pipeline(&depth_texture_view);
|
||
|
||
let heightmap_texture = with_device(|device| {
|
||
with_queue(|queue| HeightmapTexture::load(device, queue, &config.heightmap_path))
|
||
})?;
|
||
|
||
let depth_sampler = with_device(|device| {
|
||
device.create_sampler(&wgpu::SamplerDescriptor {
|
||
label: Some("Snow Depth Sampler"),
|
||
address_mode_u: wgpu::AddressMode::ClampToEdge,
|
||
address_mode_v: wgpu::AddressMode::ClampToEdge,
|
||
mag_filter: wgpu::FilterMode::Nearest,
|
||
min_filter: wgpu::FilterMode::Nearest,
|
||
..Default::default()
|
||
})
|
||
});
|
||
|
||
let displacement_bind_group = Self::create_displacement_bind_group(
|
||
&heightmap_texture,
|
||
&depth_texture_view,
|
||
&depth_sampler,
|
||
);
|
||
|
||
let clipmap_config = ClipmapConfig::default();
|
||
let levels = Vec::new();
|
||
|
||
Ok(Self {
|
||
depth_texture,
|
||
depth_texture_view,
|
||
depth_bind_group,
|
||
width,
|
||
height,
|
||
deform_bind_group,
|
||
deform_pipeline,
|
||
deform_params_buffer,
|
||
terrain_size: config.terrain_size,
|
||
clipmap_config,
|
||
levels,
|
||
displacement_bind_group,
|
||
heightmap_texture,
|
||
depth_sampler,
|
||
})
|
||
}
|
||
|
||
fn load_depth_map(path: &str) -> anyhow::Result<(Vec<f32>, u32, u32)>
|
||
{
|
||
println!("Loading snow depth map from: {}", path);
|
||
|
||
let image = exr::prelude::read()
|
||
.no_deep_data()
|
||
.largest_resolution_level()
|
||
.all_channels()
|
||
.all_layers()
|
||
.all_attributes()
|
||
.from_file(path)?;
|
||
|
||
let layer = &image.layer_data[0];
|
||
let width = layer.size.width() as u32;
|
||
let height = layer.size.height() as u32;
|
||
|
||
println!(" Layer size: {}×{}", width, height);
|
||
println!(
|
||
" Available channels: {:?}",
|
||
layer
|
||
.channel_data
|
||
.list
|
||
.iter()
|
||
.map(|c| &c.name)
|
||
.collect::<Vec<_>>()
|
||
);
|
||
|
||
let channel = layer
|
||
.channel_data
|
||
.list
|
||
.iter()
|
||
.find(|c| format!("{:?}", c.name).contains("\"R\""))
|
||
.or_else(|| layer.channel_data.list.first())
|
||
.ok_or_else(|| anyhow::anyhow!("No channels found in EXR"))?;
|
||
|
||
println!(" Using channel: {:?}", channel.name);
|
||
|
||
let depths: Vec<f32> = channel.sample_data.values_as_f32().collect();
|
||
|
||
let min_value = depths.iter().cloned().fold(f32::INFINITY, f32::min);
|
||
let max_value = depths.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
let avg_value = depths.iter().sum::<f32>() / depths.len() as f32;
|
||
let non_zero_count = depths.iter().filter(|&&v| v > 0.0001).count();
|
||
|
||
println!(" Total values: {}", depths.len());
|
||
println!(
|
||
" Min: {:.6}, Max: {:.6}, Avg: {:.6}",
|
||
min_value, max_value, avg_value
|
||
);
|
||
println!(
|
||
" Non-zero values: {} ({:.1}%)",
|
||
non_zero_count,
|
||
(non_zero_count as f32 / depths.len() as f32) * 100.0
|
||
);
|
||
|
||
if max_value < 0.0001
|
||
{
|
||
println!(
|
||
" ⚠️ WARNING: All values are effectively zero! Snow depth map may be invalid."
|
||
);
|
||
}
|
||
else
|
||
{
|
||
println!(" ✓ Snow depth data loaded successfully");
|
||
}
|
||
|
||
Ok((depths, width, height))
|
||
}
|
||
|
||
fn create_depth_texture(
|
||
depth_data: &[f32],
|
||
width: u32,
|
||
height: u32,
|
||
) -> (wgpu::Texture, wgpu::TextureView, wgpu::BindGroup)
|
||
{
|
||
with_device(|device| {
|
||
let size = wgpu::Extent3d {
|
||
width,
|
||
height,
|
||
depth_or_array_layers: 1,
|
||
};
|
||
|
||
let texture = device.create_texture(&wgpu::TextureDescriptor {
|
||
label: Some("Snow Depth Texture"),
|
||
size,
|
||
mip_level_count: 1,
|
||
sample_count: 1,
|
||
dimension: wgpu::TextureDimension::D2,
|
||
format: wgpu::TextureFormat::R32Float,
|
||
usage: wgpu::TextureUsages::TEXTURE_BINDING
|
||
| wgpu::TextureUsages::COPY_DST
|
||
| wgpu::TextureUsages::STORAGE_BINDING,
|
||
view_formats: &[],
|
||
});
|
||
|
||
let data_bytes: &[u8] = bytemuck::cast_slice(depth_data);
|
||
|
||
with_queue(|queue| {
|
||
queue.write_texture(
|
||
wgpu::TexelCopyTextureInfo {
|
||
texture: &texture,
|
||
mip_level: 0,
|
||
origin: wgpu::Origin3d::ZERO,
|
||
aspect: wgpu::TextureAspect::All,
|
||
},
|
||
data_bytes,
|
||
wgpu::TexelCopyBufferLayout {
|
||
offset: 0,
|
||
bytes_per_row: Some(width * 4),
|
||
rows_per_image: Some(height),
|
||
},
|
||
size,
|
||
);
|
||
});
|
||
|
||
let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());
|
||
|
||
let bind_group_layout =
|
||
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||
label: Some("Snow Depth Bind Group Layout"),
|
||
entries: &[wgpu::BindGroupLayoutEntry {
|
||
binding: 0,
|
||
visibility: wgpu::ShaderStages::FRAGMENT | wgpu::ShaderStages::COMPUTE,
|
||
ty: wgpu::BindingType::Texture {
|
||
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||
view_dimension: wgpu::TextureViewDimension::D2,
|
||
multisampled: false,
|
||
},
|
||
count: None,
|
||
}],
|
||
});
|
||
|
||
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||
label: Some("Snow Depth Bind Group"),
|
||
layout: &bind_group_layout,
|
||
entries: &[wgpu::BindGroupEntry {
|
||
binding: 0,
|
||
resource: wgpu::BindingResource::TextureView(&texture_view),
|
||
}],
|
||
});
|
||
|
||
(texture, texture_view, bind_group)
|
||
})
|
||
}
|
||
|
||
fn create_deform_pipeline(
|
||
depth_texture_view: &wgpu::TextureView,
|
||
) -> (wgpu::ComputePipeline, wgpu::BindGroup, wgpu::Buffer)
|
||
{
|
||
with_device(|device| {
|
||
let shader_source = std::fs::read_to_string(&paths::shaders::snow_deform())
|
||
.expect("Failed to load snow deform shader");
|
||
|
||
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||
label: Some("Snow Deform Shader"),
|
||
source: wgpu::ShaderSource::Wgsl(shader_source.into()),
|
||
});
|
||
|
||
let params_buffer = device.create_buffer(&wgpu::BufferDescriptor {
|
||
label: Some("Snow Deform Params"),
|
||
size: 32,
|
||
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
|
||
mapped_at_creation: false,
|
||
});
|
||
|
||
let bind_group_layout =
|
||
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||
label: Some("Snow Deform Bind Group Layout"),
|
||
entries: &[
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 0,
|
||
visibility: wgpu::ShaderStages::COMPUTE,
|
||
ty: wgpu::BindingType::StorageTexture {
|
||
access: wgpu::StorageTextureAccess::ReadWrite,
|
||
format: wgpu::TextureFormat::R32Float,
|
||
view_dimension: wgpu::TextureViewDimension::D2,
|
||
},
|
||
count: None,
|
||
},
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 1,
|
||
visibility: wgpu::ShaderStages::COMPUTE,
|
||
ty: wgpu::BindingType::Buffer {
|
||
ty: wgpu::BufferBindingType::Uniform,
|
||
has_dynamic_offset: false,
|
||
min_binding_size: None,
|
||
},
|
||
count: None,
|
||
},
|
||
],
|
||
});
|
||
|
||
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||
label: Some("Snow Deform Bind Group"),
|
||
layout: &bind_group_layout,
|
||
entries: &[
|
||
wgpu::BindGroupEntry {
|
||
binding: 0,
|
||
resource: wgpu::BindingResource::TextureView(depth_texture_view),
|
||
},
|
||
wgpu::BindGroupEntry {
|
||
binding: 1,
|
||
resource: params_buffer.as_entire_binding(),
|
||
},
|
||
],
|
||
});
|
||
|
||
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
|
||
label: Some("Snow Deform Pipeline Layout"),
|
||
bind_group_layouts: &[&bind_group_layout],
|
||
immediate_size: 0,
|
||
});
|
||
|
||
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||
label: Some("Snow Deform Pipeline"),
|
||
layout: Some(&pipeline_layout),
|
||
module: &shader,
|
||
entry_point: Some("deform"),
|
||
compilation_options: Default::default(),
|
||
cache: None,
|
||
});
|
||
|
||
(pipeline, bind_group, params_buffer)
|
||
})
|
||
}
|
||
|
||
fn create_displacement_bind_group(
|
||
heightmap: &HeightmapTexture,
|
||
depth_view: &wgpu::TextureView,
|
||
depth_sampler: &wgpu::Sampler,
|
||
) -> wgpu::BindGroup
|
||
{
|
||
with_device(|device| {
|
||
let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||
label: Some("Snow Displacement Bind Group Layout"),
|
||
entries: &[
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 0,
|
||
visibility: wgpu::ShaderStages::VERTEX,
|
||
ty: wgpu::BindingType::Texture {
|
||
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||
view_dimension: wgpu::TextureViewDimension::D2,
|
||
multisampled: false,
|
||
},
|
||
count: None,
|
||
},
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 1,
|
||
visibility: wgpu::ShaderStages::VERTEX,
|
||
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
|
||
count: None,
|
||
},
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 2,
|
||
visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
|
||
ty: wgpu::BindingType::Texture {
|
||
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||
view_dimension: wgpu::TextureViewDimension::D2,
|
||
multisampled: false,
|
||
},
|
||
count: None,
|
||
},
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 3,
|
||
visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
|
||
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
|
||
count: None,
|
||
},
|
||
],
|
||
});
|
||
|
||
device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||
label: Some("Snow Displacement Bind Group"),
|
||
layout: &layout,
|
||
entries: &[
|
||
wgpu::BindGroupEntry {
|
||
binding: 0,
|
||
resource: wgpu::BindingResource::TextureView(&heightmap.view),
|
||
},
|
||
wgpu::BindGroupEntry {
|
||
binding: 1,
|
||
resource: wgpu::BindingResource::Sampler(&heightmap.sampler),
|
||
},
|
||
wgpu::BindGroupEntry {
|
||
binding: 2,
|
||
resource: wgpu::BindingResource::TextureView(depth_view),
|
||
},
|
||
wgpu::BindGroupEntry {
|
||
binding: 3,
|
||
resource: wgpu::BindingResource::Sampler(depth_sampler),
|
||
},
|
||
],
|
||
})
|
||
})
|
||
}
|
||
|
||
fn generate_clipmap_grid(&self, level: u32) -> Mesh
|
||
{
|
||
let cell_size = self.clipmap_config.base_cell_size * (1 << level) as f32;
|
||
let grid_size = self.clipmap_config.grid_size;
|
||
let half_extent = (grid_size as f32 * cell_size) / 2.0;
|
||
|
||
let mut vertices = Vec::new();
|
||
let mut indices = Vec::new();
|
||
|
||
for z in 0..=grid_size
|
||
{
|
||
for x in 0..=grid_size
|
||
{
|
||
vertices.push(Vertex {
|
||
position: [
|
||
-half_extent + x as f32 * cell_size,
|
||
0.0,
|
||
-half_extent + z as f32 * cell_size,
|
||
],
|
||
normal: [0.0, 1.0, 0.0],
|
||
uv: [0.0, 0.0],
|
||
});
|
||
}
|
||
}
|
||
|
||
let row = grid_size + 1;
|
||
for z in 0..grid_size
|
||
{
|
||
for x in 0..grid_size
|
||
{
|
||
let i0 = z * row + x;
|
||
let i1 = i0 + 1;
|
||
let i2 = i0 + row;
|
||
let i3 = i2 + 1;
|
||
indices.push(i0);
|
||
indices.push(i2);
|
||
indices.push(i1);
|
||
indices.push(i1);
|
||
indices.push(i2);
|
||
indices.push(i3);
|
||
}
|
||
}
|
||
|
||
let vertex_buffer = with_device(|device| {
|
||
device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||
label: Some(&format!("Snow Clipmap Level {} Vertex Buffer", level)),
|
||
contents: bytemuck::cast_slice(&vertices),
|
||
usage: wgpu::BufferUsages::VERTEX,
|
||
})
|
||
});
|
||
|
||
let index_buffer = with_device(|device| {
|
||
device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||
label: Some(&format!("Snow Clipmap Level {} Index Buffer", level)),
|
||
contents: bytemuck::cast_slice(&indices),
|
||
usage: wgpu::BufferUsages::INDEX,
|
||
})
|
||
});
|
||
|
||
let mut aabb_min = Vec3::splat(f32::MAX);
|
||
let mut aabb_max = Vec3::splat(f32::MIN);
|
||
for v in &vertices
|
||
{
|
||
let p = Vec3::from(v.position);
|
||
aabb_min = aabb_min.min(p);
|
||
aabb_max = aabb_max.max(p);
|
||
}
|
||
|
||
Mesh {
|
||
vertex_buffer,
|
||
index_buffer,
|
||
num_indices: indices.len() as u32,
|
||
aabb_min,
|
||
aabb_max,
|
||
cpu_positions: Vec::new(),
|
||
cpu_indices: Vec::new(),
|
||
}
|
||
}
|
||
|
||
pub fn update(&mut self, center: Vec3)
|
||
{
|
||
if self.levels.is_empty()
|
||
{
|
||
for level in 0..self.clipmap_config.num_levels
|
||
{
|
||
let mesh = self.generate_clipmap_grid(level);
|
||
let instance_buffer = create_instance_buffer();
|
||
|
||
self.levels.push(ClipmapLevel {
|
||
mesh: Rc::new(mesh),
|
||
instance_buffer,
|
||
});
|
||
}
|
||
}
|
||
|
||
with_queue(|queue| {
|
||
for (level, clipmap_level) in self.levels.iter().enumerate()
|
||
{
|
||
let cell_size = self.clipmap_config.base_cell_size * (1u32 << level) as f32;
|
||
let snapped_x = (center.x / cell_size).round() * cell_size;
|
||
let snapped_z = (center.z / cell_size).round() * cell_size;
|
||
let y_bias = -(level as f32) * 0.005;
|
||
let model =
|
||
glam::Mat4::from_translation(glam::Vec3::new(snapped_x, y_bias, snapped_z));
|
||
let instance_data = InstanceRaw {
|
||
model: model.to_cols_array_2d(),
|
||
dissolve_amount: 0.0,
|
||
_padding: [0.0; 3],
|
||
};
|
||
queue.write_buffer(
|
||
&clipmap_level.instance_buffer,
|
||
0,
|
||
bytemuck::cast_slice(&[instance_data]),
|
||
);
|
||
}
|
||
});
|
||
}
|
||
|
||
pub fn get_draw_calls(&self) -> Vec<DrawCall>
|
||
{
|
||
self.levels
|
||
.iter()
|
||
.rev()
|
||
.map(|level| DrawCall {
|
||
vertex_buffer: level.mesh.vertex_buffer.clone(),
|
||
index_buffer: level.mesh.index_buffer.clone(),
|
||
num_indices: level.mesh.num_indices,
|
||
model: glam::Mat4::IDENTITY,
|
||
pipeline: Pipeline::SnowClipmap,
|
||
instance_buffer: Some(level.instance_buffer.clone()),
|
||
num_instances: 1,
|
||
tile_scale: 2.0,
|
||
enable_dissolve: false,
|
||
enable_snow_light: true,
|
||
displacement_bind_group: Some(self.displacement_bind_group.clone()),
|
||
entity: None,
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
pub fn deform_at_position(&self, position: Vec3, radius: f32, depth: f32)
|
||
{
|
||
with_queue(|queue| {
|
||
let params_data = [
|
||
position.x,
|
||
position.z,
|
||
radius,
|
||
depth,
|
||
self.terrain_size.x,
|
||
self.terrain_size.y,
|
||
0.0,
|
||
0.0,
|
||
];
|
||
let params_bytes: &[u8] = bytemuck::cast_slice(¶ms_data);
|
||
|
||
queue.write_buffer(&self.deform_params_buffer, 0, params_bytes);
|
||
});
|
||
|
||
with_device(|device| {
|
||
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
|
||
label: Some("Snow Deform Encoder"),
|
||
});
|
||
|
||
{
|
||
let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
|
||
label: Some("Snow Deform Pass"),
|
||
timestamp_writes: None,
|
||
});
|
||
|
||
compute_pass.set_pipeline(&self.deform_pipeline);
|
||
compute_pass.set_bind_group(0, &self.deform_bind_group, &[]);
|
||
|
||
let workgroup_size = 16;
|
||
let dispatch_x = (self.width + workgroup_size - 1) / workgroup_size;
|
||
let dispatch_y = (self.height + workgroup_size - 1) / workgroup_size;
|
||
compute_pass.dispatch_workgroups(dispatch_x, dispatch_y, 1);
|
||
}
|
||
|
||
with_queue(|queue| {
|
||
queue.submit(Some(encoder.finish()));
|
||
});
|
||
});
|
||
}
|
||
}
|