Skip to main content

astrelis_render/
types.rs

1//! Typed GPU resource wrappers.
2//!
3//! This module provides lightweight wrappers around wgpu types that add
4//! type safety and metadata tracking without fully encapsulating wgpu.
5//!
6//! # Types
7//!
8//! - [`TypedBuffer`]: A buffer that tracks its element type, length, and usage
9//! - [`GpuTexture`]: A texture with cached view and metadata
10//! - [`StorageTexture`]: A texture for compute shader read/write access
11//!
12//! # Example
13//!
14//! ```ignore
15//! use astrelis_render::{GraphicsContext, TypedBuffer};
16//!
17//! let ctx = GraphicsContext::new_owned_sync().expect("Failed to create graphics context");
18//!
19//! // Create a typed buffer of f32 values
20//! let data = [1.0f32, 2.0, 3.0, 4.0];
21//! let buffer = TypedBuffer::new(
22//!     ctx.device(),
23//!     Some("My Buffer"),
24//!     &data,
25//!     wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
26//! );
27//!
28//! // Later, update the buffer
29//! buffer.write(ctx.queue(), &[5.0, 6.0, 7.0, 8.0]);
30//! ```
31
32use std::marker::PhantomData;
33
34use crate::extension::{AsWgpu, IntoWgpu};
35
36// =============================================================================
37// TypedBuffer
38// =============================================================================
39
40/// A GPU buffer with type-safe element tracking.
41///
42/// This wrapper adds:
43/// - Type safety via generics
44/// - Automatic length tracking
45/// - Convenient write operations
46///
47/// The underlying buffer is directly accessible via the `AsWgpu` trait.
48pub struct TypedBuffer<T: bytemuck::Pod> {
49    buffer: wgpu::Buffer,
50    len: u32,
51    usage: wgpu::BufferUsages,
52    _marker: PhantomData<T>,
53}
54
55impl<T: bytemuck::Pod> TypedBuffer<T> {
56    /// Create a new typed buffer with initial data.
57    ///
58    /// # Arguments
59    ///
60    /// * `device` - The wgpu device to create the buffer on
61    /// * `label` - Optional debug label for the buffer
62    /// * `data` - Initial data to populate the buffer with
63    /// * `usage` - Buffer usage flags
64    pub fn new(
65        device: &wgpu::Device,
66        label: Option<&str>,
67        data: &[T],
68        usage: wgpu::BufferUsages,
69    ) -> Self {
70        use wgpu::util::DeviceExt;
71
72        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
73            label,
74            contents: bytemuck::cast_slice(data),
75            usage,
76        });
77
78        Self {
79            buffer,
80            len: data.len() as u32,
81            usage,
82            _marker: PhantomData,
83        }
84    }
85
86    /// Create an empty typed buffer with a given capacity.
87    ///
88    /// # Arguments
89    ///
90    /// * `device` - The wgpu device to create the buffer on
91    /// * `label` - Optional debug label for the buffer
92    /// * `capacity` - Number of elements the buffer can hold
93    /// * `usage` - Buffer usage flags
94    pub fn with_capacity(
95        device: &wgpu::Device,
96        label: Option<&str>,
97        capacity: u32,
98        usage: wgpu::BufferUsages,
99    ) -> Self {
100        let size = (capacity as usize * std::mem::size_of::<T>()) as u64;
101
102        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
103            label,
104            size,
105            usage,
106            mapped_at_creation: false,
107        });
108
109        Self {
110            buffer,
111            len: 0,
112            usage,
113            _marker: PhantomData,
114        }
115    }
116
117    /// Get the number of elements in the buffer.
118    #[inline]
119    pub fn len(&self) -> u32 {
120        self.len
121    }
122
123    /// Check if the buffer is empty.
124    #[inline]
125    pub fn is_empty(&self) -> bool {
126        self.len == 0
127    }
128
129    /// Get the size of the buffer in bytes.
130    #[inline]
131    pub fn size(&self) -> u64 {
132        self.len as u64 * std::mem::size_of::<T>() as u64
133    }
134
135    /// Get the capacity of the buffer in bytes.
136    #[inline]
137    pub fn capacity_bytes(&self) -> u64 {
138        self.buffer.size()
139    }
140
141    /// Get the capacity in number of elements.
142    #[inline]
143    pub fn capacity(&self) -> u32 {
144        (self.buffer.size() / std::mem::size_of::<T>() as u64) as u32
145    }
146
147    /// Get the buffer usage flags.
148    #[inline]
149    pub fn usage(&self) -> wgpu::BufferUsages {
150        self.usage
151    }
152
153    /// Get a slice of the entire buffer.
154    #[inline]
155    pub fn slice(&self) -> wgpu::BufferSlice<'_> {
156        self.buffer.slice(..)
157    }
158
159    /// Get a slice of a portion of the buffer.
160    #[inline]
161    pub fn slice_range(&self, range: std::ops::Range<u32>) -> wgpu::BufferSlice<'_> {
162        let start = range.start as u64 * std::mem::size_of::<T>() as u64;
163        let end = range.end as u64 * std::mem::size_of::<T>() as u64;
164        self.buffer.slice(start..end)
165    }
166
167    /// Write data to the buffer.
168    ///
169    /// The buffer must have been created with `COPY_DST` usage.
170    pub fn write(&self, queue: &wgpu::Queue, data: &[T]) {
171        queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(data));
172    }
173
174    /// Write data to the buffer at an offset.
175    ///
176    /// The buffer must have been created with `COPY_DST` usage.
177    pub fn write_at(&self, queue: &wgpu::Queue, offset: u32, data: &[T]) {
178        let byte_offset = offset as u64 * std::mem::size_of::<T>() as u64;
179        queue.write_buffer(&self.buffer, byte_offset, bytemuck::cast_slice(data));
180    }
181
182    /// Get the buffer as a binding resource (for bind groups).
183    #[inline]
184    pub fn as_binding(&self) -> wgpu::BindingResource<'_> {
185        self.buffer.as_entire_binding()
186    }
187
188    /// Get a reference to the underlying buffer.
189    #[inline]
190    pub fn buffer(&self) -> &wgpu::Buffer {
191        &self.buffer
192    }
193}
194
195impl<T: bytemuck::Pod> AsWgpu for TypedBuffer<T> {
196    type WgpuType = wgpu::Buffer;
197
198    fn as_wgpu(&self) -> &Self::WgpuType {
199        &self.buffer
200    }
201}
202
203impl<T: bytemuck::Pod> IntoWgpu for TypedBuffer<T> {
204    type WgpuType = wgpu::Buffer;
205
206    fn into_wgpu(self) -> Self::WgpuType {
207        self.buffer
208    }
209}
210
211// =============================================================================
212// GpuTexture
213// =============================================================================
214
215/// A GPU texture with cached view and metadata.
216///
217/// This wrapper provides:
218/// - Automatic view creation and caching
219/// - Size and format metadata
220/// - Convenient accessors
221pub struct GpuTexture {
222    texture: wgpu::Texture,
223    view: wgpu::TextureView,
224    size: wgpu::Extent3d,
225    format: wgpu::TextureFormat,
226    sample_count: u32,
227}
228
229impl GpuTexture {
230    /// Create a new GPU texture.
231    pub fn new(device: &wgpu::Device, descriptor: &wgpu::TextureDescriptor) -> Self {
232        let texture = device.create_texture(descriptor);
233        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
234
235        Self {
236            texture,
237            view,
238            size: descriptor.size,
239            format: descriptor.format,
240            sample_count: descriptor.sample_count,
241        }
242    }
243
244    /// Create a simple 2D texture.
245    pub fn new_2d(
246        device: &wgpu::Device,
247        label: Option<&str>,
248        width: u32,
249        height: u32,
250        format: wgpu::TextureFormat,
251        usage: wgpu::TextureUsages,
252    ) -> Self {
253        Self::new(
254            device,
255            &wgpu::TextureDescriptor {
256                label,
257                size: wgpu::Extent3d {
258                    width,
259                    height,
260                    depth_or_array_layers: 1,
261                },
262                mip_level_count: 1,
263                sample_count: 1,
264                dimension: wgpu::TextureDimension::D2,
265                format,
266                usage,
267                view_formats: &[],
268            },
269        )
270    }
271
272    /// Create a texture from raw data.
273    pub fn from_data(
274        device: &wgpu::Device,
275        queue: &wgpu::Queue,
276        label: Option<&str>,
277        width: u32,
278        height: u32,
279        format: wgpu::TextureFormat,
280        data: &[u8],
281    ) -> Self {
282        let texture = Self::new_2d(
283            device,
284            label,
285            width,
286            height,
287            format,
288            wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
289        );
290
291        queue.write_texture(
292            wgpu::TexelCopyTextureInfo {
293                texture: &texture.texture,
294                mip_level: 0,
295                origin: wgpu::Origin3d::ZERO,
296                aspect: wgpu::TextureAspect::All,
297            },
298            data,
299            wgpu::TexelCopyBufferLayout {
300                offset: 0,
301                bytes_per_row: Some(width * format.block_copy_size(None).unwrap_or(4)),
302                rows_per_image: Some(height),
303            },
304            wgpu::Extent3d {
305                width,
306                height,
307                depth_or_array_layers: 1,
308            },
309        );
310
311        texture
312    }
313
314    /// Get the texture view.
315    #[inline]
316    pub fn view(&self) -> &wgpu::TextureView {
317        &self.view
318    }
319
320    /// Get the texture size.
321    #[inline]
322    pub fn size(&self) -> wgpu::Extent3d {
323        self.size
324    }
325
326    /// Get the texture width.
327    #[inline]
328    pub fn width(&self) -> u32 {
329        self.size.width
330    }
331
332    /// Get the texture height.
333    #[inline]
334    pub fn height(&self) -> u32 {
335        self.size.height
336    }
337
338    /// Get the texture format.
339    #[inline]
340    pub fn format(&self) -> wgpu::TextureFormat {
341        self.format
342    }
343
344    /// Get the sample count.
345    #[inline]
346    pub fn sample_count(&self) -> u32 {
347        self.sample_count
348    }
349
350    /// Get the texture as a binding resource.
351    #[inline]
352    pub fn as_binding(&self) -> wgpu::BindingResource<'_> {
353        wgpu::BindingResource::TextureView(&self.view)
354    }
355
356    /// Create a custom view with different parameters.
357    pub fn create_view(&self, descriptor: &wgpu::TextureViewDescriptor) -> wgpu::TextureView {
358        self.texture.create_view(descriptor)
359    }
360}
361
362impl AsWgpu for GpuTexture {
363    type WgpuType = wgpu::Texture;
364
365    fn as_wgpu(&self) -> &Self::WgpuType {
366        &self.texture
367    }
368}
369
370impl IntoWgpu for GpuTexture {
371    type WgpuType = wgpu::Texture;
372
373    fn into_wgpu(self) -> Self::WgpuType {
374        self.texture
375    }
376}
377
378// =============================================================================
379// StorageTexture
380// =============================================================================
381
382/// Access mode for storage textures in compute shaders.
383#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
384pub enum StorageTextureAccess {
385    /// Read-only access in shaders.
386    ReadOnly,
387    /// Write-only access in shaders.
388    WriteOnly,
389    /// Read-write access in shaders.
390    ReadWrite,
391}
392
393impl StorageTextureAccess {
394    /// Convert to wgpu storage texture access.
395    pub fn to_wgpu(self) -> wgpu::StorageTextureAccess {
396        match self {
397            Self::ReadOnly => wgpu::StorageTextureAccess::ReadOnly,
398            Self::WriteOnly => wgpu::StorageTextureAccess::WriteOnly,
399            Self::ReadWrite => wgpu::StorageTextureAccess::ReadWrite,
400        }
401    }
402}
403
404/// A texture for compute shader read/write access.
405///
406/// Storage textures are used when compute shaders need to write
407/// to or read from textures directly.
408pub struct StorageTexture {
409    texture: wgpu::Texture,
410    view: wgpu::TextureView,
411    size: wgpu::Extent3d,
412    format: wgpu::TextureFormat,
413    access: StorageTextureAccess,
414}
415
416impl StorageTexture {
417    /// Create a new storage texture.
418    pub fn new(
419        device: &wgpu::Device,
420        label: Option<&str>,
421        width: u32,
422        height: u32,
423        format: wgpu::TextureFormat,
424        access: StorageTextureAccess,
425    ) -> Self {
426        let size = wgpu::Extent3d {
427            width,
428            height,
429            depth_or_array_layers: 1,
430        };
431
432        let texture = device.create_texture(&wgpu::TextureDescriptor {
433            label,
434            size,
435            mip_level_count: 1,
436            sample_count: 1,
437            dimension: wgpu::TextureDimension::D2,
438            format,
439            usage: wgpu::TextureUsages::STORAGE_BINDING
440                | wgpu::TextureUsages::COPY_SRC
441                | wgpu::TextureUsages::COPY_DST,
442            view_formats: &[],
443        });
444
445        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
446
447        Self {
448            texture,
449            view,
450            size,
451            format,
452            access,
453        }
454    }
455
456    /// Get the texture view.
457    #[inline]
458    pub fn view(&self) -> &wgpu::TextureView {
459        &self.view
460    }
461
462    /// Get the texture size.
463    #[inline]
464    pub fn size(&self) -> wgpu::Extent3d {
465        self.size
466    }
467
468    /// Get the texture width.
469    #[inline]
470    pub fn width(&self) -> u32 {
471        self.size.width
472    }
473
474    /// Get the texture height.
475    #[inline]
476    pub fn height(&self) -> u32 {
477        self.size.height
478    }
479
480    /// Get the texture format.
481    #[inline]
482    pub fn format(&self) -> wgpu::TextureFormat {
483        self.format
484    }
485
486    /// Get the access mode.
487    #[inline]
488    pub fn access(&self) -> StorageTextureAccess {
489        self.access
490    }
491
492    /// Get the texture as a binding resource.
493    #[inline]
494    pub fn as_binding(&self) -> wgpu::BindingResource<'_> {
495        wgpu::BindingResource::TextureView(&self.view)
496    }
497}
498
499impl AsWgpu for StorageTexture {
500    type WgpuType = wgpu::Texture;
501
502    fn as_wgpu(&self) -> &Self::WgpuType {
503        &self.texture
504    }
505}
506
507impl IntoWgpu for StorageTexture {
508    type WgpuType = wgpu::Texture;
509
510    fn into_wgpu(self) -> Self::WgpuType {
511        self.texture
512    }
513}
514
515// =============================================================================
516// UniformBuffer (Convenience type)
517// =============================================================================
518
519/// A uniform buffer for shader uniforms.
520///
521/// This is a convenience type for buffers that contain uniform data.
522pub type UniformBuffer<T> = TypedBuffer<T>;
523
524impl<T: bytemuck::Pod> TypedBuffer<T> {
525    /// Create a new uniform buffer.
526    pub fn new_uniform(device: &wgpu::Device, label: Option<&str>, data: &T) -> Self {
527        use wgpu::util::DeviceExt;
528
529        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
530            label,
531            contents: bytemuck::bytes_of(data),
532            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
533        });
534
535        Self {
536            buffer,
537            len: 1,
538            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
539            _marker: PhantomData,
540        }
541    }
542
543    /// Write a single uniform value to the buffer.
544    pub fn write_uniform(&self, queue: &wgpu::Queue, data: &T) {
545        queue.write_buffer(&self.buffer, 0, bytemuck::bytes_of(data));
546    }
547}
548
549#[cfg(test)]
550mod tests {
551    use super::*;
552
553    #[repr(C)]
554    #[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
555    struct TestData {
556        x: f32,
557        y: f32,
558        z: f32,
559        w: f32,
560    }
561
562    #[test]
563    fn test_typed_buffer_size() {
564        // Test that size calculations are correct
565        assert_eq!(std::mem::size_of::<TestData>(), 16);
566        assert_eq!(std::mem::size_of::<f32>(), 4);
567    }
568
569    #[test]
570    fn test_storage_texture_access_conversion() {
571        assert_eq!(
572            StorageTextureAccess::ReadOnly.to_wgpu(),
573            wgpu::StorageTextureAccess::ReadOnly
574        );
575        assert_eq!(
576            StorageTextureAccess::WriteOnly.to_wgpu(),
577            wgpu::StorageTextureAccess::WriteOnly
578        );
579        assert_eq!(
580            StorageTextureAccess::ReadWrite.to_wgpu(),
581            wgpu::StorageTextureAccess::ReadWrite
582        );
583    }
584}