sbuffer.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3use core::ops::Deref;
  4
  5use kernel::{
  6    alloc::KVec,
  7    prelude::*, //
  8};
  9
 10/// A buffer abstraction for discontiguous byte slices.
 11///
 12/// This allows you to treat multiple non-contiguous `&mut [u8]` slices
 13/// of the same length as a single stream-like read/write buffer.
 14///
 15/// # Examples
 16///
 17/// ```
 18// let mut buf1 = [0u8; 5];
 19/// let mut buf2 = [0u8; 5];
 20/// let mut sbuffer = SBufferIter::new_writer([&mut buf1[..], &mut buf2[..]]);
 21///
 22/// let data = b"hi world!";
 23/// sbuffer.write_all(data)?;
 24/// drop(sbuffer);
 25///
 26/// assert_eq!(buf1, *b"hi wo");
 27/// assert_eq!(buf2, *b"rld!\0");
 28///
 29/// # Ok::<(), Error>(())
 30/// ```
 31pub(crate) struct SBufferIter<I: Iterator> {
 32    // [`Some`] if we are not at the end of the data yet.
 33    cur_slice: Option<I::Item>,
 34    // All the slices remaining after `cur_slice`.
 35    slices: I,
 36}
 37
 38impl<'a, I> SBufferIter<I>
 39where
 40    I: Iterator,
 41{
 42    /// Creates a reader buffer for a discontiguous set of byte slices.
 43    ///
 44    /// # Examples
 45    ///
 46    /// ```
 47    /// let buf1: [u8; 5] = [0, 1, 2, 3, 4];
 48    /// let buf2: [u8; 5] = [5, 6, 7, 8, 9];
 49    /// let sbuffer = SBufferIter::new_reader([&buf1[..], &buf2[..]]);
 50    /// let sum: u8 = sbuffer.sum();
 51    /// assert_eq!(sum, 45);
 52    /// ```
 53    pub(crate) fn new_reader(slices: impl IntoIterator<IntoIter = I>) -> Self
 54    where
 55        I: Iterator<Item = &'a [u8]>,
 56    {
 57        Self::new(slices)
 58    }
 59
 60    /// Creates a writeable buffer for a discontiguous set of byte slices.
 61    ///
 62    /// # Examples
 63    ///
 64    /// ```
 65    /// let mut buf1 = [0u8; 5];
 66    /// let mut buf2 = [0u8; 5];
 67    /// let mut sbuffer = SBufferIter::new_writer([&mut buf1[..], &mut buf2[..]]);
 68    /// sbuffer.write_all(&[0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9][..])?;
 69    /// drop(sbuffer);
 70    /// assert_eq!(buf1, [0, 1, 2, 3, 4]);
 71    /// assert_eq!(buf2, [5, 6, 7, 8, 9]);
 72    ///
 73    /// ```
 74    pub(crate) fn new_writer(slices: impl IntoIterator<IntoIter = I>) -> Self
 75    where
 76        I: Iterator<Item = &'a mut [u8]>,
 77    {
 78        Self::new(slices)
 79    }
 80
 81    fn new(slices: impl IntoIterator<IntoIter = I>) -> Self
 82    where
 83        I::Item: Deref<Target = [u8]>,
 84    {
 85        let mut slices = slices.into_iter();
 86
 87        Self {
 88            // Skip empty slices.
 89            cur_slice: slices.find(|s| !s.deref().is_empty()),
 90            slices,
 91        }
 92    }
 93
 94    /// Returns a slice of at most `len` bytes, or [`None`] if we are at the end of the data.
 95    ///
 96    /// If a slice shorter than `len` bytes has been returned, the caller can call this method
 97    /// again until it returns [`None`] to try and obtain the remainder of the data.
 98    ///
 99    /// The closure `f` should split the slice received in it's first parameter
100    /// at the position given in the second parameter.
101    fn get_slice_internal(
102        &mut self,
103        len: usize,
104        mut f: impl FnMut(I::Item, usize) -> (I::Item, I::Item),
105    ) -> Option<I::Item>
106    where
107        I::Item: Deref<Target = [u8]>,
108    {
109        match self.cur_slice.take() {
110            None => None,
111            Some(cur_slice) => {
112                if len >= cur_slice.len() {
113                    // Caller requested more data than is in the current slice, return it entirely
114                    // and prepare the following slice for being used. Skip empty slices to avoid
115                    // trouble.
116                    self.cur_slice = self.slices.find(|s| !s.is_empty());
117
118                    Some(cur_slice)
119                } else {
120                    // The current slice can satisfy the request, split it and return a slice of
121                    // the requested size.
122                    let (ret, next) = f(cur_slice, len);
123                    self.cur_slice = Some(next);
124
125                    Some(ret)
126                }
127            }
128        }
129    }
130
131    /// Returns whether this buffer still has data available.
132    pub(crate) fn is_empty(&self) -> bool {
133        self.cur_slice.is_none()
134    }
135}
136
137/// Provides a way to get non-mutable slices of data to read from.
138impl<'a, I> SBufferIter<I>
139where
140    I: Iterator<Item = &'a [u8]>,
141{
142    /// Returns a slice of at most `len` bytes, or [`None`] if we are at the end of the data.
143    ///
144    /// If a slice shorter than `len` bytes has been returned, the caller can call this method
145    /// again until it returns [`None`] to try and obtain the remainder of the data.
146    fn get_slice(&mut self, len: usize) -> Option<&'a [u8]> {
147        self.get_slice_internal(len, |s, pos| s.split_at(pos))
148    }
149
150    /// Ideally we would implement `Read`, but it is not available in `core`.
151    /// So mimic `std::io::Read::read_exact`.
152    #[expect(unused)]
153    pub(crate) fn read_exact(&mut self, mut dst: &mut [u8]) -> Result {
154        while !dst.is_empty() {
155            match self.get_slice(dst.len()) {
156                None => return Err(EINVAL),
157                Some(src) => {
158                    let dst_slice;
159                    (dst_slice, dst) = dst.split_at_mut(src.len());
160                    dst_slice.copy_from_slice(src);
161                }
162            }
163        }
164
165        Ok(())
166    }
167
168    /// Read all the remaining data into a [`KVec`].
169    ///
170    /// `self` will be empty after this operation.
171    pub(crate) fn flush_into_kvec(&mut self, flags: kernel::alloc::Flags) -> Result<KVec<u8>> {
172        let mut buf = KVec::<u8>::new();
173
174        if let Some(slice) = core::mem::take(&mut self.cur_slice) {
175            buf.extend_from_slice(slice, flags)?;
176        }
177        for slice in &mut self.slices {
178            buf.extend_from_slice(slice, flags)?;
179        }
180
181        Ok(buf)
182    }
183}
184
185/// Provides a way to get mutable slices of data to write into.
186impl<'a, I> SBufferIter<I>
187where
188    I: Iterator<Item = &'a mut [u8]>,
189{
190    /// Returns a mutable slice of at most `len` bytes, or [`None`] if we are at the end of the
191    /// data.
192    ///
193    /// If a slice shorter than `len` bytes has been returned, the caller can call this method
194    /// again until it returns `None` to try and obtain the remainder of the data.
195    fn get_slice_mut(&mut self, len: usize) -> Option<&'a mut [u8]> {
196        self.get_slice_internal(len, |s, pos| s.split_at_mut(pos))
197    }
198
199    /// Ideally we would implement [`Write`], but it is not available in `core`.
200    /// So mimic `std::io::Write::write_all`.
201    pub(crate) fn write_all(&mut self, mut src: &[u8]) -> Result {
202        while !src.is_empty() {
203            match self.get_slice_mut(src.len()) {
204                None => return Err(ETOOSMALL),
205                Some(dst) => {
206                    let src_slice;
207                    (src_slice, src) = src.split_at(dst.len());
208                    dst.copy_from_slice(src_slice);
209                }
210            }
211        }
212
213        Ok(())
214    }
215}
216
217impl<'a, I> Iterator for SBufferIter<I>
218where
219    I: Iterator<Item = &'a [u8]>,
220{
221    type Item = u8;
222
223    fn next(&mut self) -> Option<Self::Item> {
224        // Returned slices are guaranteed to not be empty so we can safely index the first entry.
225        self.get_slice(1).map(|s| s[0])
226    }
227}