Add raygen shader

Cargo.toml

@@ -7,7 +7,9 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.82"
+bytemuck = { version = "1.15.0", features = ["derive"] }
 crawl = { path = "../crawl" }
 env_logger = "0.11.3"
+glam = "0.27.0"
 log = "0.4.21"
 pollster = "0.3.0"

res/common/raygen.wgsl

@@ -0,0 +1,43 @@
+@group(0) @binding(0) var output: texture_storage_2d<rgba8unorm, write>;
+@group(0) @binding(1) var<uniform> camera: Camera;
+@group(0) @binding(2) var<storage, read_write> rays: array<Ray>;
+
+struct Camera {
+    projection: mat4x4<f32>,
+    view: mat4x4<f32>,
+    pos: vec3<f32>,
+};
+
+struct Ray {
+    pos: vec3<f32>,
+    dist: f32,
+    dir: vec3<f32>,
+    hit: u32,
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn compute(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let img_coord = global_id.xy;
+    let img_dims = textureDimensions(output);
+
+    // This discards the extra pixels in cases where the image size isn't perfectly divisible by the kernel.xy
+    if (img_coord.x >= img_dims.x || img_coord.y >= img_dims.y) {
+        return;
+    }
+
+    // Construct ray
+    let img_coord_frac = vec2<f32>(img_coord) / vec2<f32>(img_dims);
+    let screen_pos = img_coord_frac * 2.0 - vec2<f32>(1.0);
+    var ray_eye = camera.projection * vec4<f32>(screen_pos, -1.0, 0.0);
+    ray_eye = vec4<f32>(ray_eye.xy, -1.0, 0.0);
+    let ray_dir = normalize((camera.view * ray_eye).xyz);
+    let ray_pos = camera.pos;
+
+    let idx = img_coord.x + img_coord.y * img_dims.x;
+    let ray = Ray(ray_pos, 0.0, ray_dir, 0);
+    rays[idx] = ray;
+
+    // Temp
+    var color = vec4<f32>(f32(img_coord.x) / f32(img_dims.x), f32(img_coord.y) / f32(img_dims.y), 0.5, 1.0);
+    textureStore(output, img_coord, color);
+}

src/camera.rs

@@ -0,0 +1,260 @@
+use crawl::{
+    wgpu::{self, util::DeviceExt},
+    winit::{
+        event::{ElementState, KeyEvent, WindowEvent},
+        keyboard::{KeyCode, PhysicalKey},
+    },
+    Context,
+};
+use std::time::Duration;
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+pub struct CameraUniform {
+    projection: [[f32; 4]; 4],
+    view: [[f32; 4]; 4],
+    pos: [f32; 3],
+    _pad: f32,
+}
+
+impl Default for CameraUniform {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl CameraUniform {
+    pub fn new() -> Self {
+        Self {
+            projection: glam::Mat4::IDENTITY.to_cols_array_2d(),
+            view: glam::Mat4::IDENTITY.to_cols_array_2d(),
+            pos: glam::Vec3::ZERO.to_array(),
+            _pad: 0.0,
+        }
+    }
+
+    pub fn update(&mut self, view: glam::Mat4, projection: glam::Mat4, pos: glam::Vec3) {
+        self.view = view.to_cols_array_2d();
+        self.projection = projection.to_cols_array_2d();
+        self.pos = pos.to_array();
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Camera {
+    pub position: glam::Vec3,
+    pub yaw: f32,
+    pub pitch: f32,
+}
+
+impl Camera {
+    pub fn new(position: glam::Vec3, yaw: f32, pitch: f32) -> Self {
+        Self {
+            position,
+            yaw,
+            pitch,
+        }
+    }
+
+    pub fn get_view_matrix(&self) -> glam::Mat4 {
+        glam::Mat4::look_to_rh(
+            self.position,
+            glam::vec3(
+                self.pitch.cos() * self.yaw.cos(),
+                self.pitch.sin(),
+                self.pitch.cos() * self.yaw.sin(),
+            )
+            .normalize(),
+            glam::Vec3::Y,
+        )
+        .transpose()
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Projection {
+    aspect: f32,
+    fov_y: f32,
+    z_near: f32,
+    z_far: f32,
+}
+
+impl Projection {
+    pub fn new(width: u32, height: u32, fov_y: f32, z_near: f32, z_far: f32) -> Self {
+        let aspect = height as f32 / width as f32;
+        Self {
+            aspect,
+            fov_y,
+            z_near,
+            z_far,
+        }
+    }
+
+    pub fn resize(&mut self, width: u32, height: u32) {
+        self.aspect = height as f32 / width as f32;
+    }
+
+    pub fn get_matrix(&self) -> glam::Mat4 {
+        glam::Mat4::perspective_rh(self.fov_y, self.aspect, self.z_near, self.z_far).transpose()
+    }
+}
+
+#[derive(Debug)]
+pub struct CameraController {
+    camera: Camera,
+    projection: Projection,
+    uniform: CameraUniform,
+    buffer: wgpu::Buffer,
+    move_speed: f32,
+    mouse_sensitivity: f32,
+    move_dirs_pressed: glam::IVec3,
+    rot_dirs_pressed: glam::IVec2,
+}
+
+impl CameraController {
+    pub fn new(
+        context: &Context,
+        camera: Camera,
+        projection: Projection,
+        move_speed: f32,
+        mouse_sensitivity: f32,
+    ) -> Self {
+        let mut uniform = CameraUniform::new();
+        uniform.update(
+            camera.get_view_matrix(),
+            projection.get_matrix(),
+            camera.position,
+        );
+
+        let buffer = context
+            .device
+            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
+                label: None,
+                contents: bytemuck::cast_slice(&[uniform]),
+                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+            });
+
+        Self {
+            camera,
+            projection,
+            uniform,
+            buffer,
+            move_speed,
+            mouse_sensitivity,
+            move_dirs_pressed: glam::ivec3(0, 0, 0),
+            rot_dirs_pressed: glam::ivec2(0, 0),
+        }
+    }
+
+    pub fn process_events(&mut self, event: &WindowEvent) -> bool {
+        let mut handled = true;
+        match event {
+            WindowEvent::Resized(physical_size) => {
+                self.projection
+                    .resize(physical_size.width, physical_size.height);
+            }
+            WindowEvent::KeyboardInput {
+                event:
+                    KeyEvent {
+                        state,
+                        physical_key: PhysicalKey::Code(keycode),
+                        ..
+                    },
+                ..
+            } => {
+                let val = match state {
+                    ElementState::Pressed => 1,
+                    ElementState::Released => 0,
+                };
+
+                match keycode {
+                    KeyCode::KeyW => {
+                        self.move_dirs_pressed.z = val;
+                    }
+                    KeyCode::KeyS => {
+                        self.move_dirs_pressed.z = -val;
+                    }
+                    KeyCode::KeyA => {
+                        self.move_dirs_pressed.x = -val;
+                    }
+                    KeyCode::KeyD => {
+                        self.move_dirs_pressed.x = val;
+                    }
+                    KeyCode::KeyQ => {
+                        self.move_dirs_pressed.y = val;
+                    }
+                    KeyCode::KeyE => {
+                        self.move_dirs_pressed.y = -val;
+                    }
+                    KeyCode::ArrowUp => {
+                        self.rot_dirs_pressed.y = val;
+                    }
+                    KeyCode::ArrowDown => {
+                        self.rot_dirs_pressed.y = -val;
+                    }
+                    KeyCode::ArrowLeft => {
+                        self.rot_dirs_pressed.x = -val;
+                    }
+                    KeyCode::ArrowRight => {
+                        self.rot_dirs_pressed.x = val;
+                    }
+                    _ => handled = false,
+                }
+            }
+            _ => handled = false,
+        }
+
+        handled
+    }
+
+    pub fn update(&mut self, dt: Duration) {
+        let dt = dt.as_secs_f32();
+
+        // Calculate look vectors
+        let pitch = self.camera.pitch;
+        let yaw = self.camera.yaw;
+        let front = glam::vec3(
+            pitch.cos() * yaw.cos(),
+            pitch.sin(),
+            pitch.cos() * yaw.sin(),
+        )
+        .normalize();
+        let right = front.cross(glam::Vec3::Y).normalize();
+        let up = right.cross(front).normalize();
+
+        // Apply movement
+        let ms = self.move_speed * dt;
+        self.camera.position += front * ms * self.move_dirs_pressed.z as f32;
+        self.camera.position += right * ms * self.move_dirs_pressed.x as f32;
+        self.camera.position += up * ms * self.move_dirs_pressed.y as f32;
+
+        // Apply rotation
+        let cam_ms = (self.move_speed * self.move_speed).to_radians() * dt;
+        let max_pitch = 85_f32.to_radians();
+        self.camera.yaw += cam_ms * self.rot_dirs_pressed.x as f32;
+        self.camera.pitch += cam_ms * self.rot_dirs_pressed.y as f32;
+        self.camera.pitch = self.camera.pitch.clamp(-max_pitch, max_pitch);
+
+        // Debug log
+        // log::info!("Camera Front: {:?}", front);
+        // log::info!("Move Speed: {:?} {:?} {:?}", self.move_speed, ms, dt);
+        // log::info!("Camera Position: {:?}", self.camera.position);
+        // log::info!("Camera Yaw: {:?}", self.camera.yaw);
+        // log::info!("Camera Pitch: {:?}", self.camera.pitch);
+    }
+
+    pub fn update_buffer(&mut self, context: &Context) {
+        self.uniform.update(
+            self.camera.get_view_matrix(),
+            self.projection.get_matrix(),
+            self.camera.position,
+        );
+        context
+            .queue
+            .write_buffer(&self.buffer, 0, bytemuck::cast_slice(&[self.uniform]));
+    }
+
+    pub fn get_buffer(&self) -> &wgpu::Buffer {
+        &self.buffer
+    }
+}

src/main.rs

@@ -1,3 +1,4 @@
+mod camera;
 mod renderer;
 
 use anyhow::Result;
@@ -12,7 +13,30 @@ fn main() -> Result<()> {
             .with_vsync(false)
             .build(),
     )?;
-    let renderer = VoxelRenderer::new(&context)?;
+
+    let mut camera_controller = camera::CameraController::new(
+        &context,
+        camera::Camera::new(
+            glam::Vec3 {
+                x: 4.01,
+                y: 4.01,
+                z: 20.0,
+            },
+            -90.0_f32.to_radians(),
+            0.0_f32.to_radians(),
+        ),
+        camera::Projection::new(
+            context.size.width,
+            context.size.height,
+            90.0_f32.to_radians(),
+            0.01,
+            100.0,
+        ),
+        10.0,
+        0.25,
+    );
+
+    let renderer = VoxelRenderer::new(&context, &camera_controller)?;
     let passes: Vec<Box<dyn crawl::Pass>> = vec![Box::new(renderer)];
 
     event_loop.run(|event, elwt| match event {
@@ -21,6 +45,10 @@ fn main() -> Result<()> {
                 return;
             }
 
+            if camera_controller.process_events(&event) {
+                return;
+            }
+
             if let WindowEvent::RedrawRequested = event {
                 if let Err(err) = context.render(passes.as_slice()) {
                     log::error!("{err:?}");

src/renderer.rs

@@ -1,14 +1,27 @@
 use anyhow::Result;
 use crawl::wgpu;
 
+use crate::camera::CameraController;
+
+#[repr(C)]
+#[derive(Debug, Default, Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+pub struct Ray {
+    pos: [f32; 3],
+    dist: f32,
+    dir: [f32; 3],
+    hit: u32,
+}
+
 pub struct VoxelRenderer {
     render_texture: crawl::Texture,
     render_pipeline: wgpu::RenderPipeline,
+    raygen_bind_group: wgpu::BindGroup,
+    raygen_pipeline: wgpu::ComputePipeline,
 }
 
 impl VoxelRenderer {
-    pub fn new(context: &crawl::Context) -> Result<Self> {
-        let shader = context
+    pub fn new(context: &crawl::Context, camera: &CameraController) -> Result<Self> {
+        let render_shader = context
             .device
             .create_shader_module(wgpu::include_wgsl!(concat!(
                 env!("CARGO_MANIFEST_DIR"),
@@ -26,20 +39,68 @@ impl VoxelRenderer {
             .with_shader_visibility(wgpu::ShaderStages::FRAGMENT | wgpu::ShaderStages::COMPUTE)
             .build(context)?;
 
-        let render_pipeline = crawl::RenderPipelineBuilder::new("Raycast Quad", &shader)
+        let render_pipeline = crawl::RenderPipelineBuilder::new("Raycast Quad", &render_shader)
             .with_layout("Draw", &[&render_texture.bind_group_layout], &[])
             .with_fragment_targets(&[Some(context.surface_config.format.into())])
             .build(context);
 
+        let raygen_shader = context
+            .device
+            .create_shader_module(wgpu::include_wgsl!(concat!(
+                env!("CARGO_MANIFEST_DIR"),
+                "/res/common/raygen.wgsl"
+            )));
+
+        let ray_data = vec![Ray::default(); (context.size.width * context.size.height) as usize];
+        let mut raygen_buffers = crawl::BulkBufferBuilder::new()
+            .set_usage(wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST)
+            .with_init_buffer_bm("Rays", &ray_data)
+            .build(context);
+        let raygen_layout = crawl::BindGroupLayoutBuilder::new()
+            .with_label("Raygen BGL")
+            .with_entry(
+                wgpu::ShaderStages::COMPUTE,
+                wgpu::BindingType::StorageTexture {
+                    access: wgpu::StorageTextureAccess::WriteOnly,
+                    format: render_texture.attributes.format,
+                    view_dimension: wgpu::TextureViewDimension::D2,
+                },
+                None,
+            )
+            .with_uniform_entry(wgpu::ShaderStages::COMPUTE)
+            .with_rw_storage_entry(wgpu::ShaderStages::COMPUTE)
+            .build(context);
+
+        let raygen_bind_group = crawl::BindGroupBuilder::new()
+            .with_label("Raygen BG")
+            .with_layout(&raygen_layout)
+            .with_entry(wgpu::BindingResource::TextureView(&render_texture.view))
+            .with_entry(camera.get_buffer().as_entire_binding())
+            .with_entry(raygen_buffers.remove(0).as_entire_binding())
+            .build(context)?;
+
+        let raygen_pipeline = crawl::ComputePipelineBuilder::new("Raygen P", &raygen_shader)
+            .with_layout("Raygen PL", &[&raygen_layout], &[])
+            .build(context);
+
         Ok(Self {
             render_texture,
             render_pipeline,
+            raygen_pipeline,
+            raygen_bind_group,
         })
     }
 }
 
 impl crawl::Pass for VoxelRenderer {
     fn execute(&self, encoder: &mut wgpu::CommandEncoder, view: &wgpu::TextureView) {
+        let size = self.render_texture.attributes.size;
+        let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
+        compute_pass.set_pipeline(&self.raygen_pipeline);
+        compute_pass.set_bind_group(0, &self.raygen_bind_group, &[]);
+        compute_pass.dispatch_workgroups(size.width / 8, size.height / 8, 1);
+        drop(compute_pass);
+
         let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
             label: Some("Render Pass"),
             color_attachments: &[Some(wgpu::RenderPassColorAttachment {
@@ -52,5 +113,6 @@ impl crawl::Pass for VoxelRenderer {
         render_pass.set_pipeline(&self.render_pipeline);
         render_pass.set_bind_group(0, &self.render_texture.bind_group, &[]);
         render_pass.draw(0..6, 0..1);
+        drop(render_pass);
     }
 }