ha_ndarray/host/
buffer.rs

1use std::borrow::Borrow;
2use std::fmt;
3use std::ops::Deref;
4
5use smallvec::SmallVec;
6
7use crate::buffer::{BufferConverter, BufferInstance, BufferMut};
8use crate::{Error, Number};
9
10use super::VEC_MIN_SIZE;
11
12/// A stack-allocated buffer.
13pub type StackVec<T> = SmallVec<[T; VEC_MIN_SIZE]>;
14
15impl<T: Number> BufferInstance<T> for StackVec<T> {
16    fn read(&self) -> BufferConverter<'_, T> {
17        self.as_slice().into()
18    }
19
20    fn read_value(&self, offset: usize) -> Result<T, Error> {
21        BufferInstance::read_value(&self.as_slice(), offset)
22    }
23
24    fn len(&self) -> usize {
25        StackVec::len(self)
26    }
27}
28
29impl<T: Number> BufferMut<T> for StackVec<T> {
30    fn write<'a>(&mut self, data: BufferConverter<'a, T>) -> Result<(), Error> {
31        self.as_mut_slice().write(data)
32    }
33
34    fn write_value(&mut self, value: T) -> Result<(), Error> {
35        self.as_mut_slice().write_value(value)
36    }
37
38    fn write_value_at(&mut self, offset: usize, value: T) -> Result<(), Error> {
39        self.as_mut_slice().write_value_at(offset, value)
40    }
41}
42
43impl<T: Number> BufferInstance<T> for Vec<T> {
44    fn read(&self) -> BufferConverter<'_, T> {
45        self.as_slice().into()
46    }
47
48    fn read_value(&self, offset: usize) -> Result<T, Error> {
49        BufferInstance::read_value(&self.as_slice(), offset)
50    }
51
52    fn len(&self) -> usize {
53        Vec::len(self)
54    }
55}
56
57impl<T: Number> BufferMut<T> for Vec<T> {
58    fn write<'a>(&mut self, data: BufferConverter<'a, T>) -> Result<(), Error> {
59        self.as_mut_slice().write(data)
60    }
61
62    fn write_value(&mut self, value: T) -> Result<(), Error> {
63        self.as_mut_slice().write_value(value)
64    }
65
66    fn write_value_at(&mut self, offset: usize, value: T) -> Result<(), Error> {
67        self.as_mut_slice().write_value_at(offset, value)
68    }
69}
70
71impl<T: Number> BufferInstance<T> for &[T] {
72    fn read(&self) -> BufferConverter<'_, T> {
73        (*self).into()
74    }
75
76    fn read_value(&self, offset: usize) -> Result<T, Error> {
77        self.get(offset).copied().ok_or_else(|| {
78            Error::bounds(format!(
79                "invalid offset {offset} for a buffer of length {}",
80                self.len()
81            ))
82        })
83    }
84
85    fn len(&self) -> usize {
86        <[T]>::len(self)
87    }
88}
89
90impl<T: Number> BufferInstance<T> for &mut [T] {
91    fn read(&self) -> BufferConverter<'_, T> {
92        (&**self).into()
93    }
94
95    fn read_value(&self, offset: usize) -> Result<T, Error> {
96        BufferInstance::read_value(&&**self, offset)
97    }
98
99    fn len(&self) -> usize {
100        <[T]>::len(self)
101    }
102}
103
104impl<T: Number> BufferMut<T> for &mut [T] {
105    fn write<'b>(&mut self, data: BufferConverter<'b, T>) -> Result<(), Error> {
106        if data.len() == self.len() {
107            let data = data.to_slice()?;
108            self.copy_from_slice(&data);
109            Ok(())
110        } else {
111            Err(Error::bounds(format!(
112                "cannot overwrite a buffer of size {} with one of size {}",
113                self.len(),
114                data.len()
115            )))
116        }
117    }
118
119    fn write_value(&mut self, value: T) -> Result<(), Error> {
120        self.fill(value);
121        Ok(())
122    }
123
124    fn write_value_at(&mut self, offset: usize, value: T) -> Result<(), Error> {
125        if offset < self.len() {
126            self[offset] = value;
127            Ok(())
128        } else {
129            Err(Error::bounds(format!(
130                "invalid offset {offset} for a buffer of length {}",
131                self.len()
132            )))
133        }
134    }
135}
136
137#[derive(Clone, Eq, PartialEq, Debug)]
138pub enum Buffer<T> {
139    Heap(Vec<T>),
140    Stack(StackVec<T>),
141}
142
143impl<T: Clone> Buffer<T> {
144    pub fn into_vec(self) -> Vec<T> {
145        match self {
146            Self::Heap(data) => data,
147            Self::Stack(data) => data.into_vec(),
148        }
149    }
150
151    pub fn to_vec(&self) -> Vec<T> {
152        match self {
153            Self::Heap(data) => data.to_vec(),
154            Self::Stack(data) => data.to_vec(),
155        }
156    }
157}
158
159impl<T> Borrow<[T]> for Buffer<T> {
160    fn borrow(&self) -> &[T] {
161        match self {
162            Self::Heap(buf) => buf.borrow(),
163            Self::Stack(buf) => buf.borrow(),
164        }
165    }
166}
167
168impl<T> AsMut<[T]> for Buffer<T> {
169    fn as_mut(&mut self) -> &mut [T] {
170        match self {
171            Self::Heap(buf) => buf.as_mut_slice(),
172            Self::Stack(buf) => buf.as_mut_slice(),
173        }
174    }
175}
176
177impl<T: Number> BufferInstance<T> for Buffer<T> {
178    fn read(&self) -> BufferConverter<'_, T> {
179        BufferConverter::Host(self.into())
180    }
181
182    fn read_value(&self, offset: usize) -> Result<T, Error> {
183        match self {
184            Self::Heap(buf) => buf.read_value(offset),
185            Self::Stack(buf) => buf.read_value(offset),
186        }
187    }
188
189    fn len(&self) -> usize {
190        match self {
191            Self::Heap(buf) => buf.len(),
192            Self::Stack(buf) => buf.len(),
193        }
194    }
195}
196
197impl<T: Number> BufferMut<T> for Buffer<T> {
198    fn write<'a>(&mut self, data: BufferConverter<'a, T>) -> Result<(), Error> {
199        match self {
200            Self::Heap(buf) => buf.write(data),
201            Self::Stack(buf) => buf.write(data),
202        }
203    }
204
205    fn write_value(&mut self, value: T) -> Result<(), Error> {
206        match self {
207            Self::Heap(buf) => buf.write_value(value),
208            Self::Stack(buf) => buf.write_value(value),
209        }
210    }
211
212    fn write_value_at(&mut self, offset: usize, value: T) -> Result<(), Error> {
213        match self {
214            Self::Heap(buf) => buf.write_value_at(offset, value),
215            Self::Stack(buf) => buf.write_value_at(offset, value),
216        }
217    }
218}
219
220impl<T> From<StackVec<T>> for Buffer<T> {
221    fn from(buf: StackVec<T>) -> Self {
222        Self::Stack(buf)
223    }
224}
225
226impl<T> From<Vec<T>> for Buffer<T> {
227    fn from(buf: Vec<T>) -> Self {
228        Self::Heap(buf)
229    }
230}
231
232#[derive(Clone)]
233/// A buffer in host memory, either borrowed or owned
234pub enum SliceConverter<'a, T> {
235    Heap(Vec<T>),
236    Stack(StackVec<T>),
237    Slice(&'a [T]),
238}
239
240impl<'a, T> SliceConverter<'a, T> {
241    /// Return the number of elements in this buffer.
242    pub fn len(&self) -> usize {
243        match self {
244            Self::Heap(vec) => vec.len(),
245            Self::Stack(vec) => vec.len(),
246            Self::Slice(slice) => slice.len(),
247        }
248    }
249
250    pub fn is_empty(&self) -> bool {
251        self.len() == 0
252    }
253}
254
255impl<'a, T: Copy> SliceConverter<'a, T> {
256    /// Return this buffer as an owned [`Vec`].
257    /// This will allocate a new [`Vec`] if this buffer is a [`StackVec`] or borrowed slice.
258    pub fn into_vec(self) -> Vec<T> {
259        match self {
260            Self::Heap(vec) => vec,
261            Self::Stack(vec) => vec.into_vec(),
262            Self::Slice(slice) => slice.to_vec(),
263        }
264    }
265
266    /// Return this buffer as an owned [`StackVec`].
267    pub fn into_stackvec(self) -> StackVec<T> {
268        match self {
269            Self::Heap(vec) => vec.into(),
270            Self::Stack(vec) => vec,
271            Self::Slice(slice) => StackVec::from_slice(slice),
272        }
273    }
274
275    /// Return this buffer as an owned host [`Buffer`].
276    pub fn into_buffer(self) -> Buffer<T> {
277        match self {
278            Self::Heap(vec) => Buffer::Heap(vec),
279            Self::Stack(vec) => Buffer::Stack(vec),
280            Self::Slice(slice) => {
281                if slice.len() < VEC_MIN_SIZE {
282                    Buffer::Stack(StackVec::from_slice(slice))
283                } else {
284                    Buffer::Heap(slice.to_vec())
285                }
286            }
287        }
288    }
289}
290
291impl<T> From<Buffer<T>> for SliceConverter<'static, T> {
292    fn from(buf: Buffer<T>) -> Self {
293        match buf {
294            Buffer::Heap(buf) => SliceConverter::Heap(buf),
295            Buffer::Stack(buf) => SliceConverter::Stack(buf),
296        }
297    }
298}
299
300impl<'a, T> From<&'a Buffer<T>> for SliceConverter<'a, T> {
301    fn from(buf: &'a Buffer<T>) -> Self {
302        match buf {
303            Buffer::Heap(slice) => slice.as_slice().into(),
304            Buffer::Stack(slice) => slice.as_slice().into(),
305        }
306    }
307}
308
309impl<T> From<StackVec<T>> for SliceConverter<'static, T> {
310    fn from(vec: StackVec<T>) -> Self {
311        Self::Stack(vec)
312    }
313}
314
315impl<T> From<Vec<T>> for SliceConverter<'static, T> {
316    fn from(vec: Vec<T>) -> Self {
317        Self::Heap(vec)
318    }
319}
320
321impl<'a, T> From<&'a [T]> for SliceConverter<'a, T> {
322    fn from(slice: &'a [T]) -> Self {
323        Self::Slice(slice)
324    }
325}
326
327impl<'a, T> Deref for SliceConverter<'a, T> {
328    type Target = [T];
329
330    fn deref(&self) -> &Self::Target {
331        match self {
332            Self::Heap(data) => data.as_slice(),
333            Self::Stack(data) => data.as_slice(),
334            Self::Slice(slice) => slice,
335        }
336    }
337}
338
339impl<'a, T: fmt::Debug> fmt::Debug for SliceConverter<'a, T> {
340    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
341        fmt::Debug::fmt(self.deref(), f)
342    }
343}
ha_ndarray/host/buffer.rs

ha_ndarray/host/
buffer.rs
