notedeck

One damus client to rule them all
git clone git://jb55.com/notedeck
Log | Files | Refs | README | LICENSE

avatar.rs (15383B)


      1 use std::num::NonZeroU64;
      2 
      3 use crate::{Quaternion, Vec3};
      4 use eframe::egui_wgpu::{self, wgpu};
      5 use egui::{Rect, Response};
      6 use rand::Rng;
      7 
      8 pub struct DaveAvatar {
      9     rotation: Quaternion,
     10     rot_dir: Vec3,
     11 }
     12 
     13 // Matrix utilities for perspective projection
     14 fn perspective_matrix(fovy_radians: f32, aspect: f32, near: f32, far: f32) -> [f32; 16] {
     15     let f = 1.0 / (fovy_radians / 2.0).tan();
     16     let nf = 1.0 / (near - far);
     17 
     18     // Column-major for WGPU
     19     [
     20         f / aspect,
     21         0.0,
     22         0.0,
     23         0.0,
     24         0.0,
     25         f,
     26         0.0,
     27         0.0,
     28         0.0,
     29         0.0,
     30         (far + near) * nf,
     31         -1.0,
     32         0.0,
     33         0.0,
     34         2.0 * far * near * nf,
     35         0.0,
     36     ]
     37 }
     38 
     39 // Combine two 4x4 matrices (column-major)
     40 fn matrix_multiply(a: &[f32; 16], b: &[f32; 16]) -> [f32; 16] {
     41     let mut result = [0.0; 16];
     42 
     43     for row in 0..4 {
     44         for col in 0..4 {
     45             let mut sum = 0.0;
     46             for i in 0..4 {
     47                 sum += a[row + i * 4] * b[i + col * 4];
     48             }
     49             result[row + col * 4] = sum;
     50         }
     51     }
     52 
     53     result
     54 }
     55 
     56 impl DaveAvatar {
     57     pub fn new(wgpu_render_state: &egui_wgpu::RenderState) -> Self {
     58         let device = &wgpu_render_state.device;
     59 
     60         // Create shader module with improved shader code
     61         let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
     62             label: Some("cube_shader"),
     63             source: wgpu::ShaderSource::Wgsl(
     64                 r#"
     65 struct Uniforms {
     66     model_view_proj: mat4x4<f32>,
     67     model: mat4x4<f32>,    // Added model matrix for correct normal transformation
     68 };
     69 
     70 @group(0) @binding(0)
     71 var<uniform> uniforms: Uniforms;
     72 
     73 struct VertexOutput {
     74     @builtin(position) position: vec4<f32>,
     75     @location(0) normal: vec3<f32>,
     76     @location(1) world_pos: vec3<f32>,
     77 };
     78 
     79 @vertex
     80 fn vs_main(@builtin(vertex_index) vertex_index: u32) -> VertexOutput {
     81     // Define cube vertices (-0.5 to 0.5 in each dimension)
     82     var positions = array<vec3<f32>, 8>(
     83         vec3<f32>(-0.5, -0.5, -0.5),  // 0: left bottom back
     84         vec3<f32>(0.5, -0.5, -0.5),   // 1: right bottom back
     85         vec3<f32>(-0.5, 0.5, -0.5),   // 2: left top back
     86         vec3<f32>(0.5, 0.5, -0.5),    // 3: right top back
     87         vec3<f32>(-0.5, -0.5, 0.5),   // 4: left bottom front
     88         vec3<f32>(0.5, -0.5, 0.5),    // 5: right bottom front
     89         vec3<f32>(-0.5, 0.5, 0.5),    // 6: left top front
     90         vec3<f32>(0.5, 0.5, 0.5)      // 7: right top front
     91     );
     92     
     93     // Define indices for the 12 triangles (6 faces * 2 triangles)
     94     var indices = array<u32, 36>(
     95         // back face (Z-)
     96         0, 2, 1, 1, 2, 3,
     97         // front face (Z+)
     98         4, 5, 6, 5, 7, 6,
     99         // left face (X-)
    100         0, 4, 2, 2, 4, 6,
    101         // right face (X+)
    102         1, 3, 5, 3, 7, 5,
    103         // bottom face (Y-)
    104         0, 1, 4, 1, 5, 4,
    105         // top face (Y+)
    106         2, 6, 3, 3, 6, 7
    107     );
    108     
    109     // Define normals for each face
    110     var face_normals = array<vec3<f32>, 6>(
    111         vec3<f32>(0.0, 0.0, -1.0),  // back face (Z-)
    112         vec3<f32>(0.0, 0.0, 1.0),   // front face (Z+)
    113         vec3<f32>(-1.0, 0.0, 0.0),  // left face (X-)
    114         vec3<f32>(1.0, 0.0, 0.0),   // right face (X+)
    115         vec3<f32>(0.0, -1.0, 0.0),  // bottom face (Y-)
    116         vec3<f32>(0.0, 1.0, 0.0)    // top face (Y+)
    117     );
    118 
    119     var output: VertexOutput;
    120     
    121     // Get vertex from indices
    122     let index = indices[vertex_index];
    123     let position = positions[index];
    124     
    125     // Determine which face this vertex belongs to
    126     let face_index = vertex_index / 6u;
    127     
    128     // Apply transformations
    129     output.position = uniforms.model_view_proj * vec4<f32>(position, 1.0);
    130     
    131     // Transform normal to world space
    132     // Extract the 3x3 rotation part from the 4x4 model matrix
    133     let normal_matrix = mat3x3<f32>(
    134         uniforms.model[0].xyz,
    135         uniforms.model[1].xyz,
    136         uniforms.model[2].xyz
    137     );
    138     output.normal = normalize(normal_matrix * face_normals[face_index]);
    139     
    140     // Pass world position for lighting calculations
    141     output.world_pos = (uniforms.model * vec4<f32>(position, 1.0)).xyz;
    142     
    143     return output;
    144 }
    145 
    146 @fragment
    147 fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    148     // Material properties
    149     let material_color = vec3<f32>(1.0, 1.0, 1.0);  // White color
    150     let ambient_strength = 0.2;
    151     let diffuse_strength = 0.7;
    152     let specular_strength = 0.2;
    153     let shininess = 20.0;
    154     
    155     // Light properties
    156     let light_pos = vec3<f32>(2.0, 2.0, 2.0);  // Light positioned diagonally above and to the right
    157     let light_color = vec3<f32>(1.0, 1.0, 1.0); // White light
    158     
    159     // View position (camera)
    160     let view_pos = vec3<f32>(0.0, 0.0, 3.0);   // Camera position
    161     
    162     // Calculate ambient lighting
    163     let ambient = ambient_strength * light_color;
    164     
    165     // Calculate diffuse lighting
    166     let normal = normalize(in.normal);  // Renormalize the interpolated normal
    167     let light_dir = normalize(light_pos - in.world_pos);
    168     let diff = max(dot(normal, light_dir), 0.0);
    169     let diffuse = diffuse_strength * diff * light_color;
    170     
    171     // Calculate specular lighting
    172     let view_dir = normalize(view_pos - in.world_pos);
    173     let reflect_dir = reflect(-light_dir, normal);
    174     let spec = pow(max(dot(view_dir, reflect_dir), 0.0), shininess);
    175     let specular = specular_strength * spec * light_color;
    176     
    177     // Combine lighting components
    178     let result = (ambient + diffuse + specular) * material_color;
    179     
    180     return vec4<f32>(result, 1.0);
    181 }
    182 "#
    183                 .into(),
    184             ),
    185         });
    186 
    187         // Create uniform buffer for MVP matrix and model matrix
    188         let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
    189             label: Some("cube_uniform_buffer"),
    190             size: 128, // Two 4x4 matrices of f32 (2 * 16 * 4 bytes)
    191             usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    192             mapped_at_creation: false,
    193         });
    194 
    195         // Create bind group layout
    196         let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    197             label: Some("cube_bind_group_layout"),
    198             entries: &[wgpu::BindGroupLayoutEntry {
    199                 binding: 0,
    200                 visibility: wgpu::ShaderStages::VERTEX,
    201                 ty: wgpu::BindingType::Buffer {
    202                     ty: wgpu::BufferBindingType::Uniform,
    203                     has_dynamic_offset: false,
    204                     min_binding_size: NonZeroU64::new(128),
    205                 },
    206                 count: None,
    207             }],
    208         });
    209 
    210         // Create bind group
    211         let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
    212             label: Some("cube_bind_group"),
    213             layout: &bind_group_layout,
    214             entries: &[wgpu::BindGroupEntry {
    215                 binding: 0,
    216                 resource: uniform_buffer.as_entire_binding(),
    217             }],
    218         });
    219 
    220         // Create pipeline layout
    221         let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    222             label: Some("cube_pipeline_layout"),
    223             bind_group_layouts: &[&bind_group_layout],
    224             push_constant_ranges: &[],
    225         });
    226 
    227         // Create render pipeline
    228         let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
    229             label: Some("cube_pipeline"),
    230             layout: Some(&pipeline_layout),
    231             vertex: wgpu::VertexState {
    232                 module: &shader,
    233                 entry_point: Some("vs_main"),
    234                 buffers: &[], // No vertex buffer - vertices are in the shader
    235                 compilation_options: wgpu::PipelineCompilationOptions::default(),
    236             },
    237             fragment: Some(wgpu::FragmentState {
    238                 module: &shader,
    239                 entry_point: Some("fs_main"),
    240                 targets: &[Some(wgpu::ColorTargetState {
    241                     format: wgpu_render_state.target_format,
    242                     blend: Some(wgpu::BlendState::ALPHA_BLENDING),
    243                     write_mask: wgpu::ColorWrites::ALL,
    244                 })],
    245                 compilation_options: wgpu::PipelineCompilationOptions::default(),
    246             }),
    247             primitive: wgpu::PrimitiveState {
    248                 topology: wgpu::PrimitiveTopology::TriangleList,
    249                 strip_index_format: None,
    250                 front_face: wgpu::FrontFace::Ccw,
    251                 cull_mode: Some(wgpu::Face::Back),
    252                 polygon_mode: wgpu::PolygonMode::Fill,
    253                 unclipped_depth: false,
    254                 conservative: false,
    255             },
    256             depth_stencil: Some(wgpu::DepthStencilState {
    257                 format: wgpu::TextureFormat::Depth24Plus,
    258                 depth_write_enabled: true,
    259                 depth_compare: wgpu::CompareFunction::Less,
    260                 stencil: wgpu::StencilState::default(),
    261                 bias: wgpu::DepthBiasState::default(),
    262             }),
    263             multisample: wgpu::MultisampleState::default(),
    264             multiview: None,
    265             cache: None,
    266         });
    267 
    268         // Store resources in renderer
    269         wgpu_render_state
    270             .renderer
    271             .write()
    272             .callback_resources
    273             .insert(CubeRenderResources {
    274                 pipeline,
    275                 bind_group,
    276                 uniform_buffer,
    277             });
    278 
    279         let initial_rot = {
    280             let x_rotation = Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), 0.5);
    281             let y_rotation = Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), 0.5);
    282 
    283             // Apply rotations (order matters)
    284             y_rotation.multiply(&x_rotation)
    285         };
    286         Self {
    287             rotation: initial_rot,
    288             rot_dir: Vec3::new(0.0, 0.0, 0.0),
    289         }
    290     }
    291 }
    292 
    293 #[inline]
    294 fn apply_friction(val: f32, friction: f32, clamp: f32) -> f32 {
    295     if val < clamp {
    296         0.0
    297     } else {
    298         val * friction
    299     }
    300 }
    301 
    302 impl DaveAvatar {
    303     pub fn random_nudge(&mut self) {
    304         self.random_nudge_with(1.0);
    305     }
    306 
    307     pub fn random_nudge_with(&mut self, force: f32) {
    308         let mut rng = rand::rng();
    309 
    310         let nudge = Vec3::new(
    311             rng.random::<f32>() * force,
    312             rng.random::<f32>() * force,
    313             rng.random::<f32>() * force,
    314         )
    315         .normalize();
    316 
    317         self.rot_dir.x += nudge.x;
    318         self.rot_dir.y += nudge.y;
    319         self.rot_dir.z += nudge.z;
    320     }
    321 
    322     pub fn render(&mut self, rect: Rect, ui: &mut egui::Ui) -> Response {
    323         let response = ui.allocate_rect(rect, egui::Sense::CLICK | egui::Sense::DRAG);
    324 
    325         // Update rotation based on drag or animation
    326         if response.dragged() {
    327             // Create rotation quaternions based on drag
    328             let dx = response.drag_delta().x;
    329             let dy = response.drag_delta().y;
    330             let x_rotation = Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), dy * 0.01);
    331             let y_rotation = Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), dx * 0.01);
    332 
    333             self.rot_dir = Vec3::new(dx, dy, 0.0);
    334 
    335             // Apply rotations (order matters)
    336             self.rotation = y_rotation.multiply(&x_rotation).multiply(&self.rotation);
    337         } else if response.clicked() {
    338             self.random_nudge_with(1.0);
    339         } else {
    340             // Continuous rotation - reduced speed and simplified axis
    341             let friction = 0.95;
    342             let clamp = 0.1;
    343             self.rot_dir.x = apply_friction(self.rot_dir.x, friction, clamp);
    344             self.rot_dir.y = apply_friction(self.rot_dir.y, friction, clamp);
    345             self.rot_dir.z = apply_friction(self.rot_dir.y, friction, clamp);
    346 
    347             // we only need to render if we're still spinning
    348             if self.rot_dir.x > clamp || self.rot_dir.y > clamp || self.rot_dir.z > clamp {
    349                 let x_rotation =
    350                     Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), self.rot_dir.y * 0.03);
    351                 let y_rotation =
    352                     Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), self.rot_dir.x * 0.03);
    353                 let z_rotation =
    354                     Quaternion::from_axis_angle(&Vec3::new(0.0, 0.0, 1.0), self.rot_dir.z * 0.03);
    355 
    356                 self.rotation = y_rotation
    357                     .multiply(&x_rotation)
    358                     .multiply(&z_rotation)
    359                     .multiply(&self.rotation);
    360 
    361                 tracing::trace!("repainting due to avatar rotation");
    362                 ui.ctx().request_repaint();
    363             }
    364         }
    365 
    366         // Create model matrix from rotation quaternion
    367         let model_matrix = self.rotation.to_matrix4();
    368 
    369         // Create projection matrix with proper depth range
    370         // Adjust aspect ratio based on rect dimensions
    371         let aspect = rect.width() / rect.height();
    372         let projection = perspective_matrix(std::f32::consts::PI / 4.0, aspect, 0.1, 100.0);
    373 
    374         // Create view matrix (move camera back a bit)
    375         let view_matrix = [
    376             1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, -3.0, 1.0,
    377         ];
    378 
    379         // Combine matrices: projection * view * model
    380         let mv_matrix = matrix_multiply(&view_matrix, &model_matrix);
    381         let mvp_matrix = matrix_multiply(&projection, &mv_matrix);
    382 
    383         // Add paint callback
    384         ui.painter().add(egui_wgpu::Callback::new_paint_callback(
    385             rect,
    386             CubeCallback {
    387                 mvp_matrix,
    388                 model_matrix,
    389             },
    390         ));
    391 
    392         response
    393     }
    394 }
    395 
    396 // Callback implementation
    397 struct CubeCallback {
    398     mvp_matrix: [f32; 16],   // Model-View-Projection matrix
    399     model_matrix: [f32; 16], // Model matrix for lighting calculations
    400 }
    401 
    402 impl egui_wgpu::CallbackTrait for CubeCallback {
    403     fn prepare(
    404         &self,
    405         _device: &wgpu::Device,
    406         queue: &wgpu::Queue,
    407         _screen_descriptor: &egui_wgpu::ScreenDescriptor,
    408         _egui_encoder: &mut wgpu::CommandEncoder,
    409         resources: &mut egui_wgpu::CallbackResources,
    410     ) -> Vec<wgpu::CommandBuffer> {
    411         let resources: &CubeRenderResources = resources.get().unwrap();
    412 
    413         // Create a combined uniform buffer with both matrices
    414         let mut uniform_data = [0.0f32; 32]; // Space for two 4x4 matrices
    415 
    416         // Copy MVP matrix to first 16 floats
    417         uniform_data[0..16].copy_from_slice(&self.mvp_matrix);
    418 
    419         // Copy model matrix to next 16 floats
    420         uniform_data[16..32].copy_from_slice(&self.model_matrix);
    421 
    422         // Update uniform buffer with both matrices
    423         queue.write_buffer(
    424             &resources.uniform_buffer,
    425             0,
    426             bytemuck::cast_slice(&uniform_data),
    427         );
    428 
    429         Vec::new()
    430     }
    431 
    432     fn paint(
    433         &self,
    434         _info: egui::PaintCallbackInfo,
    435         render_pass: &mut wgpu::RenderPass,
    436         resources: &egui_wgpu::CallbackResources,
    437     ) {
    438         let resources: &CubeRenderResources = resources.get().unwrap();
    439 
    440         render_pass.set_pipeline(&resources.pipeline);
    441         render_pass.set_bind_group(0, &resources.bind_group, &[]);
    442         render_pass.draw(0..36, 0..1); // 36 vertices for a cube (6 faces * 2 triangles * 3 vertices)
    443     }
    444 }
    445 
    446 // Simple resources struct
    447 struct CubeRenderResources {
    448     pipeline: wgpu::RenderPipeline,
    449     bind_group: wgpu::BindGroup,
    450     uniform_buffer: wgpu::Buffer,
    451 }