Skip to main content

dma_api/
array.rs

1use core::{alloc::Layout, marker::PhantomData, ptr::NonNull};
2
3use crate::{DeviceDma, DmaAddr, DmaDirection, DmaError, DmaPod, common::DmaAllocation};
4
5pub struct CoherentArray<T: DmaPod> {
6    data: DmaAllocation,
7    _phantom: PhantomData<T>,
8}
9
10unsafe impl<T: DmaPod + Send> Send for CoherentArray<T> {}
11unsafe impl<T: DmaPod + Sync> Sync for CoherentArray<T> {}
12
13impl<T: DmaPod> CoherentArray<T> {
14    pub(crate) fn new_zero_with_align(
15        os: &DeviceDma,
16        len: usize,
17        align: usize,
18    ) -> Result<Self, DmaError> {
19        let layout = array_layout::<T>(len, align)?;
20        Ok(Self {
21            data: DmaAllocation::new_zero_coherent(os, layout)?,
22            _phantom: PhantomData,
23        })
24    }
25
26    pub(crate) fn new_zero(os: &DeviceDma, len: usize) -> Result<Self, DmaError> {
27        Self::new_zero_with_align(os, len, core::mem::align_of::<T>())
28    }
29
30    pub fn dma_addr(&self) -> DmaAddr {
31        self.data.handle.dma_addr()
32    }
33
34    pub fn len(&self) -> usize {
35        len_from_bytes::<T>(self.data.handle.size())
36    }
37
38    pub fn is_empty(&self) -> bool {
39        self.len() == 0
40    }
41
42    pub fn bytes_len(&self) -> usize {
43        self.data.handle.size()
44    }
45
46    pub fn read_cpu(&self, index: usize) -> Option<T> {
47        read_at(self.as_ptr(), self.len(), index)
48    }
49
50    pub fn set_cpu(&mut self, index: usize, value: T) {
51        write_at(self.as_ptr(), self.len(), index, value);
52    }
53
54    pub fn copy_from_slice_cpu(&mut self, src: &[T]) {
55        copy_from_slice(self.as_ptr(), self.len(), src);
56    }
57
58    pub fn iter_cpu(&self) -> ArrayCpuIter<'_, T, Self> {
59        ArrayCpuIter {
60            array: self,
61            index: 0,
62            _phantom: PhantomData,
63        }
64    }
65
66    pub fn write_with_cpu<R>(&mut self, len: usize, f: impl FnOnce(&mut [T]) -> R) -> R {
67        assert!(len <= self.len(), "range out of bounds");
68        let data = unsafe { self.as_mut_slice_cpu() };
69        f(&mut data[..len])
70    }
71
72    pub fn read_with_cpu<R>(&self, len: usize, f: impl FnOnce(&[T]) -> R) -> R {
73        assert!(len <= self.len(), "range out of bounds");
74        let data = unsafe { core::slice::from_raw_parts(self.as_ptr().as_ptr(), len) };
75        f(data)
76    }
77
78    pub fn as_ptr(&self) -> NonNull<T> {
79        self.data.handle.as_ptr().cast::<T>()
80    }
81
82    pub fn as_slice_cpu(&self) -> &[T] {
83        unsafe { core::slice::from_raw_parts(self.as_ptr().as_ptr(), self.len()) }
84    }
85
86    /// # Safety
87    ///
88    /// The caller must ensure the device is not concurrently accessing this
89    /// memory in a way that races with CPU writes.
90    pub unsafe fn as_mut_slice_cpu(&mut self) -> &mut [T] {
91        unsafe { core::slice::from_raw_parts_mut(self.as_ptr().as_ptr(), self.len()) }
92    }
93}
94
95pub struct ContiguousArray<T: DmaPod> {
96    data: DmaAllocation,
97    _phantom: PhantomData<T>,
98}
99
100unsafe impl<T: DmaPod + Send> Send for ContiguousArray<T> {}
101unsafe impl<T: DmaPod + Sync> Sync for ContiguousArray<T> {}
102
103impl<T: DmaPod> ContiguousArray<T> {
104    pub(crate) fn new_zero_with_align(
105        os: &DeviceDma,
106        len: usize,
107        align: usize,
108        direction: DmaDirection,
109    ) -> Result<Self, DmaError> {
110        let layout = array_layout::<T>(len, align)?;
111        Ok(Self {
112            data: DmaAllocation::new_zero_contiguous(os, layout, direction)?,
113            _phantom: PhantomData,
114        })
115    }
116
117    pub(crate) fn new_zero(
118        os: &DeviceDma,
119        len: usize,
120        direction: DmaDirection,
121    ) -> Result<Self, DmaError> {
122        Self::new_zero_with_align(os, len, core::mem::align_of::<T>(), direction)
123    }
124
125    pub fn dma_addr(&self) -> DmaAddr {
126        self.data.handle.dma_addr()
127    }
128
129    pub fn len(&self) -> usize {
130        len_from_bytes::<T>(self.data.handle.size())
131    }
132
133    pub fn is_empty(&self) -> bool {
134        self.len() == 0
135    }
136
137    pub fn bytes_len(&self) -> usize {
138        self.data.handle.size()
139    }
140
141    pub fn read_cpu(&self, index: usize) -> Option<T> {
142        read_at(self.as_ptr(), self.len(), index)
143    }
144
145    pub fn set_cpu(&mut self, index: usize, value: T) {
146        write_at(self.as_ptr(), self.len(), index, value);
147    }
148
149    pub fn copy_from_slice_cpu(&mut self, src: &[T]) {
150        copy_from_slice(self.as_ptr(), self.len(), src);
151    }
152
153    pub fn iter_cpu(&self) -> ArrayCpuIter<'_, T, Self> {
154        ArrayCpuIter {
155            array: self,
156            index: 0,
157            _phantom: PhantomData,
158        }
159    }
160
161    pub fn sync_for_device(&self, offset: usize, size: usize) {
162        self.check_range(offset, size);
163        self.data.sync_for_device(offset, size);
164    }
165
166    pub fn sync_for_cpu(&self, offset: usize, size: usize) {
167        self.check_range(offset, size);
168        self.data.sync_for_cpu(offset, size);
169    }
170
171    pub fn sync_for_device_all(&self) {
172        self.data.sync_for_device(0, self.bytes_len());
173    }
174
175    pub fn sync_for_cpu_all(&self) {
176        self.data.sync_for_cpu(0, self.bytes_len());
177    }
178
179    pub fn prepare_for_device(&self, offset: usize, size: usize) {
180        self.sync_for_device(offset, size);
181    }
182
183    pub fn prepare_for_device_all(&self) {
184        self.sync_for_device_all();
185    }
186
187    pub fn complete_for_cpu(&self, offset: usize, size: usize) {
188        self.sync_for_cpu(offset, size);
189    }
190
191    pub fn complete_for_cpu_all(&self) {
192        self.sync_for_cpu_all();
193    }
194
195    pub fn write_for_device<R>(&mut self, len: usize, f: impl FnOnce(&mut [T]) -> R) -> R {
196        let ret = self.write_with_cpu(len, f);
197        self.prepare_for_device(0, len * core::mem::size_of::<T>());
198        ret
199    }
200
201    pub fn read_from_device<R>(&self, len: usize, f: impl FnOnce(&[T]) -> R) -> R {
202        let size = len * core::mem::size_of::<T>();
203        self.complete_for_cpu(0, size);
204        self.read_with_cpu(len, f)
205    }
206
207    pub fn copy_to_device_from_slice(&mut self, src: &[T]) {
208        self.copy_from_slice_cpu(src);
209        self.prepare_for_device(0, core::mem::size_of_val(src));
210    }
211
212    pub fn copy_from_device_to_slice(&self, dst: &mut [T]) {
213        self.read_from_device(dst.len(), |src| dst.copy_from_slice(src));
214    }
215
216    pub fn write_with_cpu<R>(&mut self, len: usize, f: impl FnOnce(&mut [T]) -> R) -> R {
217        assert!(len <= self.len(), "range out of bounds");
218        {
219            let data = unsafe { self.as_mut_slice_cpu() };
220            f(&mut data[..len])
221        }
222    }
223
224    pub fn read_with_cpu<R>(&self, len: usize, f: impl FnOnce(&[T]) -> R) -> R {
225        assert!(len <= self.len(), "range out of bounds");
226        let data = unsafe { core::slice::from_raw_parts(self.as_ptr().as_ptr(), len) };
227        f(data)
228    }
229
230    pub fn as_ptr(&self) -> NonNull<T> {
231        self.data.handle.as_ptr().cast::<T>()
232    }
233
234    pub fn as_slice_cpu(&self) -> &[T] {
235        unsafe { core::slice::from_raw_parts(self.as_ptr().as_ptr(), self.len()) }
236    }
237
238    /// # Safety
239    ///
240    /// The caller must ensure the device is not concurrently accessing this
241    /// memory in a way that races with CPU writes.
242    pub unsafe fn as_mut_slice_cpu(&mut self) -> &mut [T] {
243        unsafe { core::slice::from_raw_parts_mut(self.as_ptr().as_ptr(), self.len()) }
244    }
245
246    fn check_range(&self, offset: usize, size: usize) {
247        assert!(
248            offset <= self.bytes_len() && size <= self.bytes_len().saturating_sub(offset),
249            "range out of bounds, offset: {}, size: {}, bytes_len: {}",
250            offset,
251            size,
252            self.bytes_len()
253        );
254    }
255}
256
257pub trait DmaArrayCpuRead<T: DmaPod> {
258    fn len(&self) -> usize;
259    fn is_empty(&self) -> bool;
260    fn read_cpu(&self, index: usize) -> Option<T>;
261}
262
263impl<T: DmaPod> DmaArrayCpuRead<T> for CoherentArray<T> {
264    fn len(&self) -> usize {
265        CoherentArray::len(self)
266    }
267
268    fn is_empty(&self) -> bool {
269        CoherentArray::is_empty(self)
270    }
271
272    fn read_cpu(&self, index: usize) -> Option<T> {
273        CoherentArray::read_cpu(self, index)
274    }
275}
276
277impl<T: DmaPod> DmaArrayCpuRead<T> for ContiguousArray<T> {
278    fn len(&self) -> usize {
279        ContiguousArray::len(self)
280    }
281
282    fn is_empty(&self) -> bool {
283        ContiguousArray::is_empty(self)
284    }
285
286    fn read_cpu(&self, index: usize) -> Option<T> {
287        ContiguousArray::read_cpu(self, index)
288    }
289}
290
291pub struct ArrayCpuIter<'a, T: DmaPod, A: DmaArrayCpuRead<T>> {
292    array: &'a A,
293    index: usize,
294    _phantom: PhantomData<T>,
295}
296
297impl<'a, T: DmaPod, A: DmaArrayCpuRead<T>> Iterator for ArrayCpuIter<'a, T, A> {
298    type Item = T;
299
300    fn next(&mut self) -> Option<Self::Item> {
301        if self.index >= self.array.len() {
302            return None;
303        }
304        let value = self.array.read_cpu(self.index);
305        self.index += 1;
306        value
307    }
308}
309
310fn array_layout<T>(len: usize, align: usize) -> Result<Layout, DmaError> {
311    let size = len
312        .checked_mul(core::mem::size_of::<T>())
313        .ok_or(DmaError::LayoutError(
314            Layout::from_size_align(usize::MAX, 1).unwrap_err(),
315        ))?;
316    Ok(Layout::from_size_align(
317        size,
318        align.max(core::mem::align_of::<T>()),
319    )?)
320}
321
322fn len_from_bytes<T>(bytes: usize) -> usize {
323    if core::mem::size_of::<T>() == 0 {
324        0
325    } else {
326        bytes / core::mem::size_of::<T>()
327    }
328}
329
330fn read_at<T: DmaPod>(ptr: NonNull<T>, len: usize, index: usize) -> Option<T> {
331    if index >= len {
332        return None;
333    }
334    Some(unsafe { ptr.add(index).read() })
335}
336
337fn write_at<T: DmaPod>(ptr: NonNull<T>, len: usize, index: usize, value: T) {
338    assert!(
339        index < len,
340        "index out of range, index: {}, len: {}",
341        index,
342        len
343    );
344    unsafe { ptr.add(index).write(value) };
345}
346
347fn copy_from_slice<T: DmaPod>(ptr: NonNull<T>, len: usize, src: &[T]) {
348    assert!(
349        src.len() <= len,
350        "source slice is larger than DMA array, src len: {}, array len: {}",
351        src.len(),
352        len
353    );
354    unsafe {
355        ptr.as_ptr()
356            .copy_from_nonoverlapping(src.as_ptr(), src.len());
357    }
358}