Skip to main content

dma_api/
array.rs

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