diff --git a/Cargo.toml b/Cargo.toml index c9e548b..6d0fc01 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,13 @@ egui = "0.33.3" egui-wgpu = "0.33.3" egui-winit = "0.33.3" env_logger = "0.11.8" +indicatif = "0.18.3" +itertools = "0.14.0" pollster = "0.4.0" +rand = "0.9.2" wgpu = "27.0.1" winit = "0.30.12" + +[profile.release] +opt-level = 3 + diff --git a/shaders/cube.wgsl b/shaders/cube.wgsl new file mode 100644 index 0000000..6f149df --- /dev/null +++ b/shaders/cube.wgsl @@ -0,0 +1,439 @@ +struct PushConstants +{ + view_projection: mat4x4, + transform: mat4x4, + eye_position: vec3, + root_color: vec4, + root_subdivided: u32, +} + +var constants: PushConstants; + +struct VertexOutput +{ + @builtin(position) pos: vec4, + @location(0) eye_pos: vec3, + @location(1) world_pos: vec3, + @location(2) cube_pos: vec3, + @location(3) root_color: vec4, + @location(4) root_subdivided: u32 +} + +@vertex +fn vertex_main(@builtin(vertex_index) index: u32, @builtin(instance_index) instance_index: u32) -> VertexOutput +{ + let side_length = u32(100); + let offset = vec3(vec3( + instance_index % side_length, + (instance_index / side_length) % side_length, + instance_index / (side_length * side_length) + )); + + let cube_vertices = array, 8>( + vec3(0., 0., 0.), + vec3(0., 0., 1.), + vec3(1., 0., 1.), + vec3(1., 0., 0.), + + vec3(0., 1., 0.), + vec3(0., 1., 1.), + vec3(1., 1., 1.), + vec3(1., 1., 0.), + ); + + let cube_faces = array( + // Bottom face + 1, 0, 2, 3, + + // Top face + 4, 5, 7, 6, + + // Side faces + 0, 1, 4, 5, + 1, 2, 5, 6, + 2, 3, 6, 7, + 3, 0, 7, 4, + ); + + let quad_index = index / (3 * 2); + let triangle_index = index % (3 * 2); + let triangle_map = array( + 0, 1, 2, 1, 3, 2 + ); + + let vertex = cube_vertices[cube_faces[quad_index * 4 + triangle_map[triangle_index]]]; + + var output: VertexOutput; + output.pos = constants.view_projection * constants.transform * vec4(vertex + offset, 1.); + output.eye_pos = constants.eye_position; + output.world_pos = (constants.transform * vec4(vertex + offset, 1.)).xyz; + output.cube_pos = offset; + + output.root_color = constants.root_color; + output.root_subdivided = constants.root_subdivided; + return output; +} + +fn is_voxel(p: vec3) -> bool +{ + //return true; + return length(vec3(p) - vec3(16.)) < 16.; +} + +const N: u32 = 4; +const N3: u32 = N * N * N; +const MAX_DEPTH: u32 = 4; + +struct StructureTile +{ + children: array, +} + +struct ColorTile +{ + colors: array, 64> +} + +// BG +@group(0) @binding(0) var structure_tiles : array; +@group(0) @binding(1) var color_tiles : array; + +//@fragment +fn traverse(max_depth: u32, root_color: vec4, root_subdivided: bool, eye_pos: vec3, ray_dir: vec3) -> vec4 +{ + if max_depth == 0 || !root_subdivided + { + return root_color; + } + let chunk_size = N * N * N * N; + let depth_child_size_lut = array(N*N*N, N*N, N, 1); + + var stack_nodes = array(0, 0, 0, 0); + var stack_child_pos = array, 4>(vec3(0), vec3(0), vec3(0), vec3(0)); + var stack_node_offset = array, 4>(vec3(0), vec3(0), vec3(0), vec3(0)); + var stack_ptr = 0; + var current_child_size = chunk_size / N; + var current_child_pos = vec3(0); + var current_node_offset = vec3(0); + var current_depth = 1; + + var current_tile_index = 0; + var current_children_data = structure_tiles[current_tile_index]; + var current_children_colors = color_tiles[current_tile_index]; + + // Intersection parameters + let dist = 1. / ray_dir; + var offset = - (eye_pos * f32(chunk_size)) * dist; + + // Interesect with root + let t0 = fma(dist, vec3(0.), offset); + let t1 = fma(dist, vec3(chunk_size), offset); + + let tmin = min(t0, t1); + let tmax = max(t0, t1); + + let t_near = max(max(tmin.x, max(tmin.y, tmin.z)), 0.); + let t_far = min(tmax.x, min(tmax.y, tmax.z)); + + let step = select(vec3(-1), vec3(1), ray_dir > vec3(0.)); + + var start_pos = (eye_pos * f32(chunk_size)) + t_near * ray_dir; + var hit_pos = start_pos; + offset = - start_pos * dist; + var t = 0.; + + current_child_pos = vec3( + clamp(i32(floor(start_pos.x)), 0, i32(chunk_size) - 1), + clamp(i32(floor(start_pos.y)), 0, i32(chunk_size) - 1), + clamp(i32(floor(start_pos.z)), 0, i32(chunk_size) - 1), + ) / i32(current_child_size); + + //return vec4(vec3(current_child_pos) / 4., 1.); + + for(var iter = 0; iter < 300; iter++) + { + // Retrieve current child information + let child_index = current_child_pos.x + current_child_pos.y * i32(N) + current_child_pos.z * i32(N * N); + + let child_u32 = current_children_data.children[child_index]; + let child_subdivided = (child_u32 >> 31) == 1; + let child_color = current_children_colors.colors[child_index]; + + if child_color.w != 0. // Child is solid + && (max_depth == u32(current_depth) || !child_subdivided) + { + // Sample mat + let voxel_pos = current_node_offset + current_child_pos * i32(current_child_size); + //return vec4(child_color.xyz, 1.); + return child_color+ vec4(vec3(f32(iter) / 200.), 1.); + } + + // Advance + // Project current child + let global_child_pos = current_child_pos * i32(current_child_size) + current_node_offset; + + let t0 = fma(dist, vec3(global_child_pos), offset); + let t1 = fma(dist, vec3(global_child_pos) + vec3(current_child_size), offset); + + let tmin = min(t0, t1); + let tmax = max(t0, t1); + + let t_near = max(max(tmin.x, max(tmin.y, tmin.z)), 0.); + let t_far = min(tmax.x, min(tmax.y, tmax.z)); + + if child_subdivided + { + // Push operation + + stack_nodes[stack_ptr] = current_tile_index; + stack_child_pos[stack_ptr] = current_child_pos; + stack_node_offset[stack_ptr] = current_node_offset; + stack_ptr ++; + + // Retrieve child information + current_tile_index = i32(child_u32 & 0x3fffffff); + current_children_data = structure_tiles[current_tile_index]; + current_children_colors = color_tiles[current_tile_index]; + + // Determine child of the child + let hit_pos = start_pos + ray_dir * t_near; + + let next_node_offset = current_node_offset + current_child_pos * i32(current_child_size); + let next_child_size = current_child_size / N; + current_child_pos = + (clamp(vec3(floor(hit_pos)), global_child_pos, global_child_pos + vec3(i32(current_child_size - 1))) - next_node_offset) / i32(next_child_size); + + current_child_size = next_child_size; + current_node_offset = next_node_offset; + current_depth ++; + + } + else + { + // ADVANCE + let advance_mask = min_mask3i32(tmax); + let next_child = current_child_pos + advance_mask * step; + if any(next_child < vec3(0)) || any(next_child >= vec3(i32(N))) + { + let aligned_child = select(vec3(0), vec3(i32(N)), vec3(step) > vec3(0)); + let masked_aligned = advance_mask * ((aligned_child * i32(current_child_size)) + current_node_offset); + let exiting_axis = masked_aligned.x + masked_aligned.y + masked_aligned.z + 256; + + // HARDCODED FOR N = 4 + let ctz = countTrailingZeros(exiting_axis) / 2; + let exiting_depth = 4 - ctz; + + if exiting_depth == 0 // Getting out of root + { + return vec4(f32(iter) / 200.); + discard; + } + + // Restore destination depth + current_depth = exiting_depth; + stack_ptr = current_depth - 1; + + current_tile_index = stack_nodes[stack_ptr]; + current_children_data = structure_tiles[current_tile_index]; + current_children_colors = color_tiles[current_tile_index]; + + current_node_offset = stack_node_offset[stack_ptr]; + current_child_pos = stack_child_pos[stack_ptr] + step * advance_mask; + + current_child_size = depth_child_size_lut[current_depth - 1]; + }else{ + current_child_pos = next_child; + } + } + } + + return vec4(100., 0., 100., 100.); +} + +fn sample_mat(pos: vec3) -> f32 +{ + var voxel = pos; + var div = 1; + var overlay = 1.; + for(var i = 1; i <= 4; i++) + { + let x = (voxel.x / div + voxel.y / div + voxel.z / div) % 2 == 0; + overlay -= select(0., 1. / (f32(i) * 2.5), x); + div *= 4; + } + return overlay; +} + +@fragment +fn fragment_tree_main(in: VertexOutput) -> @location(0) vec4 +{ + var hit_pos = vec3(0.); + let dir = normalize(in.world_pos.xyz - in.eye_pos); + let chunk_size = 4 * 4 * 4 * 4; + if all(in.eye_pos > in.cube_pos) && all(in.eye_pos < (in.cube_pos + vec3(1.))) + { + hit_pos = in.eye_pos - in.cube_pos; + } + else + { + let aabb = intersectAABB(in.eye_pos, dir, in.cube_pos, in.cube_pos + vec3(1.)); + hit_pos = in.eye_pos + aabb.x * dir - in.cube_pos; + } + + let cube_color = vec3(1.); + + var pos = hit_pos * f32(chunk_size); + let step = vec3( + select(-1, 1, dir.x > 0.), + select(-1, 1, dir.y > 0.), + select(-1, 1, dir.z > 0.) + ); + + var voxel = vec3( + clamp(i32(floor(pos.x)), 0, chunk_size - 1), + clamp(i32(floor(pos.y)), 0, chunk_size - 1), + clamp(i32(floor(pos.z)), 0, chunk_size - 1), + ); + + var div = 1; + var overlay = 1.; + for(var i = 1; i <= 4; i++) + { + let x = (voxel.x / div + voxel.y / div + voxel.z / div) % 2 == 0; + overlay -= select(0., 1. / (f32(i) * 2.5), x); + div *= 4; + } + overlay = 1.; + return overlay * traverse(4, in.root_color, in.root_subdivided != 0, in.eye_pos - in.cube_pos, dir); +} + +@fragment +fn fragment_main(in: VertexOutput) -> @location(0) vec4 +{ + let chunk_size = 32; + + let dir = normalize(in.world_pos.xyz - in.eye_pos); + + var hit_pos = vec3(0.); + if all(in.eye_pos > in.cube_pos) && all(in.eye_pos < (in.cube_pos + vec3(1.))) + { + hit_pos = in.eye_pos - in.cube_pos; + } + else + { + let aabb = intersectAABB(in.eye_pos, dir, in.cube_pos, in.cube_pos + vec3(1.)); + hit_pos = in.eye_pos + aabb.x * dir - in.cube_pos; + } + + let cube_color = vec3(1.); + + var pos = hit_pos * f32(chunk_size); + let step = vec3( + select(-1, 1, dir.x > 0.), + select(-1, 1, dir.y > 0.), + select(-1, 1, dir.z > 0.) + ); + + var voxel = vec3( + clamp(i32(floor(pos.x)), 0, chunk_size - 1), + clamp(i32(floor(pos.y)), 0, chunk_size - 1), + clamp(i32(floor(pos.z)), 0, chunk_size - 1), + ); + + let tDelta = vec3(1.) / abs(dir); + var dist = vec3( + select(pos.x - f32(voxel.x), f32(voxel.x) + 1. - pos.x, step.x > 0), + select(pos.y - f32(voxel.y), f32(voxel.y) + 1. - pos.y, step.y > 0), + select(pos.z - f32(voxel.z), f32(voxel.z) + 1. - pos.z, step.z > 0), + ); + var tMax = dist * tDelta; + //var tMax = (ceil(vec3(step) * pos) - vec3(step) * pos) * tDelta; + var t = 0.; + + + // Loop + loop + { + if any(voxel >= vec3(chunk_size)) || any(voxel < vec3(0)) + { + discard; + //break; + } + // Sample + if is_voxel(voxel) + { + // Compute normal + let voxel_center = vec3(voxel) + vec3(0.5); + let pos = pos + t * dir; + + let norm_dir = normalize(pos - voxel_center); + let norm_dir_max = max_mask3 (abs(norm_dir)); + let norm = sign(norm_dir * norm_dir_max); + + let color = (1.2 + dot(norm, vec3(0., 1., 0.))) * 0.5; + return vec4(vec3(color) * cube_color, 1.); + return vec4(vec3(color), 1.); + } + + // Select which to step + let mask = min_mask3(tMax); + let delta = tDelta * mask; + + let next_t_vec = tMax * mask; + t = next_t_vec.x + next_t_vec.y + next_t_vec.z; + + tMax += delta; + + voxel += step * vec3(mask); + + } + + // Ray direction + return vec4(1., 0., 1., 1.); +} + +fn min_mask3(v: vec3) -> vec3 +{ + let min = min(v.x, min(v.y, v.z)); + + return vec3 + ( + select(0., 1., v.x == min), + select(0., 1., v.y == min), + select(0., 1., v.z == min), + ); +} + +fn min_mask3i32(v: vec3) -> vec3 +{ + let min = min(v.x, min(v.y, v.z)); + + return vec3 + ( + select(0, 1, v.x == min), + select(0, 1, v.y == min), + select(0, 1, v.z == min), + ); +} + +fn max_mask3(v: vec3) -> vec3 +{ + let max = max(v.x, max(v.y, v.z)); + + return vec3 + ( + select(0., 1., v.x == max), + select(0., 1., v.y == max), + select(0., 1., v.z == max), + ); +} + +fn intersectAABB(rayOrigin: vec3, rayDir: vec3, boxMin: vec3, boxMax: vec3) -> vec2 { + let tMin = (boxMin - rayOrigin) / rayDir; + let tMax = (boxMax - rayOrigin) / rayDir; + let t1 = min(tMin, tMax); + let t2 = max(tMin, tMax); + let tNear = max(max(t1.x, t1.y), t1.z); + let tFar = min(min(t2.x, t2.y), t2.z); + return vec2(tNear, tFar); +}; diff --git a/src/lib.rs b/src/lib.rs index 84abb7a..2cebb1e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,14 +1,17 @@ +#![allow(incomplete_features)] +#![feature(generic_const_exprs)] + pub mod egui_renderer; pub mod state; +pub mod voxel; use std::sync::Arc; -use winit::{ - application::ApplicationHandler, - event::WindowEvent, - event_loop::{self, EventLoop}, - window::Window, -}; +use winit::application::ApplicationHandler; +use winit::event::WindowEvent; +use winit::event_loop::EventLoop; +use winit::event_loop::{self}; +use winit::window::Window; use crate::state::State; @@ -45,6 +48,10 @@ impl ApplicationHandler for App { self.state = Some(state); window.request_redraw(); + window.set_cursor_visible(false); + window + .set_cursor_grab(winit::window::CursorGrabMode::Locked) + .unwrap(); } fn window_event( @@ -68,6 +75,29 @@ impl ApplicationHandler for App { WindowEvent::Resized(size) => { state.resize(size); } + + WindowEvent::MouseWheel { delta, .. } => { + state.mouse_wheel(delta); + } + + _ => {} + } + } + + fn device_event( + &mut self, + _event_loop: &event_loop::ActiveEventLoop, + _device_id: winit::event::DeviceId, + event: winit::event::DeviceEvent, + ) { + let state = self.state.as_mut().unwrap(); + + #[allow(clippy::single_match)] + match event { + winit::event::DeviceEvent::MouseMotion { delta } => { + state.cursor_moved(delta.0 as f32, delta.1 as f32); + } + _ => {} } } diff --git a/src/state.rs b/src/state.rs index 7b208c5..316b817 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,14 +1,52 @@ +use std::collections::HashSet; +use std::num::NonZero; use std::sync::Arc; +use std::time::Instant; -use cgmath::{Matrix4, Vector3}; +use cgmath::EuclideanSpace; +use cgmath::InnerSpace; +use cgmath::Matrix4; +use cgmath::Point3; +use cgmath::SquareMatrix; +use cgmath::Vector3; +use cgmath::Vector4; use crevice::std430::AsStd430; +use egui::ProgressBar; +use egui::menu::bar; use egui_wgpu::ScreenDescriptor; -use wgpu::{Features, FeaturesWGPU, FeaturesWebGPU}; -use winit::{event::WindowEvent, window::Window}; +use indicatif::ProgressIterator; +use itertools::Itertools; +use wgpu::BindGroup; +use wgpu::BindGroupDescriptor; +use wgpu::BindGroupEntry; +use wgpu::BindGroupLayoutDescriptor; +use wgpu::BindGroupLayoutEntry; +use wgpu::Buffer; +use wgpu::BufferUsages; +use wgpu::Extent3d; +use wgpu::FeaturesWGPU; +use wgpu::FeaturesWebGPU; +use wgpu::PushConstantRange; +use wgpu::RenderPipeline; +use wgpu::ShaderStages; +use wgpu::TextureDescriptor; +use wgpu::TextureFormat; +use wgpu::TextureUsages; +use wgpu::TextureView; +use wgpu::include_wgsl; +use wgpu::util::DeviceExt; +use winit::event::MouseScrollDelta; +use winit::event::WindowEvent; +use winit::keyboard::KeyCode; +use winit::window::Window; use crate::egui_renderer::EguiState; +use crate::voxel::Color; +use crate::voxel::GPUStructureTile; +use crate::voxel::NTree; -pub struct State { +pub struct State +{ window: Arc, device: wgpu::Device, queue: wgpu::Queue, @@ -16,10 +54,51 @@ pub struct State { surface: wgpu::Surface<'static>, surface_format: wgpu::TextureFormat, egui_state: EguiState, + camera: Camera, + + // Pipelines + rm_pipeline: RenderPipeline, + + depth_buffer_view: TextureView, + + // Input + pressed_set: HashSet, + + // Frame time + last_frame: Instant, + + // Data + ntree: NTree<4>, + + // Tree structure + structure_tiles_buf: Buffer, + color_tiles_buf: Buffer, + tree_bind_group: BindGroup, + root_color: Color, } -impl State { - pub async fn new(window: Arc) -> State { +#[derive(Clone, Copy)] +pub struct Camera +{ + pub eye: cgmath::Point3, + pub up: cgmath::Vector3, + pub aspect: f32, + pub fovy: f32, + pub znear: f32, + pub zfar: f32, + + pub radius: f32, + + pub yaw: f32, + pub pitch: f32, + pub rotation_speed: f32, + pub speed: f32, +} + +impl State +{ + pub async fn new(window: Arc) -> State + { let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default()); let adapter = instance .request_adapter(&wgpu::RequestAdapterOptions::default()) @@ -27,13 +106,16 @@ impl State { .unwrap(); let (device, queue) = adapter .request_device(&wgpu::DeviceDescriptor { - required_features: Features { - features_wgpu: FeaturesWGPU::PUSH_CONSTANTS, + required_features: wgpu::Features { + features_wgpu: FeaturesWGPU::PUSH_CONSTANTS + | FeaturesWGPU::BUFFER_BINDING_ARRAY + | FeaturesWGPU::STORAGE_RESOURCE_BINDING_ARRAY, features_webgpu: FeaturesWebGPU::empty(), }, required_limits: wgpu::Limits { max_push_constant_size: RayMarchingPushConstants::std430_size_static() as u32, - ..Default::default() + max_binding_array_elements_per_shader_stage: 2, + ..wgpu::Limits::downlevel_defaults() }, ..Default::default() }) @@ -46,15 +128,229 @@ impl State { let cap = surface.get_capabilities(&adapter); let surface_format = cap.formats[0]; + // Depth buffer + let depth_buffer = device.create_texture(&TextureDescriptor { + label: Some("depth_buffer"), + size: Extent3d { + width: size.width, + height: size.height, + depth_or_array_layers: 1, + }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: TextureFormat::Depth24PlusStencil8, + usage: TextureUsages::RENDER_ATTACHMENT, + view_formats: &[], + }); + + let depth_buffer_view = depth_buffer.create_view(&Default::default()); + + // Rendering pipeline + let cube_shader_module = device.create_shader_module(include_wgsl!("../shaders/cube.wgsl")); + let tree_bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor { + label: Some("ntree_bind_group_layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: NonZero::new(0), + }, + count: NonZero::new(1), + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only: true }, + has_dynamic_offset: false, + min_binding_size: NonZero::new(0), + }, + count: NonZero::new(1), + }, + ], + }); + let rm_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("Render3D Mesh Pipeline Layout"), + bind_group_layouts: &[&tree_bind_group_layout], + push_constant_ranges: &[PushConstantRange { + stages: ShaderStages::VERTEX, + range: 0..RayMarchingPushConstants::std430_size_static() as u32, + }], + }); + + let rm_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { + label: Some("Render3D Mesh Pipeline"), + layout: Some(&rm_pipeline_layout), + vertex: wgpu::VertexState { + module: &cube_shader_module, + entry_point: Some("vertex_main"), + buffers: &[], + compilation_options: wgpu::PipelineCompilationOptions::default(), + }, + fragment: Some(wgpu::FragmentState { + module: &cube_shader_module, + entry_point: Some("fragment_tree_main"), + targets: &[Some(wgpu::ColorTargetState { + format: surface_format, + blend: Some(wgpu::BlendState::REPLACE), + write_mask: wgpu::ColorWrites::ALL, + })], + compilation_options: wgpu::PipelineCompilationOptions::default(), + }), + primitive: wgpu::PrimitiveState { + topology: wgpu::PrimitiveTopology::TriangleList, + cull_mode: Some(wgpu::Face::Front), + ..Default::default() + }, + depth_stencil: Some(wgpu::DepthStencilState { + format: TextureFormat::Depth24PlusStencil8, + depth_write_enabled: true, + depth_compare: wgpu::CompareFunction::Less, + stencil: wgpu::StencilState::default(), + bias: wgpu::DepthBiasState::default(), + }), + multisample: wgpu::MultisampleState { + count: 1, + mask: !0, + alpha_to_coverage_enabled: false, + }, + multiview: None, + cache: None, + }); + + let mut ntree = NTree::<4>::constant(crate::voxel::Color(0., 0., 0., 0.), 4); + + println!("Building tree"); + let width = 100; + for ((x, y), z) in (0..(width / 2)) + .progress() + .cartesian_product(0..width) + .cartesian_product(0..width) + { + let dist = + cgmath::Vector3::new(x as f32 - 128., y as f32 - 128., z as f32 - 128.).magnitude(); + if dist < 128. + { + ntree.set(x, y, z, crate::voxel::Color(0.2, 1., 0.2, 1.)); + } + } + println!("Built tree"); + + // let width = 33; + // for ((x, y), z) in (0..width) + // .cartesian_product(0..width) + // .cartesian_product(0..width) + // { + // if x == 32 || y == 32 || z == 32 + // { + // ntree.set(x, y, z, Color(0., 1., 0., 1.)); + // } + // else + // { + // ntree.set(x, y, z, Color(0., 0., 0., 0.)); + // } + // } + // + // for x in 0..256 + // { + // ntree.set(x, 40, 40, Color(0., 0., 0., 0.)); + // ntree.set(x, 41, 40, Color(1., 0., 0., 1.)); + // ntree.set(x, 39, 40, Color(1., 0., 0., 1.)); + // + // ntree.set(x, 40, 39, Color(0., 0., 1., 1.)); + // ntree.set(x, 40, 41, Color(0., 0., 1., 1.)); + // } + // + // let width = 16; + // for ((x, y), z) in (0..width) + // .cartesian_product(0..width) + // .cartesian_product(0..width) + // { + // ntree.set(x, y, z, Color(0., 0., 0., 0.)); + // } + // ntree.set(4, 0, 0, Color(1., 0., 0., 1.)); + + //ntree.set(0, 0, 0, Color(1., 0., 0., 1.)); + //ntree.set(255, 255, 255, Color(0., 0., 1., 1.)); + println!("Flattening"); + let (root_color, structure_tiles, color_tiles) = ntree.to_gpu_rep(); + println!("Falttened"); + + let structure_tiles_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("structure_tiles_buf"), + contents: unsafe { + std::slice::from_raw_parts( + structure_tiles.as_ptr() as *const u8, + structure_tiles.len() * size_of::>(), + ) + }, + usage: BufferUsages::STORAGE, + }); + let color_tiles_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("color_tiles_buf"), + contents: unsafe { + std::slice::from_raw_parts( + color_tiles.as_ptr() as *const u8, + color_tiles.len() * size_of::<[Color; 4 * 4 * 4]>(), + ) + }, + usage: BufferUsages::STORAGE, + }); + let tree_bind_group = device.create_bind_group(&BindGroupDescriptor { + label: Some("tree_bind_group"), + layout: &tree_bind_group_layout, + entries: &[ + BindGroupEntry { + binding: 0, + resource: structure_tiles_buf.as_entire_binding(), + }, + BindGroupEntry { + binding: 1, + resource: color_tiles_buf.as_entire_binding(), + }, + ], + }); + let state = State { egui_state: EguiState::new(&device, surface_format, &window), window, - device, queue, size, surface, surface_format, + camera: Camera { + eye: Point3::new(0., 0., 0.), + up: Vector3::unit_y(), + aspect: size.width as f32 / size.height as f32, + fovy: 90., + znear: 0.001, + zfar: 1000., + radius: 2., + yaw: 1., + pitch: 0., + rotation_speed: 0.005, + speed: 0.005, + }, + + rm_pipeline, + + depth_buffer_view, + + pressed_set: HashSet::new(), + last_frame: Instant::now(), + + ntree, + + root_color, + device, + structure_tiles_buf, + color_tiles_buf, + tree_bind_group, }; // Configure surface for the first time @@ -63,11 +359,13 @@ impl State { state } - pub fn get_window(&self) -> &Window { + pub fn get_window(&self) -> &Window + { &self.window } - pub fn configure_surface(&self) { + pub fn configure_surface(&self) + { let surface_config = wgpu::SurfaceConfiguration { usage: wgpu::TextureUsages::RENDER_ATTACHMENT, format: self.surface_format, @@ -82,18 +380,58 @@ impl State { self.surface.configure(&self.device, &surface_config); } - pub fn resize(&mut self, new_size: winit::dpi::PhysicalSize) { + pub fn resize(&mut self, new_size: winit::dpi::PhysicalSize) + { self.size = new_size; // reconfigure the surface + let depth_buffer = self.device.create_texture(&TextureDescriptor { + label: Some("depth_buffer"), + size: Extent3d { + width: new_size.width, + height: new_size.height, + depth_or_array_layers: 1, + }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: TextureFormat::Depth24PlusStencil8, + usage: TextureUsages::RENDER_ATTACHMENT, + view_formats: &[], + }); + + self.depth_buffer_view = depth_buffer.create_view(&Default::default()); + + self.camera.aspect = new_size.width as f32 / new_size.height as f32; self.configure_surface(); } - pub fn handle_event(&mut self, event: &WindowEvent) { + pub fn handle_event(&mut self, event: &WindowEvent) + { self.egui_state.handle_event(&self.window, event); + + if let WindowEvent::KeyboardInput { event, .. } = event + { + match (event.state, event.physical_key) + { + (winit::event::ElementState::Pressed, winit::keyboard::PhysicalKey::Code(c)) => + { + self.pressed_set.insert(c); + } + (winit::event::ElementState::Released, winit::keyboard::PhysicalKey::Code(c)) => + { + self.pressed_set.remove(&c); + } + _ => + {} + } + } } - pub fn render(&mut self) { + pub fn render(&mut self) + { + // Update camera + self.update_camera(); // Create texture view let surface_texture = self .surface @@ -111,22 +449,53 @@ impl State { // Renders a GREEN screen let mut encoder = self.device.create_command_encoder(&Default::default()); // Create the renderpass which will clear the screen. - let renderpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + let mut renderpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { label: None, color_attachments: &[Some(wgpu::RenderPassColorAttachment { view: &texture_view, depth_slice: None, resolve_target: None, ops: wgpu::Operations { - load: wgpu::LoadOp::Clear(wgpu::Color::GREEN), + load: wgpu::LoadOp::Clear(wgpu::Color::BLACK), store: wgpu::StoreOp::Store, }, })], - depth_stencil_attachment: None, + depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment { + view: &self.depth_buffer_view, + depth_ops: Some(wgpu::Operations { + load: wgpu::LoadOp::Clear(1.), + store: wgpu::StoreOp::Discard, + }), + stencil_ops: None, + }), timestamp_writes: None, occlusion_query_set: None, }); + renderpass.set_pipeline(&self.rm_pipeline); + + //renderpass.set_vertex_buffer(0, self.positions_buffer.slice(..)); + renderpass.set_bind_group(0, Some(&self.tree_bind_group), &[]); + renderpass.set_push_constants( + ShaderStages::VERTEX, + 0, + RayMarchingPushConstants { + view_projection: self.camera.view_proj(), + transform: Matrix4::identity(), + eye_position: self.camera.eye.to_vec(), + root_color: Vector4::new( + self.root_color.0, + self.root_color.1, + self.root_color.2, + self.root_color.3, + ), + root_subdivided: 1, + } + .as_std430() + .as_bytes(), + ); + renderpass.draw(0..(6 * 2 * 3), 0..1); + // End the renderpass. drop(renderpass); // Egui @@ -140,7 +509,14 @@ impl State { egui::Window::new("Hello Window").resizable(true).show( self.egui_state.context(), |ui| { - ui.label("Hello, world."); + ui.label(format!( + "Frame time: {:.2} ms", + (Instant::now() - self.last_frame).as_secs_f32() * 1000. + )); + ui.label(format!( + "Frame rate: {:.2} fps", + 1. / (Instant::now() - self.last_frame).as_secs_f32() + )); }, ); @@ -158,12 +534,85 @@ impl State { self.queue.submit([encoder.finish()]); self.window.pre_present_notify(); surface_texture.present(); + + self.last_frame = Instant::now(); + } + + fn update_camera(&mut self) + { + let mut movement = cgmath::Vector3::new(0., 0., 0.); + if self.pressed_set.contains(&KeyCode::KeyW) + { + movement.z += self.camera.speed; + } + if self.pressed_set.contains(&KeyCode::KeyS) + { + movement.z -= self.camera.speed; + } + + // Left rigth + if self.pressed_set.contains(&KeyCode::KeyA) + { + movement.x += self.camera.speed; + } + if self.pressed_set.contains(&KeyCode::KeyD) + { + movement.x -= self.camera.speed; + } + + let rot_movement = cgmath::Matrix3::from_angle_y(cgmath::Deg(-self.camera.yaw)) + * cgmath::Matrix3::from_angle_x(cgmath::Deg(-self.camera.pitch)) + * movement; + self.camera.eye -= rot_movement; + } + + pub fn cursor_moved(&mut self, x: f32, y: f32) + { + const SENSIBILITY: f32 = 0.02; + let position = cgmath::Vector2::new(x, y); + let offset = position * SENSIBILITY; + + self.camera.yaw += offset.x; + self.camera.pitch += offset.y; + } + + pub fn mouse_wheel(&mut self, delta: MouseScrollDelta) + { + if let MouseScrollDelta::LineDelta(_, y) = delta + { + self.camera.speed += y * (self.camera.speed * 0.05); + } + + //self.speed += delta. } } #[derive(AsStd430)] -pub struct RayMarchingPushConstants { - inverse_projection_matrix: Matrix4, - view_matrix: Matrix4, - camera_pos: Vector3, +pub struct RayMarchingPushConstants +{ + view_projection: Matrix4, + transform: Matrix4, + eye_position: Vector3, + root_color: Vector4, + root_subdivided: u32, +} + +#[rustfmt::skip] +pub const OPENGL_TO_WGPU_MATRIX: cgmath::Matrix4 = cgmath::Matrix4::from_cols( + cgmath::Vector4::new(1.0, 0.0, 0.0, 0.0), + cgmath::Vector4::new(0.0, 1.0, 0.0, 0.0), + cgmath::Vector4::new(0.0, 0.0, 0.5, 0.0), + cgmath::Vector4::new(0.0, 0.0, 0.5, 1.0), +); + +impl Camera +{ + pub fn view_proj(&self) -> cgmath::Matrix4 + { + let view = cgmath::Matrix4::from_translation(self.eye.to_vec()) + * cgmath::Matrix4::from_angle_y(cgmath::Deg(-self.yaw)) + * cgmath::Matrix4::from_angle_x(cgmath::Deg(-self.pitch)); + let proj = cgmath::perspective(cgmath::Deg(self.fovy), self.aspect, self.znear, self.zfar); + OPENGL_TO_WGPU_MATRIX * proj * view.invert().unwrap() + } } diff --git a/src/voxel.rs b/src/voxel.rs new file mode 100644 index 0000000..3912d9a --- /dev/null +++ b/src/voxel.rs @@ -0,0 +1,548 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::hash::Hash; +use std::vec; + +use itertools::Itertools; + +#[derive(Clone, Copy, PartialEq, Debug)] +pub struct Color(pub f32, pub f32, pub f32, pub f32); + +pub struct NTree +{ + structure: HashMap, NTreeNode>, + depth: u32, +} + +#[derive(Clone, Copy)] +struct NTreeNode +{ + color: Color, + subdivided: bool, +} + +impl NTree +where + [(); N * N * N]:, +{ + pub fn constant(color: Color, depth: u32) -> Self + { + let mut structure = HashMap::new(); + structure.insert( + NTreeLocator::root(), + NTreeNode { + color, + subdivided: false, + }, + ); + Self { structure, depth } + } + + pub fn get(&self, x: usize, y: usize, z: usize) -> Color + { + let mut local_x = x; + let mut local_y = y; + let mut local_z = z; + + let mut size = N.pow(self.depth); + assert!(x < size && y < size && z < size); + + let mut current_loc = NTreeLocator::::root(); + let mut current_node = *self.structure.get(¤t_loc).unwrap(); + + loop + { + if !current_node.subdivided + { + return current_node.color; + } + + size /= N; + + // Descend + let child_x = local_x / size; + let child_y = local_y / size; + let child_z = local_z / size; + + // Node is subdivided, descend + current_loc = current_loc.get_child(child_x, child_y, child_z); + current_node = *self.structure.get(¤t_loc).unwrap(); + + local_x -= child_x * size; + local_y -= child_y * size; + local_z -= child_z * size; + } + } + + pub fn set(&mut self, x: usize, y: usize, z: usize, color: Color) + { + let mut local_x = x; + let mut local_y = y; + let mut local_z = z; + + let mut size = N.pow(self.depth); + assert!(x < size && y < size && z < size); + + let mut current_loc = NTreeLocator::::root(); + let mut current_node = *self.structure.get(¤t_loc).unwrap(); + let mut current_depth = 0; + + loop + { + // Subnodes alread with correct color + if current_node.color == color + { + return; + } + + if current_depth == self.depth + { + current_node.color = color; + self.structure.get_mut(¤t_loc).unwrap().color = color; + break; + } + + if !current_node.subdivided + { + // Have to subdivide + let sub_child = NTreeNode:: { + color: current_node.color, + subdivided: false, + }; + // Set node as subdivided + self.structure.get_mut(¤t_loc).unwrap().subdivided = true; + current_node.subdivided = true; + + // Insert new children + current_loc.iter_children().for_each(|x| { + self.structure.insert(x, sub_child); + }); + } + + if current_node.subdivided + { + size /= N; + let child_x = local_x / size; + let child_y = local_y / size; + let child_z = local_z / size; + + // Node is subdivided, descend + current_loc = current_loc.get_child(child_x, child_y, child_z); + current_node = *self.structure.get(¤t_loc).unwrap(); + current_depth += 1; + + local_x -= child_x * size; + local_y -= child_y * size; + local_z -= child_z * size; + } + } + + // Insertion has been done, travel back up to optimise + loop + { + // Try to simplify, compute average + if current_node.subdivided + { + let (compressable, color) = current_loc + .iter_children() + .map(|x| self.structure.get(&x).unwrap()) + .fold((true, None), |state, child| match state + { + (_, None) => (!child.subdivided, Some(child.color)), + (b, Some(c)) => (b && !child.subdivided && c == child.color, Some(c)), + }); + + if compressable + { + // Compress + *self.structure.get_mut(¤t_loc).unwrap() = NTreeNode { + color: color.unwrap(), + subdivided: false, + }; + + // Remove children + current_loc.iter_children().for_each(|x| { + self.structure.remove(&x).unwrap(); + }); + } + else + { + // Update average color + let colors = current_loc + .iter_children() + .map(|x| self.structure.get(&x).unwrap().color) + .collect::>(); + self.structure.get_mut(¤t_loc).unwrap().color = + Color::average(colors.as_slice()); + } + } + + // Travel up + if current_depth == 0 + { + break; // Finished + } + + current_loc = current_loc.get_parent(); + current_node = *self.structure.get(¤t_loc).unwrap(); + current_depth -= 1; + } + } + + pub fn to_gpu_rep(&self) -> (Color, Vec>, Vec<[Color; N * N * N]>) + { + // No root + group by child group + let tile_count = (self.structure.len() - 1) / (N * N * N); + if tile_count == 0 + { + return ( + self.structure + .get(&NTreeLocator::::root()) + .unwrap() + .color, + vec![], + vec![], + ); + } + + let mut structure_tiles = vec![GPUStructureTile::::zero(); tile_count]; + let mut color_tiles = vec![[Color(0., 0., 0., 0.); N * N * N]; tile_count]; + + let mut current_tile = 0usize; + let mut queue = VecDeque::new(); + queue.push_back((NTreeLocator::::root(), current_tile)); + current_tile += 1; + while !queue.is_empty() + { + let (loc, dest_tile) = queue.pop_front().unwrap(); + + let children = loc + .iter_children() + .map(|x| (x, self.structure.get(&x).unwrap())) + .collect::>(); + + let mut structure_tile = GPUStructureTile::::zero(); + let mut color_tile = [Color(0., 0., 0., 0.); N * N * N]; + + for (i, (child_loc, child)) in children.into_iter().enumerate() + { + structure_tile.children[i] = GPUStructureTileIndex::new( + child.subdivided, + child.subdivided, + current_tile as u32, + ); + color_tile[i] = child.color; + + if child.subdivided + { + queue.push_back((child_loc, current_tile)); + current_tile += 1; + } + } + + structure_tiles[dest_tile] = structure_tile; + color_tiles[dest_tile] = color_tile; + } + + println!( + "node count: {}, current_tile: {}", + self.structure.len(), + current_tile + ); + ( + self.structure + .get(&NTreeLocator::::root()) + .unwrap() + .color, + structure_tiles, + color_tiles, + ) + } +} + +#[derive(Clone, Copy)] +pub struct GPUStructureTile +where + [(); N * N * N]:, +{ + children: [GPUStructureTileIndex; N * N * N], +} + +impl GPUStructureTile +where + [(); N * N * N]:, +{ + pub fn zero() -> Self + { + GPUStructureTile { + children: [GPUStructureTileIndex::zero(); N * N * N], + } + } +} + +#[derive(Clone, Copy)] +pub struct GPUStructureTileIndex(u32); + +impl GPUStructureTileIndex +{ + pub fn zero() -> Self + { + GPUStructureTileIndex(0) + } + + pub fn new(subdivided: bool, leaf: bool, ptr: u32) -> Self + { + assert_eq!((ptr >> 30), 0); + GPUStructureTileIndex(ptr | ((leaf as u32) << 30) | ((subdivided as u32) << 31)) + } + + pub fn index(&self) -> u32 + { + self.0 & 0x3FFFFFFFu32 + } + + pub fn subdivided(&self) -> bool + { + (self.0 >> 31) == 1 + } + + pub fn leaf(&self) -> bool + { + ((self.0 >> 30) & 1) == 1 + } +} + +impl Color +{ + pub fn average(colors: &[Color]) -> Color + { + let sum = colors + .iter() + .copied() + .reduce(|Color(r_a, g_a, b_a, a_a), Color(r_b, g_b, b_b, a_b)| { + Color(r_a + r_b, g_a + g_b, b_a + b_b, a_a + a_b) + }) + .unwrap_or(Color(0., 0., 0., 0.)); + let len = colors.len().max(1) as f32; + + Color(sum.0 / len, sum.1 / len, sum.2 / len, sum.3 / len) + } +} + +#[derive(Clone, Copy, Hash, PartialEq, Eq)] +pub struct NTreeLocator(usize); + +impl NTreeLocator +{ + pub fn root() -> Self + { + Self(1) + } + + pub fn get_child(&self, child_x: usize, child_y: usize, child_z: usize) -> Self + { + assert!(child_x < N && child_y < N && child_z < N); + let mut new_loc = self.0; + + // Shift to left three times + new_loc *= N; + new_loc += child_x; + + new_loc *= N; + new_loc += child_y; + + new_loc *= N; + new_loc += child_z; + + Self(new_loc) + } + + pub fn iter_children(&self) -> impl Iterator> + { + (0..N) + .cartesian_product(0..N) + .cartesian_product(0..N) + .map(|((x, y), z)| self.get_child(x, y, z)) + } + + pub fn get_parent(&self) -> Self + { + if self.0 == 0 + { + *self + } + else + { + Self(self.0 / (N * N * N)) + } + } + + pub fn get_local_location(&self) -> (usize, usize, usize) + { + let mut loc = self.0; + let z = loc % N; + loc /= N; + + let y = loc % N; + loc /= N; + + let x = loc % N; + + (x, y, z) + } + + pub fn is_root(&self) -> bool + { + self.0 == 1 + } +} + +#[cfg(test)] +mod test +{ + + use itertools::Itertools; + use rand::Rng; + + use crate::voxel::Color; + use crate::voxel::NTree; + + #[test] + pub fn constant() + { + let color = Color(0.5, 0.3, 0.5, 1.); + let depth = 5; + const N: usize = 3; + let width = N.pow(depth); + let ntree = NTree::::constant(color, depth); + + for ((x, y), z) in (0..width) + .cartesian_product(0..width) + .cartesian_product(0..width) + { + assert_eq!(ntree.get(x, y, z), color); + } + } + + #[test] + pub fn full_insert() + { + const DEPTH: u32 = 4; + const N: usize = 3; + const WIDTH: usize = N.pow(DEPTH); + let mut ntree = NTree::::constant(Color(1., 0., 0., 0.), DEPTH); + + let mut rng = rand::rng(); + let mut storage = vec![vec![vec![Color(0., 0., 0., 0.); WIDTH]; WIDTH]; WIDTH]; + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + let color = Color(rng.random(), rng.random(), rng.random(), rng.random()); + storage[x][y][z] = color; + ntree.set(x, y, z, color); + } + + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + let color = ntree.get(x, y, z); + assert_eq!(storage[x][y][z], color); + } + println!("Total nodes {}", ntree.structure.len()); + } + + #[test] + pub fn full_insert_const() + { + const DEPTH: u32 = 4; + const N: usize = 3; + const WIDTH: usize = N.pow(DEPTH); + const NEW_COLOR: Color = Color(0., 1., 0., 1.); + let mut ntree = NTree::::constant(Color(1., 0., 1., 0.), DEPTH); + + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + ntree.set(x, y, z, NEW_COLOR); + } + + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + let color = ntree.get(x, y, z); + assert_eq!(NEW_COLOR, color); + } + + println!("Total nodes {}", ntree.structure.len()); + } + + #[test] + pub fn full_insert_quantized() + { + const DEPTH: u32 = 4; + const N: usize = 3; + const WIDTH: usize = N.pow(DEPTH); + let mut ntree = NTree::::constant(Color(0., 0., 0., 0.), DEPTH); + + let mut rng = rand::rng(); + let mut storage = vec![vec![vec![Color(0., 0., 0., 0.); WIDTH]; WIDTH]; WIDTH]; + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + let nbr = rng.random::(); + let mut color = Color(0., 1., 0., 0.); + // Only 1 percent of different color + if nbr > 0.01 + { + color = Color(1., 0., 0., 0.); + } + storage[x][y][z] = color; + ntree.set(x, y, z, color); + } + + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + let color = ntree.get(x, y, z); + assert_eq!(storage[x][y][z], color); + } + println!("Total nodes {}", ntree.structure.len()); + } + + #[test] + pub fn gpu_rep_quantized() + { + const DEPTH: u32 = 4; + const N: usize = 3; + const WIDTH: usize = N.pow(DEPTH); + let mut ntree = NTree::::constant(Color(0., 0., 0., 0.), DEPTH); + + let mut rng = rand::rng(); + let mut storage = vec![vec![vec![Color(0., 0., 0., 0.); WIDTH]; WIDTH]; WIDTH]; + for ((x, y), z) in (0..WIDTH) + .cartesian_product(0..WIDTH) + .cartesian_product(0..WIDTH) + { + let nbr = rng.random::(); + let mut color = Color(0., 1., 0., 0.); + // Only 1 percent of different color + if nbr > 0.01 + { + color = Color(1., 0., 0., 0.); + } + storage[x][y][z] = color; + ntree.set(x, y, z, color); + } + + let (_color, _a, _b) = ntree.to_gpu_rep(); + + drop(_a); + drop(_b); + } +}