projects/linux·blob master · refs

224 lines · 6.9 KiB

1// SPDX-License-Identifier: GPL-2.0
2
3use core::ops::Deref;
4
5use 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/// ```
28pub(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
35impl<'a, I> SBufferIter<I>
36where
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.
135impl<'a, I> SBufferIter<I>
136where
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.
183impl<'a, I> SBufferIter<I>
184where
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
214impl<'a, I> Iterator for SBufferIter<I>
215where
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}