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

226