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