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