src/workshop.rs
Repository source snapshot · Plain text · GitHub
1//! Eight executable milestones. No window, DOM, file I/O or event loop here.2use bytemuck::{Pod, Zeroable};3use glam::{Mat4, Vec3};4use wgpu::util::DeviceExt;56// ANCHOR: layout7#[repr(C)]8#[derive(Clone, Copy, Pod, Zeroable)]9struct Uniforms {10 view_projection: [[f32; 4]; 4], // bytes 0..64, column-major11 viewport_stage: [f32; 4], // bytes 64..80: width, height, stage, padding12}1314#[repr(C)]15#[derive(Clone, Copy, Pod, Zeroable)]16struct Instance {17 center_alpha: [f32; 4], // bytes 0..1618 scale_pad: [f32; 4], // bytes 16..32, activated standard deviations19 color_pad: [f32; 4], // bytes 32..48, direct RGB (not SH coefficients)20}21// END: layout2223#[derive(Clone, Copy, Debug)]24pub struct Frame {25 pub stage: u32,26 pub time: f32,27 pub yaw: f32,28 pub pitch: f32,29 pub distance: f32,30 pub reverse_order: bool,31}3233impl Default for Frame {34 fn default() -> Self {35 Self {36 stage: 6,37 time: 0.5,38 yaw: 0.0,39 pitch: 0.15,40 distance: 5.0,41 reverse_order: false,42 }43 }44}4546impl Frame {47 fn validate(self) -> Result<(), String> {48 if self.stage > 749 || ![self.time, self.yaw, self.pitch, self.distance]50 .iter()51 .all(|x| x.is_finite())52 || !(0.0..=1.0).contains(&self.time)53 || !(1.5..=60.0).contains(&self.distance)54 || self.pitch.abs() > 1.455 {56 return Err("Invalid workshop stage, normalized time or camera".into());57 }58 Ok(())59 }60}6162// ANCHOR: camera63pub fn camera(frame: Frame, width: u32, height: u32) -> (Mat4, Mat4) {64 // The research capture uses a different world orientation from the toy scene.65 let (target, up, sign, distance) = if frame.stage == 7 {66 (67 Vec3::new(0.0, 3.5, 14.0),68 Vec3::NEG_Y,69 -1.0,70 frame.distance * 3.0,71 )72 } else {73 (Vec3::ZERO, Vec3::Y, 1.0, frame.distance)74 };75 let orbit = Vec3::new(76 frame.yaw.sin() * frame.pitch.cos(),77 frame.pitch.sin(),78 sign * frame.yaw.cos() * frame.pitch.cos(),79 );80 let view = Mat4::look_at_rh(target + orbit * distance, target, up);81 let projection = Mat4::perspective_rh(82 50_f32.to_radians(),83 width as f32 / height as f32,84 0.05,85 200.0,86 );87 (view, projection)88}89// END: camera9091fn instances(frame: Frame, view: Mat4) -> Vec<Instance> {92 let make = |p: [f32; 3], a, scale: [f32; 3], rgb: [f32; 3]| Instance {93 center_alpha: [p[0], p[1], p[2], a],94 scale_pad: [scale[0], scale[1], scale[2], 0.0],95 color_pad: [rgb[0], rgb[1], rgb[2], 0.0],96 };97 let mut values = if frame.stage < 4 {98 vec![make(99 [0.0, 0.0, 0.0],100 0.8,101 [0.48, 0.22, 0.16],102 [1.0, 0.45, 0.08],103 )]104 } else {105 vec![106 make([0.0, 0.0, 0.5], 0.5, [0.45, 0.2, 0.15], [1.0, 0.0, 0.0]),107 make([0.0, 0.0, -0.5], 0.5, [0.2, 0.5, 0.2], [0.0, 0.25, 1.0]),108 ]109 };110 // ANCHOR: motion111 if frame.stage == 6 {112 let dt = frame.time - 0.5;113 values[0].center_alpha[0] += 2.5 * dt;114 values[0].center_alpha[1] += 1.5 * dt * dt;115 values[0].center_alpha[3] *= (-(dt / 0.45).powi(2)).exp();116 values[1].center_alpha[0] -= 1.5 * dt;117 }118 // Sorting happens AFTER evaluating the current time and camera.119 values.sort_by(|a, b| {120 let depth = |s: &Instance| -(view * Vec3::from_slice(&s.center_alpha[..3]).extend(1.0)).z;121 depth(b).total_cmp(&depth(a))122 });123 if frame.reverse_order {124 values.reverse();125 }126 // END: motion127 values128}129130pub struct WorkshopPass {131 pipeline: wgpu::RenderPipeline,132 bindings: wgpu::BindGroup,133 uniforms: wgpu::Buffer,134 instances: wgpu::Buffer,135 stage: u32,136 count: u32,137 real: Option<crate::stg_pass::StgPass>,138 format: wgpu::TextureFormat,139}140141impl WorkshopPass {142 pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {143 // ANCHOR: buffers144 let uniforms = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {145 label: Some("workshop-uniforms-80-bytes"),146 contents: bytemuck::bytes_of(&Uniforms::zeroed()),147 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,148 });149 let instances = device.create_buffer(&wgpu::BufferDescriptor {150 label: Some("workshop-instances-2-times-48-bytes"),151 size: 96,152 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,153 mapped_at_creation: false,154 });155 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {156 label: Some("workshop-bindings"),157 entries: &[158 wgpu::BindGroupLayoutEntry {159 binding: 0,160 visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,161 ty: wgpu::BindingType::Buffer {162 ty: wgpu::BufferBindingType::Uniform,163 has_dynamic_offset: false,164 min_binding_size: None,165 },166 count: None,167 },168 wgpu::BindGroupLayoutEntry {169 binding: 1,170 visibility: wgpu::ShaderStages::VERTEX,171 ty: wgpu::BindingType::Buffer {172 ty: wgpu::BufferBindingType::Storage { read_only: true },173 has_dynamic_offset: false,174 min_binding_size: None,175 },176 count: None,177 },178 ],179 });180 let bindings = device.create_bind_group(&wgpu::BindGroupDescriptor {181 label: Some("workshop-bind-group"),182 layout: &layout,183 entries: &[184 wgpu::BindGroupEntry {185 binding: 0,186 resource: uniforms.as_entire_binding(),187 },188 wgpu::BindGroupEntry {189 binding: 1,190 resource: instances.as_entire_binding(),191 },192 ],193 });194 // END: buffers195 // ANCHOR: pipeline196 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {197 label: Some("workshop-shader"),198 source: wgpu::ShaderSource::Wgsl(include_str!("workshop.wgsl").into()),199 });200 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {201 label: Some("workshop-pipeline-layout"),202 bind_group_layouts: &[Some(&layout)],203 immediate_size: 0,204 });205 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {206 label: Some("workshop-pipeline"),207 layout: Some(&pipeline_layout),208 vertex: wgpu::VertexState {209 module: &shader,210 entry_point: Some("vs_main"),211 compilation_options: Default::default(),212 buffers: &[],213 },214 fragment: Some(wgpu::FragmentState {215 module: &shader,216 entry_point: Some("fs_main"),217 compilation_options: Default::default(),218 targets: &[Some(wgpu::ColorTargetState {219 format,220 blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),221 write_mask: wgpu::ColorWrites::ALL,222 })],223 }),224 primitive: Default::default(),225 depth_stencil: None,226 multisample: Default::default(),227 multiview_mask: None,228 cache: None,229 });230 // END: pipeline231 Self {232 pipeline,233 bindings,234 uniforms,235 instances,236 stage: 0,237 count: 0,238 real: None,239 format,240 }241 }242243 pub fn load_model(&mut self, device: &wgpu::Device, ply: &[u8]) -> Result<(), String> {244 self.real = Some(crate::stg_pass::StgPass::new(245 device,246 ply,247 self.format,248 None,249 )?);250 Ok(())251 }252253 // ANCHOR: prepare254 pub fn prepare(255 &mut self,256 queue: &wgpu::Queue,257 frame: Frame,258 size: [u32; 2],259 ) -> Result<u32, String> {260 frame.validate()?;261 if size.contains(&0) {262 return Err("Cannot prepare a zero-sized target".into());263 }264 let (view, projection) = camera(frame, size[0], size[1]);265 self.stage = frame.stage;266 if frame.stage == 7 {267 let pass = self.real.as_mut().ok_or("Load the STG-Lite model first")?;268 self.count = pass269 .prepare(270 queue,271 frame.time,272 view.to_cols_array_2d(),273 projection.to_cols_array_2d(),274 size,275 )?276 .visible_count as u32;277 } else {278 let data = instances(frame, view);279 self.count = data.len() as u32;280 queue.write_buffer(&self.instances, 0, bytemuck::cast_slice(&data));281 queue.write_buffer(282 &self.uniforms,283 0,284 bytemuck::bytes_of(&Uniforms {285 view_projection: (projection * view).to_cols_array_2d(),286 viewport_stage: [size[0] as f32, size[1] as f32, frame.stage as f32, 0.0],287 }),288 );289 }290 Ok(if frame.stage == 0 { 0 } else { self.count })291 }292293 pub fn draw(&self, pass: &mut wgpu::RenderPass<'_>) {294 if self.stage == 0 {295 return;296 }297 if self.stage == 7 {298 if let Some(real) = &self.real {299 real.draw(pass);300 }301 return;302 }303 pass.set_pipeline(&self.pipeline);304 pass.set_bind_group(0, &self.bindings, &[]);305 pass.draw(0..if self.stage == 1 { 3 } else { 6 }, 0..self.count);306 }307 // END: prepare308}309310#[cfg(test)]311mod tests {312 use super::*;313 #[test]314 fn host_shader_contract_and_invalid_input() {315 assert_eq!(std::mem::size_of::<Uniforms>(), 80);316 assert_eq!(std::mem::size_of::<Instance>(), 48);317 assert!(Frame {318 time: f32::NAN,319 ..Frame::default()320 }321 .validate()322 .is_err());323 assert!(Frame {324 stage: 8,325 ..Frame::default()326 }327 .validate()328 .is_err());329 }330 #[test]331 fn camera_and_time_change_the_ordered_scene() {332 let frame = Frame {333 stage: 6,334 ..Frame::default()335 };336 let (view, _) = camera(frame, 640, 480);337 let base = instances(frame, view);338 assert!(base[0].center_alpha[2] < base[1].center_alpha[2]);339 let moved = instances(Frame { time: 0.9, ..frame }, view);340 assert_ne!(base[1].center_alpha, moved[1].center_alpha);341 let reversed = instances(342 Frame {343 reverse_order: true,344 ..frame345 },346 view,347 );348 assert_eq!(base[0].center_alpha, reversed[1].center_alpha);349 }350}351