Skip to main content

encase/core/
rw.rs

1use super::ShaderType;
2use alloc::boxed::Box;
3use alloc::vec::Vec;
4use core::mem::MaybeUninit;
5use thiserror::Error;
6
7#[derive(Clone, Copy, Debug, Error)]
8pub enum Error {
9    #[error("could not read/write {expected} bytes from/into {found} byte sized buffer")]
10    BufferTooSmall { expected: u64, found: u64 },
11}
12
13pub type Result<T> = core::result::Result<T, Error>;
14
15pub struct WriteContext {
16    /// length of the contained runtime sized array
17    ///
18    /// used by the derive macro
19    pub rts_array_length: Option<u32>,
20}
21
22pub struct Writer<B: BufferMut> {
23    pub ctx: WriteContext,
24    cursor: Cursor<B>,
25}
26
27impl<B: BufferMut> Writer<B> {
28    #[inline]
29    pub fn new<T: ?Sized + ShaderType>(data: &T, buffer: B, offset: usize) -> Result<Self> {
30        let mut cursor = Cursor::new(buffer, offset);
31        let size = data.size().get();
32        if cursor.try_enlarge(offset + size as usize).is_err() {
33            Err(Error::BufferTooSmall {
34                expected: size,
35                found: cursor.capacity() as u64,
36            })
37        } else {
38            Ok(Self {
39                ctx: WriteContext {
40                    rts_array_length: None,
41                },
42                cursor,
43            })
44        }
45    }
46
47    #[inline]
48    pub fn advance(&mut self, amount: usize) {
49        self.cursor.advance(amount);
50    }
51
52    #[inline]
53    pub fn write<const N: usize>(&mut self, val: &[u8; N]) {
54        self.cursor.write(val);
55    }
56
57    #[inline]
58    pub fn write_slice(&mut self, val: &[u8]) {
59        self.cursor.write_slice(val);
60    }
61}
62
63pub struct ReadContext {
64    /// max elements to read into the contained runtime sized array
65    ///
66    /// used by the derive macro
67    pub rts_array_max_el_to_read: Option<u32>,
68}
69
70pub struct Reader<B: BufferRef> {
71    pub ctx: ReadContext,
72    cursor: Cursor<B>,
73}
74
75impl<B: BufferRef> Reader<B> {
76    #[inline]
77    pub fn new<T: ?Sized + ShaderType>(buffer: B, offset: usize) -> Result<Self> {
78        let cursor = Cursor::new(buffer, offset);
79        if cursor.remaining() < T::min_size().get() as usize {
80            Err(Error::BufferTooSmall {
81                expected: T::min_size().get(),
82                found: cursor.remaining() as u64,
83            })
84        } else {
85            Ok(Self {
86                ctx: ReadContext {
87                    rts_array_max_el_to_read: None,
88                },
89                cursor,
90            })
91        }
92    }
93
94    #[inline]
95    pub fn advance(&mut self, amount: usize) {
96        self.cursor.advance(amount);
97    }
98
99    #[inline]
100    pub fn read<const N: usize>(&mut self) -> &[u8; N] {
101        self.cursor.read()
102    }
103
104    #[inline]
105    pub fn read_slice(&mut self, val: &mut [u8]) {
106        self.cursor.read_slice(val);
107    }
108
109    #[inline]
110    pub fn remaining(&self) -> usize {
111        self.cursor.remaining()
112    }
113}
114
115struct Cursor<B> {
116    buffer: B,
117    pos: usize,
118}
119
120impl<B> Cursor<B> {
121    #[inline]
122    fn new(buffer: B, offset: usize) -> Self {
123        Self {
124            buffer,
125            pos: offset,
126        }
127    }
128    #[inline]
129    fn advance(&mut self, amount: usize) {
130        self.pos += amount;
131    }
132}
133
134impl<B: BufferRef> Cursor<B> {
135    #[inline]
136    fn remaining(&self) -> usize {
137        self.buffer.len().saturating_sub(self.pos)
138    }
139
140    #[inline]
141    fn read<const N: usize>(&mut self) -> &[u8; N] {
142        let res = self.buffer.read(self.pos);
143        self.pos += N;
144        res
145    }
146
147    #[inline]
148    fn read_slice(&mut self, val: &mut [u8]) {
149        self.buffer.read_slice(self.pos, val);
150        self.pos += val.len();
151    }
152}
153
154impl<B: BufferMut> Cursor<B> {
155    #[inline]
156    fn capacity(&self) -> usize {
157        self.buffer.capacity().saturating_sub(self.pos)
158    }
159
160    #[inline]
161    fn write<const N: usize>(&mut self, val: &[u8; N]) {
162        self.buffer.write(self.pos, val);
163        self.pos += N;
164    }
165
166    #[inline]
167    fn write_slice(&mut self, val: &[u8]) {
168        self.buffer.write_slice(self.pos, val);
169        self.pos += val.len();
170    }
171
172    #[inline]
173    fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
174        self.buffer.try_enlarge(wanted)
175    }
176}
177
178#[derive(Clone, Copy, Debug, Error)]
179#[error("could not enlarge buffer")]
180pub struct EnlargeError;
181
182impl From<alloc::collections::TryReserveError> for EnlargeError {
183    fn from(_: alloc::collections::TryReserveError) -> Self {
184        Self
185    }
186}
187
188#[allow(clippy::len_without_is_empty)]
189pub trait BufferRef {
190    fn len(&self) -> usize;
191
192    fn read<const N: usize>(&self, offset: usize) -> &[u8; N];
193
194    fn read_slice(&self, offset: usize, val: &mut [u8]);
195}
196
197pub trait BufferMut {
198    fn capacity(&self) -> usize;
199
200    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]);
201
202    fn write_slice(&mut self, offset: usize, val: &[u8]);
203
204    #[inline]
205    fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
206        if wanted > self.capacity() {
207            Err(EnlargeError)
208        } else {
209            Ok(())
210        }
211    }
212}
213
214impl BufferRef for [u8] {
215    fn len(&self) -> usize {
216        self.len()
217    }
218
219    #[inline]
220    fn read<const N: usize>(&self, offset: usize) -> &[u8; N] {
221        use crate::utils::SliceExt;
222        self.array(offset)
223    }
224
225    #[inline]
226    fn read_slice(&self, offset: usize, val: &mut [u8]) {
227        val.copy_from_slice(&self[offset..offset + val.len()]);
228    }
229}
230
231impl<const LEN: usize> BufferRef for [u8; LEN] {
232    #[inline]
233    fn len(&self) -> usize {
234        <[u8] as BufferRef>::len(self)
235    }
236
237    #[inline]
238    fn read<const N: usize>(&self, offset: usize) -> &[u8; N] {
239        <[u8] as BufferRef>::read(self, offset)
240    }
241
242    #[inline]
243    fn read_slice(&self, offset: usize, val: &mut [u8]) {
244        <[u8] as BufferRef>::read_slice(self, offset, val);
245    }
246}
247
248impl BufferRef for Vec<u8> {
249    #[inline]
250    fn len(&self) -> usize {
251        <[u8] as BufferRef>::len(self)
252    }
253
254    #[inline]
255    fn read<const N: usize>(&self, offset: usize) -> &[u8; N] {
256        <[u8] as BufferRef>::read(self, offset)
257    }
258
259    #[inline]
260    fn read_slice(&self, offset: usize, val: &mut [u8]) {
261        <[u8] as BufferRef>::read_slice(self, offset, val);
262    }
263}
264
265impl BufferMut for [u8] {
266    #[inline]
267    fn capacity(&self) -> usize {
268        self.len()
269    }
270
271    #[inline]
272    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
273        use crate::utils::SliceExt;
274        *self.array_mut(offset) = *val;
275    }
276
277    #[inline]
278    fn write_slice(&mut self, offset: usize, val: &[u8]) {
279        self[offset..offset + val.len()].copy_from_slice(val);
280    }
281}
282
283impl BufferMut for [MaybeUninit<u8>] {
284    #[inline]
285    fn capacity(&self) -> usize {
286        self.len()
287    }
288
289    #[inline]
290    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
291        use crate::utils::SliceExt;
292        // SAFETY: &[u8; N] and &[MaybeUninit<u8>; N] have the same layout
293        let val: &[MaybeUninit<u8>; N] = unsafe { core::mem::transmute(val) };
294        *self.array_mut(offset) = *val;
295    }
296
297    #[inline]
298    fn write_slice(&mut self, offset: usize, val: &[u8]) {
299        // SAFETY: &[u8] and &[MaybeUninit<u8>] have the same layout
300        let val: &[MaybeUninit<u8>] = unsafe { core::mem::transmute(val) };
301        self[offset..offset + val.len()].copy_from_slice(val);
302    }
303}
304
305impl<const LEN: usize> BufferMut for [u8; LEN] {
306    #[inline]
307    fn capacity(&self) -> usize {
308        <[u8] as BufferMut>::capacity(self)
309    }
310
311    #[inline]
312    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
313        <[u8] as BufferMut>::write(self, offset, val);
314    }
315
316    #[inline]
317    fn write_slice(&mut self, offset: usize, val: &[u8]) {
318        <[u8] as BufferMut>::write_slice(self, offset, val);
319    }
320}
321
322impl<const LEN: usize> BufferMut for [MaybeUninit<u8>; LEN] {
323    #[inline]
324    fn capacity(&self) -> usize {
325        <[MaybeUninit<u8>] as BufferMut>::capacity(self)
326    }
327
328    #[inline]
329    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
330        <[MaybeUninit<u8>] as BufferMut>::write(self, offset, val);
331    }
332
333    #[inline]
334    fn write_slice(&mut self, offset: usize, val: &[u8]) {
335        <[MaybeUninit<u8>] as BufferMut>::write_slice(self, offset, val);
336    }
337}
338
339impl BufferMut for Vec<u8> {
340    #[inline]
341    fn capacity(&self) -> usize {
342        self.capacity()
343    }
344
345    #[inline]
346    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
347        <[u8] as BufferMut>::write(self, offset, val);
348    }
349
350    #[inline]
351    fn write_slice(&mut self, offset: usize, val: &[u8]) {
352        <[u8] as BufferMut>::write_slice(self, offset, val);
353    }
354
355    #[inline]
356    fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
357        use crate::utils::ByteVecExt;
358        self.try_extend(wanted).map_err(EnlargeError::from)
359    }
360}
361
362impl BufferMut for Vec<MaybeUninit<u8>> {
363    #[inline]
364    fn capacity(&self) -> usize {
365        self.capacity()
366    }
367
368    #[inline]
369    fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
370        <[MaybeUninit<u8>] as BufferMut>::write(self, offset, val);
371    }
372
373    #[inline]
374    fn write_slice(&mut self, offset: usize, val: &[u8]) {
375        <[MaybeUninit<u8>] as BufferMut>::write_slice(self, offset, val);
376    }
377
378    #[inline]
379    fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
380        use crate::utils::ByteVecExt;
381        self.try_extend(wanted).map_err(EnlargeError::from)
382    }
383}
384
385macro_rules! impl_buffer_ref_for_wrappers {
386    ($($type:ty),*) => {$(
387        impl<T: ?Sized + BufferRef> BufferRef for $type {
388            #[inline]
389            fn len(&self) -> usize {
390                T::len(self)
391            }
392
393            #[inline]
394            fn read<const N: usize>(&self, offset: usize) -> &[u8; N] {
395                T::read(self, offset)
396            }
397
398            #[inline]
399            fn read_slice(&self, offset: usize, val: &mut [u8]) {
400                T::read_slice(self, offset, val)
401            }
402        }
403    )*};
404}
405
406impl_buffer_ref_for_wrappers!(&T, &mut T, Box<T>, alloc::rc::Rc<T>, alloc::sync::Arc<T>);
407
408macro_rules! impl_buffer_mut_for_wrappers {
409    ($($type:ty),*) => {$(
410        impl<T: ?Sized + BufferMut> BufferMut for $type {
411            #[inline]
412            fn capacity(&self) -> usize {
413                T::capacity(self)
414            }
415
416            #[inline]
417            fn write<const N: usize>(&mut self, offset: usize, val: &[u8; N]) {
418                T::write(self, offset, val)
419            }
420
421            #[inline]
422            fn write_slice(&mut self, offset: usize, val: &[u8]) {
423                T::write_slice(self, offset, val)
424            }
425
426            #[inline]
427            fn try_enlarge(&mut self, wanted: usize) -> core::result::Result<(), EnlargeError> {
428                T::try_enlarge(self, wanted)
429            }
430        }
431    )*};
432}
433
434impl_buffer_mut_for_wrappers!(&mut T, Box<T>);
435
436#[cfg(test)]
437mod buffer_ref {
438    use super::BufferRef;
439    use alloc::vec::Vec;
440
441    #[test]
442    fn array() {
443        let arr = [0, 1, 2, 3, 4, 5];
444
445        assert_eq!(BufferRef::len(&arr), 6);
446        assert_eq!(BufferRef::read(&arr, 3), &[3, 4]);
447    }
448
449    #[test]
450    fn vec() {
451        let vec = Vec::from([0, 1, 2, 3, 4, 5]);
452
453        assert_eq!(BufferRef::len(&vec), 6);
454        assert_eq!(BufferRef::read(&vec, 3), &[3, 4]);
455    }
456}
457
458#[cfg(test)]
459mod buffer_mut {
460    use super::BufferMut;
461    use crate::core::EnlargeError;
462    use alloc::vec::Vec;
463
464    #[test]
465    fn array() {
466        let mut arr = [0, 1, 2, 3, 4, 5];
467
468        assert_eq!(BufferMut::capacity(&arr), 6);
469
470        BufferMut::write(&mut arr, 3, &[9, 1]);
471        assert_eq!(arr, [0, 1, 2, 9, 1, 5]);
472
473        assert!(matches!(BufferMut::try_enlarge(&mut arr, 6), Ok(())));
474        assert!(matches!(
475            BufferMut::try_enlarge(&mut arr, 7),
476            Err(EnlargeError)
477        ));
478    }
479
480    #[test]
481    fn vec() {
482        let mut vec = Vec::from([0, 1, 2, 3, 4, 5]);
483
484        assert_eq!(BufferMut::capacity(&vec), vec.capacity());
485
486        BufferMut::write(&mut vec, 3, &[9, 1]);
487        assert_eq!(vec, Vec::from([0, 1, 2, 9, 1, 5]));
488
489        assert!(matches!(BufferMut::try_enlarge(&mut vec, 100), Ok(())));
490        assert!(matches!(
491            BufferMut::try_enlarge(&mut vec, usize::MAX),
492            Err(EnlargeError)
493        ));
494    }
495}
496
497#[cfg(test)]
498mod error {
499    use super::Error;
500    use alloc::format;
501
502    #[test]
503    fn derived_traits() {
504        let err = Error::BufferTooSmall {
505            expected: 4,
506            found: 2,
507        };
508
509        {
510            use core::error::Error;
511            assert!(err.source().is_none());
512        }
513
514        assert_eq!(
515            format!("{}", err.clone()),
516            "could not read/write 4 bytes from/into 2 byte sized buffer"
517        );
518
519        assert_eq!(
520            format!("{:?}", err.clone()),
521            "BufferTooSmall { expected: 4, found: 2 }"
522        );
523    }
524}
525
526#[cfg(test)]
527mod enlarge_error {
528    use super::EnlargeError;
529    use alloc::format;
530    use alloc::vec::Vec;
531
532    #[test]
533    fn derived_traits() {
534        // can't construct a TryReserveError due to TryReserveErrorKind being unstable
535        let try_reserve_error = {
536            let mut vec = Vec::<u8>::new();
537            vec.try_reserve(usize::MAX).err().unwrap()
538        };
539        let err = EnlargeError::from(try_reserve_error);
540
541        use core::error::Error;
542        assert!(err.source().is_none());
543
544        assert_eq!(format!("{}", err.clone()), "could not enlarge buffer");
545
546        assert_eq!(format!("{:?}", err.clone()), "EnlargeError");
547    }
548}