librashader_runtime/
uniforms.rs

1use librashader_reflect::reflect::semantics::{MemberOffset, UniformMemberBlock};
2use std::marker::PhantomData;
3use std::ops::{Deref, DerefMut};
4
5/// A scalar value that is valid as a uniform member
6pub trait UniformScalar: Copy + bytemuck::Pod {}
7impl UniformScalar for f32 {}
8impl UniformScalar for i32 {}
9impl UniformScalar for u32 {}
10
11/// A trait for a binder that binds the given value and context into the uniform for a shader pass.
12pub trait BindUniform<C, T, D> {
13    /// Bind the given value to the shader uniforms given the input context.
14    ///
15    /// A `BindUniform` implementation should not write to a backing buffer from a [`UniformStorage`].
16    /// If the binding is successful and no writes to a backing buffer is necessary, this function should return `Some(())`.
17    /// If this function returns `None`, then the value will instead be written to the backing buffer.
18    fn bind_uniform(block: UniformMemberBlock, value: T, ctx: C, device: &D) -> Option<()>;
19}
20
21/// A trait to access the raw pointer to a backing uniform storage.
22pub trait UniformStorageAccess {
23    /// Get a pointer to the backing UBO storage. This pointer must be valid for the lifetime
24    /// of the implementing struct.
25    fn ubo_pointer(&self) -> *const u8;
26
27    /// Get a pointer to the backing UBO storage. This pointer must be valid for the lifetime
28    /// of the implementing struct.
29    fn ubo_slice(&self) -> &[u8];
30
31    /// Get a pointer to the backing Push Constant buffer storage.
32    /// This pointer must be valid for the lifetime of the implementing struct.
33    fn push_pointer(&self) -> *const u8;
34
35    /// Get a slice to the backing Push Constant buffer storage.
36    /// This pointer must be valid for the lifetime of the implementing struct.
37    fn push_slice(&self) -> &[u8];
38}
39
40impl<D, T, H, U, P> UniformStorageAccess for UniformStorage<T, H, U, P, D>
41where
42    U: Deref<Target = [u8]> + DerefMut,
43    P: Deref<Target = [u8]> + DerefMut,
44{
45    fn ubo_pointer(&self) -> *const u8 {
46        self.ubo.as_ptr()
47    }
48
49    fn ubo_slice(&self) -> &[u8] {
50        &self.ubo
51    }
52
53    fn push_pointer(&self) -> *const u8 {
54        self.push.as_ptr()
55    }
56
57    fn push_slice(&self) -> &[u8] {
58        &self.push
59    }
60}
61
62/// A uniform binder that always returns `None`, and does not do any binding of uniforms.
63/// All uniform data is thus written into the backing buffer storage.
64pub struct NoUniformBinder;
65impl<T, D> BindUniform<Option<()>, T, D> for NoUniformBinder {
66    fn bind_uniform(_: UniformMemberBlock, _: T, _: Option<()>, _: &D) -> Option<()> {
67        None
68    }
69}
70
71/// A helper to bind uniform variables to UBO or Push Constant Buffers.
72pub struct UniformStorage<H = NoUniformBinder, C = Option<()>, U = Box<[u8]>, P = Box<[u8]>, D = ()>
73where
74    U: Deref<Target = [u8]> + DerefMut,
75    P: Deref<Target = [u8]> + DerefMut,
76{
77    ubo: U,
78    push: P,
79    _h: PhantomData<H>,
80    _c: PhantomData<C>,
81    _d: PhantomData<D>,
82}
83
84impl<H, C, U, P, D> UniformStorage<H, C, U, P, D>
85where
86    U: Deref<Target = [u8]> + DerefMut,
87    P: Deref<Target = [u8]> + DerefMut,
88{
89    /// Access the backing storage for the UBO.
90    pub fn inner_ubo(&self) -> &U {
91        &self.ubo
92    }
93
94    /// Access the backing storage for the Push storage.
95    pub fn inner_push(&self) -> &P {
96        &self.push
97    }
98
99    pub(crate) fn buffer(&mut self, ty: UniformMemberBlock) -> &mut [u8] {
100        match ty {
101            UniformMemberBlock::Ubo => self.ubo.deref_mut(),
102            UniformMemberBlock::PushConstant => self.push.deref_mut(),
103        }
104    }
105}
106
107impl<H, C, U, P, D> UniformStorage<H, C, U, P, D>
108where
109    C: Copy,
110    U: Deref<Target = [u8]> + DerefMut,
111    P: Deref<Target = [u8]> + DerefMut,
112{
113    #[inline(always)]
114    fn write_scalar_inner<T: UniformScalar>(buffer: &mut [u8], value: T) {
115        let buffer = bytemuck::cast_slice_mut(buffer);
116        buffer[0] = value;
117    }
118
119    /// Bind a scalar to the given offset.
120    #[inline(always)]
121    pub fn bind_scalar<T: UniformScalar>(
122        &mut self,
123        offset: MemberOffset,
124        value: T,
125        ctx: C,
126        device: &D,
127    ) where
128        H: BindUniform<C, T, D>,
129    {
130        for ty in UniformMemberBlock::TYPES {
131            if H::bind_uniform(ty, value, ctx, device).is_some() {
132                continue;
133            }
134
135            if let Some(offset) = offset.offset(ty) {
136                let buffer = self.buffer(ty);
137                Self::write_scalar_inner(&mut buffer[offset..][..std::mem::size_of::<T>()], value)
138            }
139        }
140    }
141
142    /// Create a new `UniformStorage` with the given backing storage
143    pub fn new_with_storage(ubo: U, push: P) -> UniformStorage<H, C, U, P, D> {
144        UniformStorage {
145            ubo,
146            push,
147            _h: Default::default(),
148            _c: Default::default(),
149            _d: Default::default(),
150        }
151    }
152}
153
154impl<H, C, U, D> UniformStorage<H, C, U, Box<[u8]>, D>
155where
156    C: Copy,
157    U: Deref<Target = [u8]> + DerefMut,
158{
159    /// Create a new `UniformStorage` with the given backing storage
160    pub fn new_with_ubo_storage(
161        storage: U,
162        push_size: usize,
163    ) -> UniformStorage<H, C, U, Box<[u8]>, D> {
164        UniformStorage {
165            ubo: storage,
166            push: vec![0u8; push_size].into_boxed_slice(),
167            _h: Default::default(),
168            _c: Default::default(),
169            _d: Default::default(),
170        }
171    }
172}
173
174impl<H, C, D> UniformStorage<H, C, Box<[u8]>, Box<[u8]>, D> {
175    /// Create a new `UniformStorage` with the given size for UBO and Push Constant Buffer sizes.
176    pub fn new(ubo_size: usize, push_size: usize) -> UniformStorage<H, C, Box<[u8]>, Box<[u8]>, D> {
177        UniformStorage {
178            ubo: vec![0u8; ubo_size].into_boxed_slice(),
179            push: vec![0u8; push_size].into_boxed_slice(),
180            _h: Default::default(),
181            _c: Default::default(),
182            _d: Default::default(),
183        }
184    }
185}
186
187impl<H, C, U, P, D> UniformStorage<H, C, U, P, D>
188where
189    C: Copy,
190    U: Deref<Target = [u8]> + DerefMut,
191    P: Deref<Target = [u8]> + DerefMut,
192    H: for<'a> BindUniform<C, &'a [f32; 4], D>,
193{
194    #[inline(always)]
195    fn write_vec4_inner(buffer: &mut [u8], vec4: &[f32; 4]) {
196        let vec4 = bytemuck::cast_slice(vec4);
197        buffer.copy_from_slice(vec4);
198    }
199    /// Bind a `vec4` to the given offset.
200    #[inline(always)]
201    pub fn bind_vec4(
202        &mut self,
203        offset: MemberOffset,
204        value: impl Into<[f32; 4]>,
205        ctx: C,
206        device: &D,
207    ) {
208        let vec4 = value.into();
209
210        for ty in UniformMemberBlock::TYPES {
211            if H::bind_uniform(ty, &vec4, ctx, device).is_some() {
212                continue;
213            }
214            if let Some(offset) = offset.offset(ty) {
215                let buffer = self.buffer(ty);
216                Self::write_vec4_inner(
217                    &mut buffer[offset..][..4 * std::mem::size_of::<f32>()],
218                    &vec4,
219                );
220            }
221        }
222    }
223}
224
225impl<H, C, U, P, D> UniformStorage<H, C, U, P, D>
226where
227    C: Copy,
228    U: Deref<Target = [u8]> + DerefMut,
229    P: Deref<Target = [u8]> + DerefMut,
230    H: for<'a> BindUniform<C, &'a [f32; 16], D>,
231{
232    #[inline(always)]
233    fn write_mat4_inner(buffer: &mut [u8], mat4: &[f32; 16]) {
234        let mat4 = bytemuck::cast_slice(mat4);
235        buffer.copy_from_slice(mat4);
236    }
237
238    /// Bind a `mat4` to the given offset.
239    #[inline(always)]
240    pub fn bind_mat4(&mut self, offset: MemberOffset, value: &[f32; 16], ctx: C, device: &D) {
241        for ty in UniformMemberBlock::TYPES {
242            if H::bind_uniform(ty, value, ctx, device).is_some() {
243                continue;
244            }
245            if let Some(offset) = offset.offset(ty) {
246                let buffer = self.buffer(ty);
247                Self::write_mat4_inner(
248                    &mut buffer[offset..][..16 * std::mem::size_of::<f32>()],
249                    value,
250                );
251            }
252        }
253    }
254}