flowly_core/
memory.rs

1use std::alloc::Layout;
2
3use bytes::Bytes;
4
5#[derive(Debug, thiserror::Error)]
6pub enum MemError {
7    #[error("Mem error")]
8    CpuAllocationError,
9}
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum MemDevice {
13    Cpu,
14    Gpu(u32),
15}
16
17pub trait MemBlock: Send {
18    type Ref<'a>: MemBlock + Clone
19    where
20        Self: 'a;
21
22    fn device(&self) -> MemDevice;
23    fn borrow(&self) -> Self::Ref<'_>;
24    fn map_to_cpu(&self) -> &[u8];
25
26    // Generic non-optimal default implementations
27    #[inline]
28    fn len(&self) -> usize {
29        self.map_to_cpu().len()
30    }
31
32    #[inline]
33    fn is_empty(&self) -> bool {
34        self.len() == 0
35    }
36
37    #[inline]
38    fn layout(&self) -> Layout {
39        let mapped = self.map_to_cpu();
40        let len = mapped.len();
41        let ptr = self.map_to_cpu() as *const [u8] as *const u8 as usize;
42
43        unsafe { Layout::from_size_align_unchecked(len, 1 << ptr.trailing_zeros()) }
44    }
45
46    #[inline]
47    fn into_cpu_bytes(self) -> Bytes
48    where
49        Self: Sized,
50    {
51        Bytes::copy_from_slice(self.map_to_cpu())
52    }
53}
54
55impl MemBlock for Bytes {
56    type Ref<'a> = &'a Bytes;
57
58    #[inline]
59    fn borrow(&self) -> Self::Ref<'_> {
60        self
61    }
62
63    #[inline]
64    fn device(&self) -> MemDevice {
65        MemDevice::Cpu
66    }
67
68    #[inline]
69    fn map_to_cpu(&self) -> &[u8] {
70        self
71    }
72
73    #[inline]
74    fn into_cpu_bytes(self) -> Bytes {
75        self
76    }
77}
78
79impl MemBlock for Vec<u8> {
80    type Ref<'a> = &'a [u8];
81
82    #[inline]
83    fn borrow(&self) -> Self::Ref<'_> {
84        self
85    }
86
87    #[inline]
88    fn device(&self) -> MemDevice {
89        MemDevice::Cpu
90    }
91
92    #[inline]
93    fn map_to_cpu(&self) -> &[u8] {
94        self
95    }
96
97    #[inline]
98    fn into_cpu_bytes(self) -> Bytes {
99        self.into()
100    }
101}
102
103impl MemBlock for &[u8] {
104    type Ref<'b>
105        = &'b [u8]
106    where
107        Self: 'b;
108
109    #[inline]
110    fn borrow(&self) -> Self::Ref<'_> {
111        self
112    }
113
114    #[inline]
115    fn device(&self) -> MemDevice {
116        MemDevice::Cpu
117    }
118
119    #[inline]
120    fn map_to_cpu(&self) -> &[u8] {
121        self
122    }
123}
124
125impl MemBlock for &Bytes {
126    type Ref<'b>
127        = &'b [u8]
128    where
129        Self: 'b;
130
131    #[inline]
132    fn borrow(&self) -> Self::Ref<'_> {
133        self
134    }
135
136    #[inline]
137    fn device(&self) -> MemDevice {
138        MemDevice::Cpu
139    }
140
141    #[inline]
142    fn map_to_cpu(&self) -> &[u8] {
143        self
144    }
145}
146
147pub trait MemAlloc {
148    type Data: MemBlock;
149    type Error: std::error::Error + Send + Sync + 'static;
150
151    fn device(&self) -> MemDevice;
152    fn alloc(&self, data: &[u8]) -> Result<Self::Data, Self::Error>;
153}
154
155impl<A: MemAlloc> MemAlloc for std::sync::Arc<A> {
156    type Data = A::Data;
157    type Error = A::Error;
158
159    fn device(&self) -> MemDevice {
160        (**self).device()
161    }
162
163    fn alloc(&self, data: &[u8]) -> Result<Self::Data, Self::Error> {
164        (**self).alloc(data)
165    }
166}
167
168#[derive(Debug, Clone)]
169pub struct CpuAllocator;
170impl MemAlloc for CpuAllocator {
171    type Data = Bytes;
172    type Error = MemError;
173
174    fn device(&self) -> MemDevice {
175        MemDevice::Cpu
176    }
177
178    fn alloc(&self, data: &[u8]) -> Result<Self::Data, Self::Error> {
179        Ok(Bytes::copy_from_slice(data))
180    }
181}
flowly_core/memory.rs

flowly_core/
memory.rs
