yum/demomo

A graphics engine built on wgpu.

git clone https://git.yummers.dev/yum/demomo

yumMake camera optionalff99423

master
branchesmasterall refs →
39.0 KiB1239 linesraw
1use anyhow::Context;
2use bytemuck::NoUninit;
3use derive_builder::Builder;
4use std::borrow::Cow;
5use std::num::NonZeroU64;
6#[cfg(feature = "window")]
7use std::sync::Arc;
8use std::sync::mpsc::{Receiver, Sender};
9use wgpu::naga;
10use wgpu::util::DeviceExt;
11
12#[macro_export]
13macro_rules! include_wgsl {
14    ($filename:literal) => {
15        include_str!(concat!(env!("OUT_DIR"), "/shaders/", $filename))
16    };
17}
18
19// Magic to allow converting a primitive type, a Vec, or a slice into a slice of bytes.
20pub trait Buffer<T> {
21    fn as_bytes(&self) -> &[T];
22}
23impl<T: bytemuck::Pod> Buffer<T> for T {
24    fn as_bytes(&self) -> &[T] {
25        std::slice::from_ref(self)
26    }
27}
28impl<T> Buffer<T> for Vec<T> {
29    fn as_bytes(&self) -> &[T] {
30        self.as_slice()
31    }
32}
33impl<T> Buffer<T> for &[T] {
34    fn as_bytes(&self) -> &[T] {
35        self
36    }
37}
38
39pub struct BufferBuilder<'a> {
40    label: Option<String>,
41    size: u64,
42    usage: wgpu::BufferUsages,
43    mapped_at_creation: bool,
44    gpu: &'a Gpu,
45}
46
47impl<'a> BufferBuilder<'a> {
48    fn new<T: bytemuck::Pod>(data: &impl Buffer<T>, gpu: &'a Gpu) -> Self {
49        Self::new_by_sz(std::mem::size_of_val(data.as_bytes()), gpu)
50    }
51
52    fn new_by_sz(sz: usize, gpu: &'a Gpu) -> Self {
53        Self {
54            label: None,
55            size: wgpu::util::align_to(sz, 16) as u64,
56            usage: wgpu::BufferUsages::empty(),
57            mapped_at_creation: false,
58            gpu: gpu,
59        }
60    }
61
62    pub fn copy_dst(&mut self) -> &mut Self {
63        self.usage |= wgpu::BufferUsages::COPY_DST;
64        self
65    }
66
67    pub fn copy_src(&mut self) -> &mut Self {
68        self.usage |= wgpu::BufferUsages::COPY_SRC;
69        self
70    }
71
72    pub fn map_read(&mut self) -> &mut Self {
73        self.usage |= wgpu::BufferUsages::MAP_READ;
74        self
75    }
76
77    pub fn query_resolve(&mut self) -> &mut Self {
78        self.usage |= wgpu::BufferUsages::QUERY_RESOLVE;
79        self
80    }
81
82    pub fn storage(&mut self) -> &mut Self {
83        self.usage |= wgpu::BufferUsages::STORAGE;
84        self
85    }
86
87    pub fn uniform(&mut self) -> &mut Self {
88        self.usage |= wgpu::BufferUsages::UNIFORM;
89        self
90    }
91
92    pub fn build(&self) -> wgpu::Buffer {
93        self.gpu.device.create_buffer(&wgpu::BufferDescriptor {
94            label: self.label.as_deref(),
95            size: self.size,
96            usage: self.usage,
97            mapped_at_creation: self.mapped_at_creation,
98        })
99    }
100}
101
102pub struct GpuBuilder {
103    required_features: wgpu::Features,
104    #[cfg(feature = "window")]
105    window: Option<Arc<winit::window::Window>>,
106}
107
108impl GpuBuilder {
109    pub fn new() -> Self {
110        Self {
111            required_features: wgpu::Features::empty(),
112            #[cfg(feature = "window")]
113            window: None,
114        }
115    }
116
117    pub fn enable_profiling(mut self) -> Self {
118        self.required_features |= wgpu::Features::TIMESTAMP_QUERY;
119        self
120    }
121
122    #[cfg(feature = "window")]
123    pub fn window(mut self, window: winit::window::Window) -> Self {
124        self.window = Some(Arc::new(window));
125        self
126    }
127
128    pub async fn build(self) -> anyhow::Result<Gpu> {
129        #[cfg(feature = "window")]
130        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
131            backends: wgpu::Backends::PRIMARY,
132            flags: Default::default(),
133            memory_budget_thresholds: Default::default(),
134            backend_options: Default::default(),
135            display: None,
136        });
137        #[cfg(not(feature = "window"))]
138        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
139
140        #[cfg(feature = "window")]
141        let surface = self
142            .window
143            .as_ref()
144            .map(|window| instance.create_surface(window.clone()))
145            .transpose()?;
146        #[cfg(feature = "window")]
147        let adapter_req_opt = wgpu::RequestAdapterOptions {
148            power_preference: wgpu::PowerPreference::default(),
149            compatible_surface: surface.as_ref(),
150            force_fallback_adapter: false,
151            apply_limit_buckets: true,
152        };
153        #[cfg(not(feature = "window"))]
154        let adapter_req_opt = wgpu::RequestAdapterOptions::default();
155        let adapter = instance.request_adapter(&adapter_req_opt).await?;
156        println!("Got adapter: {:#?}", adapter.get_info());
157
158        let supported = adapter.limits();
159        let mut limits = wgpu::Limits::downlevel_defaults();
160        limits.max_compute_invocations_per_workgroup =
161            supported.max_compute_invocations_per_workgroup;
162        limits.max_compute_workgroup_storage_size = supported.max_compute_workgroup_storage_size;
163        println!(
164            "Max compute invocations per workgroup: {}",
165            limits.max_compute_invocations_per_workgroup
166        );
167
168        let (device, queue) = adapter
169            .request_device(&wgpu::DeviceDescriptor {
170                label: None,
171                required_features: self.required_features,
172                required_limits: limits,
173                experimental_features: wgpu::ExperimentalFeatures::disabled(),
174                memory_hints: wgpu::MemoryHints::MemoryUsage,
175                trace: wgpu::Trace::Off,
176            })
177            .await?;
178
179        let mut timestamps = None;
180        if self
181            .required_features
182            .contains(wgpu::Features::TIMESTAMP_QUERY)
183        {
184            timestamps = Some(device.create_query_set(&wgpu::QuerySetDescriptor {
185                label: None,
186                ty: wgpu::QueryType::Timestamp,
187                count: 2,
188            }));
189        }
190
191        Ok(Gpu {
192            instance,
193            adapter,
194            device,
195            queue,
196            #[cfg(feature = "window")]
197            window: self.window,
198            #[cfg(feature = "window")]
199            surface: surface,
200            timestamps,
201        })
202    }
203}
204
205pub struct Gpu {
206    pub instance: wgpu::Instance,
207    pub adapter: wgpu::Adapter,
208    pub device: wgpu::Device,
209    pub queue: wgpu::Queue,
210    #[cfg(feature = "window")]
211    pub window: Option<Arc<winit::window::Window>>,
212    #[cfg(feature = "window")]
213    pub surface: Option<wgpu::Surface<'static>>,
214    pub timestamps: Option<wgpu::QuerySet>,
215}
216
217impl Gpu {
218    pub fn buffer<T: bytemuck::Pod>(&self, data: &impl Buffer<T>) -> BufferBuilder<'_> {
219        BufferBuilder::new(data, self)
220    }
221
222    pub fn buffer_by_sz(&self, sz: usize) -> BufferBuilder<'_> {
223        BufferBuilder::new_by_sz(sz, self)
224    }
225}
226
227fn get_entry_point_names(module: &naga::Module) -> anyhow::Result<Vec<String>> {
228    if module.entry_points.is_empty() {
229        anyhow::bail!("Expected at least one entry point, found none.");
230    }
231    Ok(module.entry_points.iter().map(|e| e.name.clone()).collect())
232}
233
234fn get_entry_point_workgroup_size(module: &naga::Module, name: &str) -> anyhow::Result<[u32; 3]> {
235    for entry_point in module.entry_points.iter() {
236        if entry_point.name != name {
237            continue;
238        }
239        return Ok(entry_point.workgroup_size);
240    }
241    anyhow::bail!("Could not find entry point named {}", name);
242}
243
244/*
245fn for_each_global_variable<T>(
246    module: &naga::Module,
247    module_info: &naga::valid::ModuleInfo,
248    mut f: impl FnMut(
249        naga::Handle<GlobalVariable>,
250        &naga::GlobalVariable,
251        &naga::ResourceBinding,
252        &wgpu::ShaderStages,
253        bool,
254    ) -> anyhow::Result<Option<T>>,
255) -> anyhow::Result<Vec<T>> {
256    let mut result = Vec::new();
257    for (global_handle, global) in module.global_variables.iter() {
258        let Some(binding) = &global.binding else {
259            continue;
260        };
261
262        let can_write = match global.space {
263            naga::AddressSpace::Storage { access } => {
264                access.contains(naga::StorageAccess::STORE)
265                    || access.contains(naga::StorageAccess::ATOMIC)
266            }
267            naga::AddressSpace::Uniform => false,
268            _ => {
269                eprintln!(
270                    "Skipping global {:?} with address space {:?}",
271                    global.name, global.space
272                );
273                continue;
274            }
275        };
276        let mut stages = wgpu::ShaderStages::empty();
277        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
278            let usage = module_info.get_entry_point(entry_idx)[global_handle];
279            if usage.is_empty() {
280                continue;
281            }
282
283            stages |= match entry.stage {
284                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
285                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
286                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
287                _ => continue,
288            }
289        }
290        if stages.is_empty() {
291            eprintln!(
292                "Skipping bound global {:?} which is not referenced by any stage",
293                global.name
294            );
295            continue;
296        }
297        if let Some(value) = f(global_handle, global, binding, &stages, can_write)? {
298            result.push(value);
299        }
300    }
301    Ok(result)
302}
303*/
304
305fn get_bind_group_layout(
306    gpu: &Gpu,
307    module: &naga::Module,
308    module_info: &naga::valid::ModuleInfo,
309    module_layout: &naga::proc::Layouter,
310) -> anyhow::Result<wgpu::BindGroupLayout> {
311    let mut entries = Vec::<wgpu::BindGroupLayoutEntry>::new();
312    for (global_handle, global) in module.global_variables.iter() {
313        let Some(binding) = &global.binding else {
314            continue;
315        };
316
317        let can_write = match global.space {
318            naga::AddressSpace::Storage { access } => {
319                access.contains(naga::StorageAccess::STORE)
320                    || access.contains(naga::StorageAccess::ATOMIC)
321            }
322            naga::AddressSpace::Uniform => false,
323            _ => {
324                eprintln!(
325                    "Skipping global {:?} with address space {:?}",
326                    global.name, global.space
327                );
328                continue;
329            }
330        };
331
332        let mut stages = wgpu::ShaderStages::empty();
333        for (entry_idx, entry) in module.entry_points.iter().enumerate() {
334            let usage = module_info.get_entry_point(entry_idx)[global_handle];
335            if usage.is_empty() {
336                continue;
337            }
338
339            stages |= match entry.stage {
340                naga::ShaderStage::Vertex => wgpu::ShaderStages::VERTEX,
341                naga::ShaderStage::Fragment => wgpu::ShaderStages::FRAGMENT,
342                naga::ShaderStage::Compute => wgpu::ShaderStages::COMPUTE,
343                _ => continue,
344            }
345        }
346        if stages.is_empty() {
347            eprintln!(
348                "Skipping bound global {:?} which is not referenced by any stage",
349                global.name
350            );
351            continue;
352        }
353
354        let global_ty = &module.types[global.ty];
355        let entry = match &module.types[global.ty].inner {
356            naga::ir::TypeInner::Array { .. } => {
357                let min_size = module_layout[global.ty].size;
358                wgpu::BindGroupLayoutEntry {
359                    binding: binding.binding,
360                    visibility: stages,
361                    ty: wgpu::BindingType::Buffer {
362                        ty: wgpu::BufferBindingType::Storage {
363                            read_only: !can_write,
364                        },
365                        min_binding_size: NonZeroU64::new(min_size as u64),
366                        has_dynamic_offset: false,
367                    },
368                    count: None,
369                }
370            }
371            naga::ir::TypeInner::Struct { .. } => {
372                let min_size = module_layout[global.ty].size;
373                wgpu::BindGroupLayoutEntry {
374                    binding: binding.binding,
375                    visibility: stages,
376                    ty: wgpu::BindingType::Buffer {
377                        ty: wgpu::BufferBindingType::Uniform,
378                        min_binding_size: NonZeroU64::new(min_size as u64),
379                        has_dynamic_offset: false,
380                    },
381                    count: None,
382                }
383            }
384            _ => {
385                anyhow::bail!(format!(
386                    "Unsupported global type: {:?} (name={:?})",
387                    global_ty.name, global_ty.inner
388                ))
389            }
390        };
391        entries.push(entry);
392    }
393    let layout = wgpu::BindGroupLayoutDescriptor {
394        label: None,
395        entries: &entries,
396    };
397    Ok(gpu.device.create_bind_group_layout(&layout))
398}
399
400pub struct ShaderEntryPoint {
401    pub name: String,
402    pub workgroup_size: [u32; 3],
403}
404
405#[derive(PartialEq, Eq)]
406pub enum VertexSemantic {
407    Position,
408    Normal,
409    Uv0,
410}
411
412pub struct VertexLayout {
413    pub semantics: Vec<VertexSemantic>,
414    pub attributes: Vec<wgpu::VertexAttribute>,
415    pub stride: u64,
416}
417
418// Vertex data is conceived of as zero or more streams of data. We need a way
419// to translate various higher-level types (f32, i32, etc) into these untyped
420// streams, and finally at the end, get an array of them. This builder handles
421// doing exactly that.
422// Example usage:
423//  let positions = [[-0.5_f32, -0.5, 0.5], [0.5, -0.5, 0.5], [0.0, 0.5, 0.5]];
424//  let normals = [[0.0_f32, 0.0, 1.0]; 3];
425//  let uvs = [[0.0_f32, 0.0], [1.0, 0.0], [0.5, 1.0]];
426//  let data_streams = VertexDataStreamBuilder::new()
427//      .stream(VertexSemantic::Position, &positions)
428//      .stream(VertexSemantic::Normal, &normals)
429//      .stream(VertexSemantic::Uv0, &uvs)
430//      .build();
431
432pub struct VertexDataStreamBuilder<'a> {
433    streams: Vec<(VertexSemantic, &'a [u8])>,
434}
435
436impl<'a> VertexDataStreamBuilder<'a> {
437    pub fn new() -> Self {
438        Self {
439            streams: Vec::new(),
440        }
441    }
442
443    pub fn stream<T: NoUninit>(mut self, semantic: VertexSemantic, data: &'a [T]) -> Self {
444        let raw_data: &[u8] = bytemuck::cast_slice(&data);
445        self.streams.push((semantic, raw_data));
446        self
447    }
448
449    pub fn build(self) -> Vec<(VertexSemantic, &'a [u8])> {
450        self.streams
451    }
452}
453
454impl VertexLayout {
455    pub fn vertex_buffer_layout(&self) -> wgpu::VertexBufferLayout<'_> {
456        wgpu::VertexBufferLayout {
457            array_stride: self.stride,
458            step_mode: wgpu::VertexStepMode::Vertex,
459            attributes: &self.attributes,
460        }
461    }
462
463    // Assembles streams of vertex data into an interleaved stream.
464    fn interleave_data(
465        &self,
466        n_vertices: usize,
467        data_streams: Vec<(VertexSemantic, &[u8])>,
468    ) -> anyhow::Result<Vec<u8>> {
469        if data_streams.len() != self.attributes.len() {
470            anyhow::bail!(
471                "Expected {} vertex data streams, got {}",
472                self.attributes.len(),
473                data_streams.len()
474            );
475        }
476        let mut streams: Vec<_> = self
477            .semantics
478            .iter()
479            .zip(&self.attributes)
480            .map(|(semantic, attr)| {
481                let (_, data) = data_streams
482                    .iter()
483                    .find(|(s, _)| s == semantic)
484                    .expect("Vertex shader given unrecognized input semantic");
485                let size = attr.format.size() as usize;
486                let start = attr.offset as usize;
487                (start..start + size, data.chunks_exact(size))
488            })
489            .collect();
490
491        let stride = self.stride as usize;
492        let mut data = vec![0; n_vertices * stride];
493
494        for vertex in data.chunks_exact_mut(stride) {
495            for (range, source) in &mut streams {
496                let value = source.next().expect("Not enough vertex data");
497                vertex[range.clone()].copy_from_slice(value);
498            }
499        }
500
501        Ok(data)
502    }
503
504    pub fn vertex_buffer(
505        &self,
506        gpu: &Gpu,
507        data_streams: Vec<(VertexSemantic, &[u8])>,
508        n_vertices: usize,
509    ) -> anyhow::Result<wgpu::Buffer> {
510        let vertex_data_bytes = self.interleave_data(n_vertices, data_streams)?;
511        Ok(gpu
512            .device
513            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
514                label: None,
515                contents: &vertex_data_bytes,
516                usage: wgpu::BufferUsages::VERTEX,
517            }))
518    }
519}
520
521pub struct Shader {
522    pub bind_group_layout: wgpu::BindGroupLayout,
523    pub naga_module: naga::Module,
524    pub shader_module: wgpu::ShaderModule,
525    pub entry_points: Vec<ShaderEntryPoint>,
526    uses_camera: bool,
527}
528
529impl Shader {
530    pub fn new(gpu: &Gpu, shader: &str) -> anyhow::Result<Self> {
531        let naga_module = naga::front::wgsl::parse_str(shader)?;
532        let module_info = naga::valid::Validator::new(
533            naga::valid::ValidationFlags::all(),
534            naga::valid::Capabilities::all(),
535        )
536        .validate(&naga_module)?;
537        let mut module_layout = naga::proc::Layouter::default();
538        module_layout.update(naga_module.to_ctx())?;
539
540        Self::ensure_compute_compatibility(gpu, &naga_module)?;
541        let bind_group_layout =
542            get_bind_group_layout(gpu, &naga_module, &module_info, &module_layout)?;
543        let uses_camera = naga_module.global_variables.iter().any(|(_, global)| {
544            global.space == naga::AddressSpace::Uniform
545                && global.binding.as_ref().is_some_and(|binding| {
546                    binding.group == 0 && binding.binding == 0
547                })
548        });
549        let entry_point_names = get_entry_point_names(&naga_module)?;
550        let entry_points = entry_point_names
551            .iter()
552            .map(|name| {
553                let workgroup_size = get_entry_point_workgroup_size(&naga_module, name)?;
554                Ok(ShaderEntryPoint {
555                    name: name.clone(),
556                    workgroup_size,
557                })
558            })
559            .collect::<anyhow::Result<Vec<_>>>()?;
560
561        let module_descriptor = wgpu::ShaderModuleDescriptor {
562            label: None,
563            source: wgpu::ShaderSource::Naga(Cow::Owned(naga_module.clone())),
564        };
565        let shader_module = gpu.device.create_shader_module(module_descriptor.clone());
566
567        Ok(Self {
568            bind_group_layout,
569            naga_module: naga_module,
570            shader_module,
571            entry_points,
572            uses_camera,
573        })
574    }
575
576    pub fn uses_camera(&self) -> bool {
577        self.uses_camera
578    }
579
580    fn ensure_compute_compatibility(gpu: &Gpu, module: &naga::Module) -> anyhow::Result<()> {
581        let mut has_compute_stage = false;
582        for entry_point in module.entry_points.iter() {
583            if entry_point.stage == naga::ShaderStage::Compute {
584                has_compute_stage = true;
585            }
586        }
587        if !has_compute_stage {
588            return Ok(());
589        }
590        let downlevel_caps = gpu.adapter.get_downlevel_capabilities();
591        if !downlevel_caps
592            .flags
593            .contains(wgpu::DownlevelFlags::COMPUTE_SHADERS)
594        {
595            anyhow::bail!(
596                "GPU {} does not support compute shaders.",
597                gpu.adapter.get_info().name
598            )
599        }
600        Ok(())
601    }
602
603    // Returns the first shader entry point for the given stage.
604    pub fn get_shader_by_stage(
605        &self,
606        stage: naga::ir::ShaderStage,
607    ) -> anyhow::Result<&naga::EntryPoint> {
608        let Some(shader) = self
609            .naga_module
610            .entry_points
611            .iter()
612            .find(|ep| ep.stage == stage)
613        else {
614            anyhow::bail!("Shader not found for stage {:?}", stage);
615        };
616        Ok(shader)
617    }
618
619    fn vertex_member_to_semantic(&self, name: &str) -> anyhow::Result<VertexSemantic> {
620        // Vertex members might be suffixed with _0 by slangc. Since none of our reserved names have underscores, we
621        // can safely split by underscore.
622        let name = name.split('_').collect::<Vec<_>>()[0];
623        match name {
624            "objPos" => Ok(VertexSemantic::Position),
625            "objNormal" => Ok(VertexSemantic::Normal),
626            "uv0" => Ok(VertexSemantic::Uv0),
627            _ => anyhow::bail!("Unknown vertex member: {}", name),
628        }
629    }
630
631    fn struct_member_to_vertex_format(
632        &self,
633        member: &naga::ir::StructMember,
634    ) -> anyhow::Result<wgpu::VertexFormat> {
635        let inner_type = &self.naga_module.types[member.ty].inner;
636        match inner_type {
637            naga::TypeInner::Vector { size, scalar } => match (scalar.kind, scalar.width) {
638                (naga::ScalarKind::Float, 4) => match size {
639                    naga::VectorSize::Bi => Ok(wgpu::VertexFormat::Float32x2),
640                    naga::VectorSize::Tri => Ok(wgpu::VertexFormat::Float32x3),
641                    naga::VectorSize::Quad => Ok(wgpu::VertexFormat::Float32x4),
642                },
643                _ => anyhow::bail!("Unsupported vertex member vector type: {:?}", inner_type),
644            },
645            _ => anyhow::bail!("Unsupported vertex member type: {:?}", inner_type),
646        }
647    }
648
649    fn struct_member_shader_location(
650        &self,
651        member: &naga::ir::StructMember,
652    ) -> anyhow::Result<wgpu::ShaderLocation> {
653        match &member.binding {
654            Some(naga::Binding::Location { location, .. }) => Ok(*location),
655            // TODO should handle built-in bindings
656            _ => anyhow::bail!("Unsupported vertex member binding: {:?}", member.binding),
657        }
658    }
659
660    pub fn get_vertex_layout(&self) -> anyhow::Result<VertexLayout> {
661        let vertex_shader = self.get_shader_by_stage(naga::ir::ShaderStage::Vertex)?;
662        if vertex_shader.function.arguments.len() != 1 {
663            anyhow::bail!(
664                "Vertex shader must have exactly one argument; got {}",
665                vertex_shader.function.arguments.len()
666            );
667        }
668        let arg = &vertex_shader.function.arguments[0];
669        let arg_ty = &self.naga_module.types[arg.ty];
670        match &arg_ty.inner {
671            naga::ir::TypeInner::Struct { members, .. } => {
672                let semantics = members
673                    .iter()
674                    .map(|m| {
675                        let Some(name) = &m.name else {
676                            // This should never happen
677                            anyhow::bail!("Vertex shader input member missing name");
678                        };
679                        self.vertex_member_to_semantic(name.as_str())
680                    })
681                    .collect::<anyhow::Result<Vec<_>>>()?;
682                let mut acc_offset = 0;
683                let attributes = members
684                    .iter()
685                    .map(|m| -> anyhow::Result<wgpu::VertexAttribute> {
686                        let vertex_fmt = self.struct_member_to_vertex_format(m)?;
687                        let location = self.struct_member_shader_location(m)?;
688                        let offset = acc_offset;
689                        acc_offset += vertex_fmt.size();
690                        Ok(wgpu::VertexAttribute {
691                            format: vertex_fmt,
692                            offset,
693                            shader_location: location,
694                        })
695                    })
696                    .collect::<anyhow::Result<Vec<_>>>()?;
697                let stride = attributes.iter().fold(0, |acc, a| acc + a.format.size());
698                Ok(VertexLayout {
699                    semantics,
700                    attributes,
701                    stride,
702                })
703            }
704            _ => {
705                anyhow::bail!("Argument type must be a struct; got {:?}", arg_ty.inner);
706            }
707        }
708    }
709}
710
711pub struct MaterialPass {
712    pub buffers: Vec<wgpu::Buffer>,
713    pub bind_group: wgpu::BindGroup,
714}
715
716pub struct MaterialPassBuilder {
717    buffers: Vec<wgpu::Buffer>,
718    bind_group: Option<wgpu::BindGroup>,
719}
720
721impl MaterialPassBuilder {
722    pub fn new() -> Self {
723        Self {
724            buffers: Vec::new(),
725            bind_group: None,
726        }
727    }
728
729    pub fn buffers(mut self, buffers: Vec<wgpu::Buffer>) -> Self {
730        self.buffers = buffers;
731        self
732    }
733
734    // Must be run after buffers(), etc.
735    pub fn bind_group(mut self, gpu: &Gpu, shader: &Shader) -> Self {
736        let mut descriptor_entries = Vec::<wgpu::BindGroupEntry>::new();
737        let entries =
738            self.buffers
739                .iter()
740                .enumerate()
741                .fold(Vec::new(), |mut entries, (index, buf)| {
742                    entries.push(wgpu::BindGroupEntry {
743                        binding: index as u32,
744                        resource: buf.as_entire_binding(),
745                    });
746                    entries
747                });
748        descriptor_entries.extend(entries);
749        let descriptor = wgpu::BindGroupDescriptor {
750            label: None,
751            layout: &shader.bind_group_layout,
752            entries: &descriptor_entries,
753        };
754        let bind_group = gpu.device.create_bind_group(&descriptor);
755        self.bind_group = Some(bind_group);
756        self
757    }
758
759    pub fn build(self) -> anyhow::Result<MaterialPass> {
760        if self.bind_group.is_none() {
761            anyhow::bail!("bind_group not set");
762        }
763        Ok(MaterialPass {
764            buffers: self.buffers,
765            bind_group: self.bind_group.unwrap(),
766        })
767    }
768}
769
770pub struct MaterialInstance {
771    pub passes: Vec<MaterialPass>,
772    pub readback: Option<wgpu::Buffer>,
773}
774
775impl MaterialInstance {
776    pub fn new(passes: Vec<MaterialPass>, readback: Option<wgpu::Buffer>) -> Self {
777        Self {
778            passes: passes,
779            readback: readback,
780        }
781    }
782}
783
784// A batch of Material work to be submitted together. This allows us to send a
785// batch of work while another is still processing on the GPU.
786pub struct MaterialBatch {
787    pub instances: Vec<MaterialInstance>,
788    pub ready_tx: Sender<(usize, Result<(), wgpu::BufferAsyncError>)>,
789    pub ready_rx: Receiver<(usize, Result<(), wgpu::BufferAsyncError>)>,
790    // Optional timestamp (ts) buffers for profiling.
791    pub timing: Option<MaterialTiming>,
792    // Whether the batch is currently processing. This is set true as soon as a batch is reserved with get_batch().
793    // The caller is responsible for setting this false
794    pub busy: bool,
795    // How many instances in the batch are currently processing. Set to
796    // batch_size (n_instancces) when the batch is reserved The caller may use
797    // this to track completion as each instance in the batch completes.
798    pub n_busy: usize,
799}
800
801impl MaterialBatch {
802    // Use MaterialBuilder instead. It handles creating batches.
803    fn new(
804        gpu: &Gpu,
805        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
806        n_instances: usize,
807        enable_timing: bool,
808    ) -> Self {
809        let mut instances = Vec::<MaterialInstance>::new();
810        instances.reserve(n_instances);
811        for _ in 0..n_instances {
812            let (passes, readback) = passes_generator();
813            instances.push(MaterialInstance::new(passes, readback));
814        }
815        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
816        let mut timing = None;
817        if enable_timing {
818            std::hint::cold_path();
819            timing = Some(MaterialTiming::new(gpu));
820        }
821        Self {
822            instances: instances,
823            ready_tx: ready_tx,
824            ready_rx: ready_rx,
825            timing: timing,
826            busy: false,
827            n_busy: 0,
828        }
829    }
830
831    pub fn enable_profiling(mut self, gpu: &Gpu) -> Self {
832        self.timing = Some(MaterialTiming::new(gpu));
833        self
834    }
835}
836
837pub struct MaterialTiming {
838    pub ts_buf: wgpu::Buffer,
839    pub ts_readback: wgpu::Buffer,
840    pub ready_tx: Sender<Result<(), wgpu::BufferAsyncError>>,
841    pub ready_rx: Receiver<Result<(), wgpu::BufferAsyncError>>,
842    pub busy: bool,
843}
844
845impl MaterialTiming {
846    pub fn new(gpu: &Gpu) -> Self {
847        let ts_buf = gpu.buffer_by_sz(16).query_resolve().copy_src().build();
848        let ts_readback = gpu.buffer_by_sz(16).copy_dst().map_read().build();
849        let (ready_tx, ready_rx) = std::sync::mpsc::channel();
850        Self {
851            ts_buf: ts_buf,
852            ts_readback: ts_readback,
853            ready_tx: ready_tx,
854            ready_rx: ready_rx,
855            busy: false,
856        }
857    }
858}
859
860pub struct Material {
861    pub batches: Vec<MaterialBatch>,
862    pub batch_size: usize,
863}
864
865impl Material {
866    pub fn get_batch(&mut self) -> Option<&mut MaterialBatch> {
867        let batch = self.batches.iter_mut().find(|batch| {
868            if batch.busy {
869                return false;
870            }
871            if let Some(timing) = &batch.timing {
872                std::hint::cold_path();
873                if timing.busy {
874                    return false;
875                }
876            }
877            true
878        })?;
879        // Found a batch.
880        batch.busy = true;
881        if let Some(timing) = &mut batch.timing {
882            std::hint::cold_path();
883            timing.busy = true;
884        }
885        batch.n_busy = self.batch_size;
886        Some(batch)
887    }
888}
889
890pub struct MaterialBuilder {
891    batch_size: usize, // instances per batch
892    n_batches: usize,
893    enable_timing: bool,
894}
895
896impl MaterialBuilder {
897    pub fn new() -> Self {
898        Self {
899            batch_size: 1,
900            n_batches: 1,
901            enable_timing: false,
902        }
903    }
904
905    pub fn batch_size(mut self, batch_size: usize) -> Self {
906        self.batch_size = batch_size;
907        self
908    }
909
910    pub fn n_batches(mut self, n_batches: usize) -> Self {
911        self.n_batches = n_batches;
912        self
913    }
914
915    pub fn enable_timing(mut self) -> Self {
916        self.enable_timing = true;
917        self
918    }
919
920    pub fn build(
921        self,
922        gpu: &Gpu,
923        mut passes_generator: impl FnMut() -> (Vec<MaterialPass>, Option<wgpu::Buffer>),
924    ) -> Material {
925        let mut batches = Vec::<MaterialBatch>::new();
926        batches.reserve(self.n_batches);
927        for _ in 0..self.n_batches {
928            batches.push(MaterialBatch::new(
929                gpu,
930                &mut passes_generator,
931                self.batch_size,
932                self.enable_timing,
933            ));
934        }
935        Material {
936            batches: batches,
937            batch_size: self.batch_size,
938        }
939    }
940}
941
942pub struct PipelineBuilder<'a> {
943    shader: &'a Shader,
944    zero_initialize: bool,
945    color_target: Option<wgpu::ColorTargetState>,
946    vertex_buffer_layout: Option<wgpu::VertexBufferLayout<'a>>,
947}
948
949impl<'a> PipelineBuilder<'a> {
950    pub fn new(shader: &'a Shader) -> Self {
951        Self {
952            shader,
953            zero_initialize: true,
954            color_target: None,
955            vertex_buffer_layout: None,
956        }
957    }
958
959    pub fn zero_initialize(mut self, zero_initialize: bool) -> Self {
960        self.zero_initialize = zero_initialize;
961        self
962    }
963
964    pub fn color_target(mut self, color_target: wgpu::ColorTargetState) -> Self {
965        self.color_target = Some(color_target);
966        self
967    }
968
969    pub fn vertex_buffer_layout(
970        mut self,
971        vertex_buffer_layout: wgpu::VertexBufferLayout<'a>,
972    ) -> Self {
973        self.vertex_buffer_layout = Some(vertex_buffer_layout);
974        self
975    }
976
977    // Build a ComputePipeline. Assumes the shader has exactly one compute shader.
978    pub fn build_compute(
979        self,
980        gpu: &Gpu,
981        entry_point: &str,
982    ) -> anyhow::Result<wgpu::ComputePipeline> {
983        let pipeline_layout = gpu
984            .device
985            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
986                label: None,
987                bind_group_layouts: &[Some(&self.shader.bind_group_layout)],
988                immediate_size: 0,
989            });
990
991        if !self
992            .shader
993            .naga_module
994            .entry_points
995            .iter()
996            .any(|ep| ep.name == entry_point && ep.stage == naga::ir::ShaderStage::Compute)
997        {
998            anyhow::bail!(
999                "Shader entry point '{}' with compute type not found.",
1000                entry_point
1001            );
1002        }
1003        let compilation_options = wgpu::PipelineCompilationOptions {
1004            zero_initialize_workgroup_memory: self.zero_initialize,
1005            ..Default::default()
1006        };
1007        Ok(gpu
1008            .device
1009            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
1010                label: None,
1011                layout: Some(&pipeline_layout),
1012                module: &self.shader.shader_module,
1013                entry_point: Some(entry_point),
1014                compilation_options: compilation_options,
1015                cache: None,
1016            }))
1017    }
1018
1019    // Build a RenderPipeline. Assumes the shader has exactly one vertex and one fragment shader.
1020    pub fn build_render(self, gpu: &Gpu) -> anyhow::Result<wgpu::RenderPipeline> {
1021        let pipeline_layout = gpu
1022            .device
1023            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
1024                label: None,
1025                bind_group_layouts: &[Some(&self.shader.bind_group_layout)],
1026                immediate_size: 0,
1027            });
1028        let vertex_shader = self
1029            .shader
1030            .get_shader_by_stage(naga::ir::ShaderStage::Vertex)?;
1031        let fragment_shader = self
1032            .shader
1033            .get_shader_by_stage(naga::ir::ShaderStage::Fragment)?;
1034        let compilation_options = wgpu::PipelineCompilationOptions {
1035            zero_initialize_workgroup_memory: self.zero_initialize,
1036            ..Default::default()
1037        };
1038        let Some(color_target) = self.color_target else {
1039            anyhow::bail!("color_target is not set");
1040        };
1041        let Some(vertex_buffer_layout) = self.vertex_buffer_layout else {
1042            anyhow::bail!("vertex_buffer_layout is not set");
1043        };
1044        Ok(gpu
1045            .device
1046            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
1047                label: None,
1048                layout: Some(&pipeline_layout),
1049                vertex: wgpu::VertexState {
1050                    module: &self.shader.shader_module,
1051                    entry_point: Some(&vertex_shader.name),
1052                    compilation_options: compilation_options.clone(),
1053                    buffers: &vec![Some(vertex_buffer_layout)],
1054                },
1055                fragment: Some(wgpu::FragmentState {
1056                    module: &self.shader.shader_module,
1057                    entry_point: Some(&fragment_shader.name),
1058                    targets: &vec![Some(color_target)],
1059                    compilation_options: compilation_options,
1060                }),
1061                primitive: wgpu::PrimitiveState::default(),
1062                depth_stencil: None,
1063                multisample: wgpu::MultisampleState::default(),
1064                multiview_mask: None,
1065                cache: None,
1066            }))
1067    }
1068}
1069
1070#[derive(Default)]
1071pub struct Transform {
1072    pub position: [f32; 3],
1073    pub scale: [f32; 3],
1074    // TODO abstract to support euler angles and quaternions
1075    pub orientation: [f32; 3],
1076}
1077
1078pub struct Object {
1079    pub transform: Transform,
1080    pub vertex_buffer: wgpu::Buffer,
1081    // Could hoist this into Scene to have more than one mesh per pipeline, but
1082    // this is good enough for prototype.
1083    pub pipeline: wgpu::RenderPipeline,
1084    pub camera_bind_group: Option<wgpu::BindGroup>,
1085}
1086
1087pub struct ObjectBuilder<'a> {
1088    transform: Transform,
1089    vertex_buffer: Option<wgpu::Buffer>,
1090    pipeline: Option<wgpu::RenderPipeline>,
1091    camera: Option<&'a Camera>,
1092}
1093
1094impl<'a> ObjectBuilder<'a> {
1095    pub fn new() -> Self {
1096        Self {
1097            transform: Transform::default(),
1098            vertex_buffer: None,
1099            pipeline: None,
1100            camera: None,
1101        }
1102    }
1103
1104    pub fn transform(mut self, transform: Transform) -> Self {
1105        self.transform = transform;
1106        self
1107    }
1108
1109    pub fn vertex_buffer(mut self, vertex_buffer: wgpu::Buffer) -> Self {
1110        self.vertex_buffer = Some(vertex_buffer);
1111        self
1112    }
1113
1114    pub fn pipeline(mut self, pipeline: wgpu::RenderPipeline) -> Self {
1115        self.pipeline = Some(pipeline);
1116        self
1117    }
1118
1119    pub fn camera(mut self, camera: &'a Camera) -> Self {
1120        self.camera = Some(camera);
1121        self
1122    }
1123
1124    pub fn build(self, gpu: &Gpu, shader: &Shader) -> anyhow::Result<Object> {
1125        let Some(vertex_buffer) = self.vertex_buffer else {
1126            anyhow::bail!("Cannot build Object without vertex_buffer");
1127        };
1128        let Some(pipeline) = self.pipeline else {
1129            anyhow::bail!("Cannot build Object without pipeline");
1130        };
1131        let camera_bind_group = if shader.uses_camera() {
1132            Some(self.camera.context("Shader requires a camera")?.bind_group(gpu, shader))
1133        } else {
1134            None
1135        };
1136        Ok(Object {
1137            transform: self.transform,
1138            vertex_buffer,
1139            pipeline,
1140            camera_bind_group,
1141        })
1142    }
1143}
1144
1145#[derive(Builder)]
1146#[builder(build_fn(skip))]
1147pub struct Camera {
1148    pub eye: glam::Vec3,
1149    pub target: glam::Vec3,
1150    pub up: glam::Vec3,
1151    pub vertical_fov: f32,
1152    pub near: f32,
1153    pub far: f32,
1154    #[builder(setter(custom))]
1155    pub buffer: wgpu::Buffer,
1156}
1157
1158impl CameraBuilder {
1159    pub fn build(&self, gpu: &Gpu, aspect: f32) -> anyhow::Result<Camera> {
1160        let buffer = gpu.buffer_by_sz(64).uniform().copy_dst().build();
1161        let camera = Camera {
1162            eye: self.eye.context("eye is not set")?,
1163            target: self.target.context("target is not set")?,
1164            up: self.up.context("up is not set")?,
1165            vertical_fov: self.vertical_fov.context("vertical_fov is not set")?,
1166            near: self.near.context("near is not set")?,
1167            far: self.far.context("far is not set")?,
1168            buffer,
1169        };
1170        camera.update(gpu, aspect);
1171        Ok(camera)
1172    }
1173}
1174
1175impl Camera {
1176    fn bind_group(&self, gpu: &Gpu, shader: &Shader) -> wgpu::BindGroup {
1177        gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
1178            label: Some("Object camera bindings"),
1179            layout: &shader.bind_group_layout,
1180            entries: &[wgpu::BindGroupEntry {
1181                binding: 0,
1182                resource: self.buffer.as_entire_binding(),
1183            }],
1184        })
1185    }
1186
1187    pub fn update(&self, gpu: &Gpu, aspect: f32) {
1188        let matrix = self.view_projection(aspect).to_cols_array_2d();
1189        gpu.queue.write_buffer(&self.buffer, 0, bytemuck::bytes_of(&matrix));
1190    }
1191
1192    fn view_projection(&self, aspect: f32) -> glam::Mat4 {
1193        // right x, up y. right handed -> near z
1194        use glam::camera::rh::{proj, view};
1195        let view = view::look_at_mat4(self.eye, self.target, self.up);
1196        // 0-1 depth.
1197        let projection = proj::directx::perspective(self.vertical_fov, aspect, self.near, self.far);
1198        projection * view
1199    }
1200}
1201
1202pub struct Scene {
1203    pub camera: Camera,
1204    pub objects: Vec<Object>,
1205}
1206
1207pub struct SceneBuilder {
1208    camera: Option<Camera>,
1209    objects: Vec<Object>,
1210}
1211
1212impl SceneBuilder {
1213    pub fn new() -> Self {
1214        Self {
1215            camera: None,
1216            objects: Vec::new(),
1217        }
1218    }
1219
1220    pub fn camera(mut self, camera: Camera) -> Self {
1221        self.camera = Some(camera);
1222        self
1223    }
1224
1225    pub fn object(mut self, object: Object) -> Self {
1226        self.objects.push(object);
1227        self
1228    }
1229
1230    pub fn build(self) -> anyhow::Result<Scene> {
1231        let Some(camera) = self.camera else {
1232            anyhow::bail!("camera is not set");
1233        };
1234        Ok(Scene {
1235            camera,
1236            objects: self.objects,
1237        })
1238    }
1239}