1use librashader_reflect::reflect::semantics::{MemberOffset, UniformMemberBlock};
2use std::marker::PhantomData;
3use std::ops::{Deref, DerefMut};
4
5pub trait UniformScalar: Copy + bytemuck::Pod {}
7impl UniformScalar for f32 {}
8impl UniformScalar for i32 {}
9impl UniformScalar for u32 {}
10
11pub trait BindUniform<C, T, D> {
13 fn bind_uniform(block: UniformMemberBlock, value: T, ctx: C, device: &D) -> Option<()>;
19}
20
21pub trait UniformStorageAccess {
23 fn ubo_pointer(&self) -> *const u8;
26
27 fn ubo_slice(&self) -> &[u8];
30
31 fn push_pointer(&self) -> *const u8;
34
35 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
62pub 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
71pub 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 pub fn inner_ubo(&self) -> &U {
91 &self.ubo
92 }
93
94 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 #[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 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 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 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 #[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 #[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}