Files
notedeck/crates/notedeck_dave/src/avatar.rs
William Casarin 2539dead1e dave: nudge avatar when you click
Signed-off-by: William Casarin <jb55@jb55.com>
2025-05-01 20:21:30 -07:00

452 lines
15 KiB
Rust

use std::num::NonZeroU64;
use crate::{Quaternion, Vec3};
use eframe::egui_wgpu::{self, wgpu};
use egui::{Rect, Response};
use rand::Rng;
pub struct DaveAvatar {
rotation: Quaternion,
rot_dir: Vec3,
}
// Matrix utilities for perspective projection
fn perspective_matrix(fovy_radians: f32, aspect: f32, near: f32, far: f32) -> [f32; 16] {
let f = 1.0 / (fovy_radians / 2.0).tan();
let nf = 1.0 / (near - far);
// Column-major for WGPU
[
f / aspect,
0.0,
0.0,
0.0,
0.0,
f,
0.0,
0.0,
0.0,
0.0,
(far + near) * nf,
-1.0,
0.0,
0.0,
2.0 * far * near * nf,
0.0,
]
}
// Combine two 4x4 matrices (column-major)
fn matrix_multiply(a: &[f32; 16], b: &[f32; 16]) -> [f32; 16] {
let mut result = [0.0; 16];
for row in 0..4 {
for col in 0..4 {
let mut sum = 0.0;
for i in 0..4 {
sum += a[row + i * 4] * b[i + col * 4];
}
result[row + col * 4] = sum;
}
}
result
}
impl DaveAvatar {
pub fn new(wgpu_render_state: &egui_wgpu::RenderState) -> Self {
let device = &wgpu_render_state.device;
// Create shader module with improved shader code
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("cube_shader"),
source: wgpu::ShaderSource::Wgsl(
r#"
struct Uniforms {
model_view_proj: mat4x4<f32>,
model: mat4x4<f32>, // Added model matrix for correct normal transformation
};
@group(0) @binding(0)
var<uniform> uniforms: Uniforms;
struct VertexOutput {
@builtin(position) position: vec4<f32>,
@location(0) normal: vec3<f32>,
@location(1) world_pos: vec3<f32>,
};
@vertex
fn vs_main(@builtin(vertex_index) vertex_index: u32) -> VertexOutput {
// Define cube vertices (-0.5 to 0.5 in each dimension)
var positions = array<vec3<f32>, 8>(
vec3<f32>(-0.5, -0.5, -0.5), // 0: left bottom back
vec3<f32>(0.5, -0.5, -0.5), // 1: right bottom back
vec3<f32>(-0.5, 0.5, -0.5), // 2: left top back
vec3<f32>(0.5, 0.5, -0.5), // 3: right top back
vec3<f32>(-0.5, -0.5, 0.5), // 4: left bottom front
vec3<f32>(0.5, -0.5, 0.5), // 5: right bottom front
vec3<f32>(-0.5, 0.5, 0.5), // 6: left top front
vec3<f32>(0.5, 0.5, 0.5) // 7: right top front
);
// Define indices for the 12 triangles (6 faces * 2 triangles)
var indices = array<u32, 36>(
// back face (Z-)
0, 2, 1, 1, 2, 3,
// front face (Z+)
4, 5, 6, 5, 7, 6,
// left face (X-)
0, 4, 2, 2, 4, 6,
// right face (X+)
1, 3, 5, 3, 7, 5,
// bottom face (Y-)
0, 1, 4, 1, 5, 4,
// top face (Y+)
2, 6, 3, 3, 6, 7
);
// Define normals for each face
var face_normals = array<vec3<f32>, 6>(
vec3<f32>(0.0, 0.0, -1.0), // back face (Z-)
vec3<f32>(0.0, 0.0, 1.0), // front face (Z+)
vec3<f32>(-1.0, 0.0, 0.0), // left face (X-)
vec3<f32>(1.0, 0.0, 0.0), // right face (X+)
vec3<f32>(0.0, -1.0, 0.0), // bottom face (Y-)
vec3<f32>(0.0, 1.0, 0.0) // top face (Y+)
);
var output: VertexOutput;
// Get vertex from indices
let index = indices[vertex_index];
let position = positions[index];
// Determine which face this vertex belongs to
let face_index = vertex_index / 6u;
// Apply transformations
output.position = uniforms.model_view_proj * vec4<f32>(position, 1.0);
// Transform normal to world space
// Extract the 3x3 rotation part from the 4x4 model matrix
let normal_matrix = mat3x3<f32>(
uniforms.model[0].xyz,
uniforms.model[1].xyz,
uniforms.model[2].xyz
);
output.normal = normalize(normal_matrix * face_normals[face_index]);
// Pass world position for lighting calculations
output.world_pos = (uniforms.model * vec4<f32>(position, 1.0)).xyz;
return output;
}
@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
// Material properties
let material_color = vec3<f32>(1.0, 1.0, 1.0); // White color
let ambient_strength = 0.2;
let diffuse_strength = 0.7;
let specular_strength = 0.2;
let shininess = 20.0;
// Light properties
let light_pos = vec3<f32>(2.0, 2.0, 2.0); // Light positioned diagonally above and to the right
let light_color = vec3<f32>(1.0, 1.0, 1.0); // White light
// View position (camera)
let view_pos = vec3<f32>(0.0, 0.0, 3.0); // Camera position
// Calculate ambient lighting
let ambient = ambient_strength * light_color;
// Calculate diffuse lighting
let normal = normalize(in.normal); // Renormalize the interpolated normal
let light_dir = normalize(light_pos - in.world_pos);
let diff = max(dot(normal, light_dir), 0.0);
let diffuse = diffuse_strength * diff * light_color;
// Calculate specular lighting
let view_dir = normalize(view_pos - in.world_pos);
let reflect_dir = reflect(-light_dir, normal);
let spec = pow(max(dot(view_dir, reflect_dir), 0.0), shininess);
let specular = specular_strength * spec * light_color;
// Combine lighting components
let result = (ambient + diffuse + specular) * material_color;
return vec4<f32>(result, 1.0);
}
"#
.into(),
),
});
// Create uniform buffer for MVP matrix and model matrix
let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("cube_uniform_buffer"),
size: 128, // Two 4x4 matrices of f32 (2 * 16 * 4 bytes)
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
// Create bind group layout
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("cube_bind_group_layout"),
entries: &[wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::VERTEX,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: NonZeroU64::new(128),
},
count: None,
}],
});
// Create bind group
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("cube_bind_group"),
layout: &bind_group_layout,
entries: &[wgpu::BindGroupEntry {
binding: 0,
resource: uniform_buffer.as_entire_binding(),
}],
});
// Create pipeline layout
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("cube_pipeline_layout"),
bind_group_layouts: &[&bind_group_layout],
push_constant_ranges: &[],
});
// Create render pipeline
let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
label: Some("cube_pipeline"),
layout: Some(&pipeline_layout),
vertex: wgpu::VertexState {
module: &shader,
entry_point: Some("vs_main"),
buffers: &[], // No vertex buffer - vertices are in the shader
compilation_options: wgpu::PipelineCompilationOptions::default(),
},
fragment: Some(wgpu::FragmentState {
module: &shader,
entry_point: Some("fs_main"),
targets: &[Some(wgpu::ColorTargetState {
format: wgpu_render_state.target_format,
blend: Some(wgpu::BlendState::ALPHA_BLENDING),
write_mask: wgpu::ColorWrites::ALL,
})],
compilation_options: wgpu::PipelineCompilationOptions::default(),
}),
primitive: wgpu::PrimitiveState {
topology: wgpu::PrimitiveTopology::TriangleList,
strip_index_format: None,
front_face: wgpu::FrontFace::Ccw,
cull_mode: Some(wgpu::Face::Back),
polygon_mode: wgpu::PolygonMode::Fill,
unclipped_depth: false,
conservative: false,
},
depth_stencil: Some(wgpu::DepthStencilState {
format: wgpu::TextureFormat::Depth24Plus,
depth_write_enabled: true,
depth_compare: wgpu::CompareFunction::Less,
stencil: wgpu::StencilState::default(),
bias: wgpu::DepthBiasState::default(),
}),
multisample: wgpu::MultisampleState::default(),
multiview: None,
cache: None,
});
// Store resources in renderer
wgpu_render_state
.renderer
.write()
.callback_resources
.insert(CubeRenderResources {
pipeline,
bind_group,
uniform_buffer,
});
let initial_rot = {
let x_rotation = Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), 0.5);
let y_rotation = Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), 0.5);
// Apply rotations (order matters)
y_rotation.multiply(&x_rotation)
};
Self {
rotation: initial_rot,
rot_dir: Vec3::new(0.0, 0.0, 0.0),
}
}
}
#[inline]
fn apply_friction(val: f32, friction: f32, clamp: f32) -> f32 {
if val < clamp {
0.0
} else {
val * friction
}
}
impl DaveAvatar {
pub fn random_nudge(&mut self) {
self.random_nudge_with(1.0);
}
pub fn random_nudge_with(&mut self, force: f32) {
let mut rng = rand::rng();
let nudge = Vec3::new(
rng.random::<f32>() * force,
rng.random::<f32>() * force,
rng.random::<f32>() * force,
)
.normalize();
self.rot_dir.x += nudge.x;
self.rot_dir.y += nudge.y;
self.rot_dir.z += nudge.z;
}
pub fn render(&mut self, rect: Rect, ui: &mut egui::Ui) -> Response {
let response = ui.allocate_rect(rect, egui::Sense::CLICK | egui::Sense::DRAG);
// Update rotation based on drag or animation
if response.dragged() {
// Create rotation quaternions based on drag
let dx = response.drag_delta().x;
let dy = response.drag_delta().y;
let x_rotation = Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), dy * 0.01);
let y_rotation = Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), dx * 0.01);
self.rot_dir = Vec3::new(dx, dy, 0.0);
// Apply rotations (order matters)
self.rotation = y_rotation.multiply(&x_rotation).multiply(&self.rotation);
} else if response.clicked() {
self.random_nudge_with(1.0);
} else {
// Continuous rotation - reduced speed and simplified axis
let friction = 0.95;
let clamp = 0.1;
self.rot_dir.x = apply_friction(self.rot_dir.x, friction, clamp);
self.rot_dir.y = apply_friction(self.rot_dir.y, friction, clamp);
self.rot_dir.z = apply_friction(self.rot_dir.y, friction, clamp);
// we only need to render if we're still spinning
if self.rot_dir.x > clamp || self.rot_dir.y > clamp || self.rot_dir.z > clamp {
let x_rotation =
Quaternion::from_axis_angle(&Vec3::new(1.0, 0.0, 0.0), self.rot_dir.y * 0.03);
let y_rotation =
Quaternion::from_axis_angle(&Vec3::new(0.0, 1.0, 0.0), self.rot_dir.x * 0.03);
let z_rotation =
Quaternion::from_axis_angle(&Vec3::new(0.0, 0.0, 1.0), self.rot_dir.z * 0.03);
self.rotation = y_rotation
.multiply(&x_rotation)
.multiply(&z_rotation)
.multiply(&self.rotation);
tracing::trace!("repainting due to avatar rotation");
ui.ctx().request_repaint();
}
}
// Create model matrix from rotation quaternion
let model_matrix = self.rotation.to_matrix4();
// Create projection matrix with proper depth range
// Adjust aspect ratio based on rect dimensions
let aspect = rect.width() / rect.height();
let projection = perspective_matrix(std::f32::consts::PI / 4.0, aspect, 0.1, 100.0);
// Create view matrix (move camera back a bit)
let view_matrix = [
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,
];
// Combine matrices: projection * view * model
let mv_matrix = matrix_multiply(&view_matrix, &model_matrix);
let mvp_matrix = matrix_multiply(&projection, &mv_matrix);
// Add paint callback
ui.painter().add(egui_wgpu::Callback::new_paint_callback(
rect,
CubeCallback {
mvp_matrix,
model_matrix,
},
));
response
}
}
// Callback implementation
struct CubeCallback {
mvp_matrix: [f32; 16], // Model-View-Projection matrix
model_matrix: [f32; 16], // Model matrix for lighting calculations
}
impl egui_wgpu::CallbackTrait for CubeCallback {
fn prepare(
&self,
_device: &wgpu::Device,
queue: &wgpu::Queue,
_screen_descriptor: &egui_wgpu::ScreenDescriptor,
_egui_encoder: &mut wgpu::CommandEncoder,
resources: &mut egui_wgpu::CallbackResources,
) -> Vec<wgpu::CommandBuffer> {
let resources: &CubeRenderResources = resources.get().unwrap();
// Create a combined uniform buffer with both matrices
let mut uniform_data = [0.0f32; 32]; // Space for two 4x4 matrices
// Copy MVP matrix to first 16 floats
uniform_data[0..16].copy_from_slice(&self.mvp_matrix);
// Copy model matrix to next 16 floats
uniform_data[16..32].copy_from_slice(&self.model_matrix);
// Update uniform buffer with both matrices
queue.write_buffer(
&resources.uniform_buffer,
0,
bytemuck::cast_slice(&uniform_data),
);
Vec::new()
}
fn paint(
&self,
_info: egui::PaintCallbackInfo,
render_pass: &mut wgpu::RenderPass,
resources: &egui_wgpu::CallbackResources,
) {
let resources: &CubeRenderResources = resources.get().unwrap();
render_pass.set_pipeline(&resources.pipeline);
render_pass.set_bind_group(0, &resources.bind_group, &[]);
render_pass.draw(0..36, 0..1); // 36 vertices for a cube (6 faces * 2 triangles * 3 vertices)
}
}
// Simple resources struct
struct CubeRenderResources {
pipeline: wgpu::RenderPipeline,
bind_group: wgpu::BindGroup,
uniform_buffer: wgpu::Buffer,
}