data_streams/source/
exact_size.rs

1// Copyright 2025 - Strixpyrr
2// SPDX-License-Identifier: Apache-2.0
3
4use core::ops::Deref;
5#[cfg(feature = "utf8")]
6use simdutf8::compat::from_utf8;
7use crate::{BufferAccess, DataSource, Result};
8#[cfg(feature = "unstable_ascii_char")]
9use crate::Error;
10use crate::markers::source::SourceSize;
11use crate::source::max_multiple_of;
12
13trait ExactSizeBuffer: Deref<Target = [u8]> {
14	fn len(&self) -> usize { (**self).len() }
15	fn consume(&mut self, count: usize);
16
17	fn read_bytes_infallible<'a>(&mut self, buf: &'a mut [u8]) -> &'a [u8] {
18		let len = self.len().min(buf.len());
19		let filled = &mut buf[..len];
20		filled.copy_from_slice(&self[..len]);
21		self.consume(len);
22		filled
23	}
24	
25	fn read_aligned_bytes_infallible<'a>(&mut self, buf: &'a mut [u8], alignment: usize) -> &'a [u8] {
26		if alignment == 0 { return &[] }
27		let len = max_multiple_of(self.len().min(buf.len()), alignment);
28		self.read_bytes_infallible(&mut buf[..len])
29	}
30}
31
32// Conflicting implementation with blanket impl, use a macro instead.
33macro_rules! impl_source {
34    ($($(#[$meta:meta])?$ty:ty);+) => {
35		$(
36		$(#[$meta])?
37		impl DataSource for $ty {
38			fn available(&self) -> usize { self.len() }
39			fn request(&mut self, count: usize) -> Result<bool> {
40				Ok(self.len() >= count)
41			}
42		
43			fn skip(&mut self, mut count: usize) -> Result<usize> {
44				count = count.min(self.len());
45				self.consume(count);
46				Ok(count)
47			}
48		
49			fn read_bytes<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a [u8]> {
50				Ok(self.read_bytes_infallible(buf))
51			}
52			
53			fn read_aligned_bytes<'a>(&mut self, buf: &'a mut [u8], alignment: usize) -> Result<&'a [u8]> {
54				Ok(self.read_aligned_bytes_infallible(buf, alignment))
55			}
56		
57			/// Reads bytes into a slice, returning them as a UTF-8 string if valid.
58			///
59			/// # Errors
60			///
61			/// Returns [`Error::Utf8`] if invalid UTF-8 is read. This implementation only
62			/// consumes valid UTF-8. `buf` is left with a valid UTF-8 string whose length
63			/// is given by the error, [`Utf8Error::valid_up_to`]. This slice can be safely
64			/// converted to a string with [`from_str_unchecked`] or [`Utf8Error::split_valid`].
65			///
66			/// [`Utf8Error::valid_up_to`]: simdutf8::compat::Utf8Error::valid_up_to
67			/// [`from_str_unchecked`]: core::str::from_utf8_unchecked
68			#[cfg(feature = "utf8")]
69			fn read_utf8<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a str> {
70				let count = buf.len().min(self.len());
71				let filled = &mut buf[..count];
72				filled.copy_from_slice(&self[..count]);
73				let (result, consumed) = match from_utf8(filled) {
74					Ok(str) => (Ok(str), count),
75					Err(error) => (Err(error.into()), error.valid_up_to())
76				};
77				self.consume(consumed);
78				result
79			}
80
81			/// Reads bytes into a slice, returning them as an ASCII slice if valid.
82			///
83			/// # Errors
84			///
85			/// Returns [`Error::Ascii`] if a non-ASCII byte is found. This
86			/// implementation consumes only valid ASCII. `buf` is left with valid
87			/// ASCII bytes with a length of [`AsciiError::valid_up_to`]. The valid
88			/// slice can be retrieved with [`AsciiError::valid_slice`].
89			#[cfg(feature = "unstable_ascii_char")]
90			fn read_ascii<'a>(&mut self, buf: &'a mut [u8]) -> Result<&'a [core::ascii::Char]> {
91				let len = buf.len().min(self.len());
92				let count = super::count_ascii(&self[..len]);
93				let bytes = self.read_bytes_infallible(&mut buf[..count]);
94				if count == len {
95					// Safety: all bytes have been checked as valid ASCII.
96					Ok(unsafe { bytes.as_ascii_unchecked() })
97				} else {
98					Err(Error::invalid_ascii(self[0], count, count))
99				}
100			}
101		})+
102	};
103}
104
105impl_source! { &[u8]; #[cfg(feature = "alloc")] alloc::vec::Vec<u8> }
106
107impl ExactSizeBuffer for &[u8] {
108	fn consume(&mut self, count: usize) {
109		*self = &self[count..];
110	}
111}
112
113// Blanket impls break everything here.
114impl BufferAccess for &[u8] {
115	fn buffer_capacity(&self) -> usize { self.len() }
116
117	fn buffer(&self) -> &[u8] { self }
118
119	fn fill_buffer(&mut self) -> Result<&[u8]> { Ok(self) }
120
121	fn drain_buffer(&mut self, count: usize) { self.consume(count); }
122}
123
124// Safety: slices produce exactly the number of bytes as their length.
125unsafe impl SourceSize for &[u8] {
126	fn lower_bound(&self) -> u64 { self.len() as u64 }
127	fn upper_bound(&self) -> Option<u64> { Some(self.len() as u64) }
128}
129
130#[cfg(feature = "alloc")]
131impl ExactSizeBuffer for alloc::vec::Vec<u8> {
132	fn consume(&mut self, count: usize) {
133		if self.len() == count {
134			self.clear();
135		} else {
136			self.drain(..count);
137		}
138	}
139}
140
141#[cfg(feature = "alloc")]
142impl BufferAccess for alloc::vec::Vec<u8> {
143	fn buffer_capacity(&self) -> usize { self.capacity() }
144
145	fn buffer(&self) -> &[u8] { self }
146
147	fn fill_buffer(&mut self) -> Result<&[u8]> { Ok(self) }
148
149	fn drain_buffer(&mut self, count: usize) { self.consume(count); }
150}
151
152#[cfg(feature = "alloc")]
153// Safety: vectors produce exactly the number of bytes as their length.
154unsafe impl SourceSize for alloc::vec::Vec<u8> {
155	fn lower_bound(&self) -> u64 { self.len() as u64 }
156	fn upper_bound(&self) -> Option<u64> { Some(self.len() as u64) }
157}
data_streams/source/exact_size.rs

data_streams/source/
exact_size.rs
