notedeck

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

avatar.rs (14669B)


      1 use std::num::NonZeroU64;
      2 
      3 use crate::mesh;
      4 use crate::{Quaternion, Vec3};
      5 use eframe::egui_wgpu::{
      6     self,
      7     wgpu::{self, util::DeviceExt},
      8 };
      9 use egui::{Rect, Response};
     10 use rand::Rng;
     11 use std::borrow::Cow;
     12 
     13 pub struct DaveAvatar {
     14     rotation: Quaternion,
     15     rot_dir: Vec3,
     16     logical_time: f32,
     17 }
     18 
     19 // Matrix utilities for perspective projection
     20 fn perspective_matrix(fovy_radians: f32, aspect: f32, near: f32, far: f32) -> [f32; 16] {
     21     let f = 1.0 / (fovy_radians / 2.0).tan();
     22     let nf = 1.0 / (near - far);
     23 
     24     // Column-major for WGPU
     25     [
     26         f / aspect,
     27         0.0,
     28         0.0,
     29         0.0,
     30         0.0,
     31         f,
     32         0.0,
     33         0.0,
     34         0.0,
     35         0.0,
     36         (far + near) * nf,
     37         -1.0,
     38         0.0,
     39         0.0,
     40         2.0 * far * near * nf,
     41         0.0,
     42     ]
     43 }
     44 
     45 // Combine two 4x4 matrices (column-major)
     46 fn matrix_multiply(a: &[f32; 16], b: &[f32; 16]) -> [f32; 16] {
     47     let mut result = [0.0; 16];
     48 
     49     for row in 0..4 {
     50         for col in 0..4 {
     51             let mut sum = 0.0;
     52             for i in 0..4 {
     53                 sum += a[row + i * 4] * b[i + col * 4];
     54             }
     55             result[row + col * 4] = sum;
     56         }
     57     }
     58 
     59     result
     60 }
     61 
     62 fn lerp3(a: [f32; 3], b: [f32; 3], t: f32) -> [f32; 3] {
     63     [
     64         a[0] + (b[0] - a[0]) * t,
     65         a[1] + (b[1] - a[1]) * t,
     66         a[2] + (b[2] - a[2]) * t,
     67     ]
     68 }
     69 
     70 fn generate_dave_instances(instance_count: u32) -> Vec<mesh::Instance> {
     71     let mut rng = rand::rng();
     72     let mut instances = Vec::with_capacity(instance_count as usize);
     73 
     74     // Logo gradient endpoints (0–1 range)
     75     const C0: [f32; 3] = [53.0 / 255.0, 77.0 / 255.0, 235.0 / 255.0]; // rgb(53, 77, 235)
     76     const C1: [f32; 3] = [229.0 / 255.0, 20.0 / 255.0, 205.0 / 255.0]; // rgb(229, 20, 205)
     77     let golden_angle = std::f32::consts::PI * (3.0 - 5.0_f32.sqrt());
     78 
     79     for i in 0..instance_count {
     80         let i_f = (i as f32) + 0.5;
     81         let n = instance_count as f32;
     82 
     83         // Fibonacci sphere (unit directions)
     84         let z = 1.0 - (2.0 * i_f) / n;
     85         let r = (1.0 - z * z).sqrt();
     86         let theta = golden_angle * i_f;
     87 
     88         // Use base_pos as *direction*; shader will normalize/scale anyway
     89         let base_pos = [r * theta.cos(), z, r * theta.sin()];
     90 
     91         let scale = 0.03;
     92 
     93         //let scale = scale + scale_var + rng.random::<f32>() * scale; // slightly smaller cubes
     94         let seed = rng.random::<f32>() * 1000.0;
     95 
     96         // damus logo gradient
     97         let t_base = (z + 1.0) * 0.5; // 0..1
     98         let t_jitter = (rng.random::<f32>() - 0.5) * 0.06; // ±0.03
     99         let t = (t_base + t_jitter).clamp(0.0, 1.0);
    100         let color = lerp3(C0, C1, t);
    101 
    102         instances.push(mesh::Instance {
    103             base_pos,
    104             scale,
    105             seed,
    106             color,
    107         });
    108     }
    109 
    110     instances
    111 }
    112 
    113 impl DaveAvatar {
    114     pub fn new(wgpu_render_state: &egui_wgpu::RenderState) -> Self {
    115         const BINDING_SIZE: u64 = 256;
    116 
    117         let device = &wgpu_render_state.device;
    118         let instance_count: u32 = 256;
    119         let instances = generate_dave_instances(instance_count);
    120 
    121         // Create shader module with improved shader code
    122         let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
    123             label: Some("cube_shader"),
    124             source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("dave.wgsl"))),
    125         });
    126 
    127         // Create uniform buffer for MVP matrix and model matrix
    128         let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
    129             label: Some("cube_uniform_buffer"),
    130             size: BINDING_SIZE, // Two 4x4 matrices of f32 (2 * 16 * 4 bytes)
    131             usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    132             mapped_at_creation: false,
    133         });
    134 
    135         // Create bind group layout
    136         let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    137             label: Some("cube_bind_group_layout"),
    138             entries: &[wgpu::BindGroupLayoutEntry {
    139                 binding: 0,
    140                 visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
    141                 ty: wgpu::BindingType::Buffer {
    142                     ty: wgpu::BufferBindingType::Uniform,
    143                     has_dynamic_offset: false,
    144                     min_binding_size: NonZeroU64::new(BINDING_SIZE),
    145                 },
    146                 count: None,
    147             }],
    148         });
    149 
    150         // Create bind group
    151         let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
    152             label: Some("cube_bind_group"),
    153             layout: &bind_group_layout,
    154             entries: &[wgpu::BindGroupEntry {
    155                 binding: 0,
    156                 resource: uniform_buffer.as_entire_binding(),
    157             }],
    158         });
    159 
    160         // Create pipeline layout
    161         let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    162             label: Some("cube_pipeline_layout"),
    163             bind_group_layouts: &[&bind_group_layout],
    164             push_constant_ranges: &[],
    165         });
    166 
    167         let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
    168             label: Some("cube_vertices"),
    169             contents: bytemuck::cast_slice(&mesh::CUBE_VERTICES),
    170             usage: wgpu::BufferUsages::VERTEX,
    171         });
    172 
    173         let instance_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
    174             label: Some("cube_instances"),
    175             contents: bytemuck::cast_slice(&instances),
    176             usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
    177         });
    178 
    179         let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
    180             label: Some("cube_indices"),
    181             contents: bytemuck::cast_slice(&mesh::CUBE_INDICES),
    182             usage: wgpu::BufferUsages::INDEX,
    183         });
    184 
    185         // Create render pipeline
    186         let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
    187             label: Some("cube_pipeline"),
    188             layout: Some(&pipeline_layout),
    189             vertex: wgpu::VertexState {
    190                 module: &shader,
    191                 entry_point: Some("vs_main"),
    192                 buffers: &[mesh::Vertex::LAYOUT, mesh::Instance::LAYOUT],
    193                 compilation_options: wgpu::PipelineCompilationOptions::default(),
    194             },
    195             fragment: Some(wgpu::FragmentState {
    196                 module: &shader,
    197                 entry_point: Some("fs_main"),
    198                 targets: &[Some(wgpu::ColorTargetState {
    199                     format: wgpu_render_state.target_format,
    200                     blend: Some(wgpu::BlendState::ALPHA_BLENDING),
    201                     write_mask: wgpu::ColorWrites::ALL,
    202                 })],
    203                 compilation_options: wgpu::PipelineCompilationOptions::default(),
    204             }),
    205             primitive: wgpu::PrimitiveState {
    206                 topology: wgpu::PrimitiveTopology::TriangleList,
    207                 strip_index_format: None,
    208                 front_face: wgpu::FrontFace::Ccw,
    209                 cull_mode: Some(wgpu::Face::Back),
    210                 polygon_mode: wgpu::PolygonMode::Fill,
    211                 unclipped_depth: false,
    212                 conservative: false,
    213             },
    214             depth_stencil: Some(wgpu::DepthStencilState {
    215                 format: wgpu::TextureFormat::Depth24Plus,
    216                 depth_write_enabled: true,
    217                 depth_compare: wgpu::CompareFunction::Less,
    218                 stencil: wgpu::StencilState::default(),
    219                 bias: wgpu::DepthBiasState::default(),
    220             }),
    221             multisample: wgpu::MultisampleState::default(),
    222             multiview: None,
    223             cache: None,
    224         });
    225 
    226         // Store resources in renderer
    227         wgpu_render_state
    228             .renderer
    229             .write()
    230             .callback_resources
    231             .insert(CubeRenderResources {
    232                 pipeline,
    233                 bind_group,
    234                 uniform_buffer,
    235                 instance_buffer,
    236                 vertex_buffer,
    237                 index_buffer,
    238                 instance_count,
    239             });
    240 
    241         let initial_rot = {
    242             let x_rotation = Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), 0.5);
    243             let y_rotation = Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), 0.5);
    244 
    245             // Apply rotations (order matters)
    246             y_rotation.multiply(&x_rotation)
    247         };
    248 
    249         Self {
    250             logical_time: 0.0,
    251             rotation: initial_rot,
    252             rot_dir: Vec3::new(0.0, 0.0, 0.0),
    253         }
    254     }
    255 }
    256 
    257 #[inline]
    258 fn apply_friction(val: f32, friction: f32, clamp: f32) -> f32 {
    259     if val < clamp {
    260         0.0
    261     } else {
    262         val * friction
    263     }
    264 }
    265 
    266 impl DaveAvatar {
    267     pub fn random_nudge(&mut self) {
    268         self.random_nudge_with(1.0);
    269     }
    270 
    271     pub fn random_nudge_with(&mut self, force: f32) {
    272         let mut rng = rand::rng();
    273 
    274         let nudge = Vec3::new(
    275             rng.random::<f32>() * force,
    276             rng.random::<f32>() * force,
    277             rng.random::<f32>() * force,
    278         )
    279         .normalize();
    280 
    281         self.rot_dir.x += nudge.x;
    282         self.rot_dir.y += nudge.y;
    283         self.rot_dir.z += nudge.z;
    284     }
    285 
    286     pub fn render(&mut self, rect: Rect, ui: &mut egui::Ui) -> Response {
    287         let response = ui.allocate_rect(rect, egui::Sense::CLICK | egui::Sense::DRAG);
    288 
    289         // Update rotation based on drag or animation
    290         if response.dragged() {
    291             // Create rotation quaternions based on drag
    292             let dx = response.drag_delta().x;
    293             let dy = response.drag_delta().y;
    294             let x_rotation = Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), dy * 0.01);
    295             let y_rotation = Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), dx * 0.01);
    296 
    297             self.rot_dir = Vec3::new(dx, dy, 0.0);
    298 
    299             // Apply rotations (order matters)
    300             self.rotation = y_rotation.multiply(&x_rotation).multiply(&self.rotation);
    301         } else if response.clicked() {
    302             self.random_nudge_with(1.0);
    303         } else {
    304             // Continuous rotation - reduced speed and simplified axis
    305             let friction = 0.95;
    306             let clamp = 0.1;
    307             self.rot_dir.x = apply_friction(self.rot_dir.x, friction, clamp);
    308             self.rot_dir.y = apply_friction(self.rot_dir.y, friction, clamp);
    309             self.rot_dir.z = apply_friction(self.rot_dir.y, friction, clamp);
    310 
    311             // we only need to render if we're still spinning
    312             if self.rot_dir.x > clamp || self.rot_dir.y > clamp || self.rot_dir.z > clamp {
    313                 let x_rotation =
    314                     Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), self.rot_dir.y * 0.03);
    315                 let y_rotation =
    316                     Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), self.rot_dir.x * 0.03);
    317                 let z_rotation =
    318                     Quaternion::from_axis_angle(&Vec3::new(0.0, 0.0, 1.0), self.rot_dir.z * 0.03);
    319 
    320                 self.rotation = y_rotation
    321                     .multiply(&x_rotation)
    322                     .multiply(&z_rotation)
    323                     .multiply(&self.rotation);
    324 
    325                 tracing::trace!("repainting due to avatar rotation");
    326                 ui.ctx().request_repaint();
    327             }
    328         }
    329 
    330         // Create model matrix from rotation quaternion
    331         let model = self.rotation.to_matrix4();
    332 
    333         // Create projection matrix with proper depth range
    334         // Adjust aspect ratio based on rect dimensions
    335         let aspect = rect.width() / rect.height();
    336         let projection = perspective_matrix(std::f32::consts::PI / 4.0, aspect, 0.1, 100.0);
    337 
    338         // Create view matrix (move camera back a bit)
    339         let camera_pos = [0.0, 0.0, 1.5];
    340 
    341         // Right-handed look-at at origin; view is a translate by -camera_pos
    342         let [cx, cy, cz] = camera_pos;
    343 
    344         #[rustfmt::skip]
    345         let view = [
    346             1.0, 0.0, 0.0, 0.0,
    347             0.0, 1.0, 0.0, 0.0,
    348             0.0, 0.0, 1.0, 0.0,
    349             -cx, -cy, -cz, 1.0,
    350         ];
    351 
    352         let view_proj = matrix_multiply(&projection, &view);
    353         let is_light = if ui.ctx().theme() == egui::Theme::Light {
    354             1.0
    355         } else {
    356             -1.0
    357         };
    358 
    359         self.logical_time += ui.ctx().input(|i| i.stable_dt.min(0.1));
    360 
    361         // Add paint callback
    362         ui.painter().add(egui_wgpu::Callback::new_paint_callback(
    363             rect,
    364             GpuData {
    365                 view_proj,
    366                 model,
    367                 camera_pos,
    368                 time: self.logical_time,
    369                 is_light: [is_light, 0.0, 0.0, 0.0],
    370             },
    371         ));
    372 
    373         response
    374     }
    375 }
    376 
    377 // Callback implementation
    378 #[repr(C)]
    379 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
    380 struct GpuData {
    381     view_proj: [f32; 16], // Model-View-Projection matrix
    382     model: [f32; 16],     // Model matrix for lighting calculations
    383     camera_pos: [f32; 3], // xyz
    384     time: f32,
    385     is_light: [f32; 4],
    386 }
    387 
    388 impl egui_wgpu::CallbackTrait for GpuData {
    389     fn prepare(
    390         &self,
    391         _device: &wgpu::Device,
    392         queue: &wgpu::Queue,
    393         _screen_descriptor: &egui_wgpu::ScreenDescriptor,
    394         _egui_encoder: &mut wgpu::CommandEncoder,
    395         resources: &mut egui_wgpu::CallbackResources,
    396     ) -> Vec<wgpu::CommandBuffer> {
    397         let resources: &CubeRenderResources = resources.get().unwrap();
    398 
    399         // Update uniform buffer with both matrices
    400         queue.write_buffer(&resources.uniform_buffer, 0, bytemuck::bytes_of(self));
    401 
    402         Vec::new()
    403     }
    404 
    405     fn paint(
    406         &self,
    407         _info: egui::PaintCallbackInfo,
    408         render_pass: &mut wgpu::RenderPass,
    409         resources: &egui_wgpu::CallbackResources,
    410     ) {
    411         let resources: &CubeRenderResources = resources.get().unwrap();
    412 
    413         render_pass.set_pipeline(&resources.pipeline);
    414         render_pass.set_bind_group(0, &resources.bind_group, &[]);
    415         render_pass.set_vertex_buffer(0, resources.vertex_buffer.slice(..));
    416         render_pass.set_vertex_buffer(1, resources.instance_buffer.slice(..));
    417         render_pass.set_index_buffer(resources.index_buffer.slice(..), wgpu::IndexFormat::Uint16);
    418         render_pass.draw_indexed(
    419             0..mesh::CUBE_INDICES.len() as u32,
    420             0,
    421             0..resources.instance_count,
    422         );
    423     }
    424 }
    425 
    426 // Simple resources struct
    427 struct CubeRenderResources {
    428     pipeline: wgpu::RenderPipeline,
    429     bind_group: wgpu::BindGroup,
    430     uniform_buffer: wgpu::Buffer,
    431     instance_buffer: wgpu::Buffer,
    432     vertex_buffer: wgpu::Buffer,
    433     index_buffer: wgpu::Buffer,
    434     instance_count: u32,
    435 }