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 layout(&self) -> Layout {
34        let mapped = self.map_to_cpu();
35        let len = mapped.len();
36        let ptr = self.map_to_cpu() as *const [u8] as *const u8 as usize;
37
38        unsafe { Layout::from_size_align_unchecked(len, 1 << ptr.trailing_zeros()) }
39    }
40
41    #[inline]
42    fn into_cpu_bytes(self) -> Bytes
43    where
44        Self: Sized,
45    {
46        Bytes::copy_from_slice(self.map_to_cpu())
47    }
48}
49
50impl MemBlock for Bytes {
51    type Ref<'a> = &'a Bytes;
52
53    #[inline]
54    fn borrow(&self) -> Self::Ref<'_> {
55        todo!()
56    }
57
58    #[inline]
59    fn device(&self) -> MemDevice {
60        MemDevice::Cpu
61    }
62
63    #[inline]
64    fn map_to_cpu(&self) -> &[u8] {
65        self
66    }
67
68    #[inline]
69    fn into_cpu_bytes(self) -> Bytes {
70        self
71    }
72}
73
74impl<'a> MemBlock for &'a [u8] {
75    type Ref<'b>
76        = &'b [u8]
77    where
78        Self: 'b;
79
80    #[inline]
81    fn borrow(&self) -> Self::Ref<'_> {
82        self
83    }
84
85    #[inline]
86    fn device(&self) -> MemDevice {
87        MemDevice::Cpu
88    }
89
90    #[inline]
91    fn map_to_cpu(&self) -> &[u8] {
92        self
93    }
94}
95
96impl<'a> MemBlock for &'a Bytes {
97    type Ref<'b>
98        = &'b [u8]
99    where
100        Self: 'b;
101
102    #[inline]
103    fn borrow(&self) -> Self::Ref<'_> {
104        self
105    }
106
107    #[inline]
108    fn device(&self) -> MemDevice {
109        MemDevice::Cpu
110    }
111
112    #[inline]
113    fn map_to_cpu(&self) -> &[u8] {
114        self
115    }
116}
117
118pub trait MemAlloc {
119    type Data: MemBlock;
120    type Error: std::error::Error + Send + Sync + 'static;
121
122    fn device(&self) -> MemDevice;
123    fn alloc(&self, data: &[u8]) -> Result<Self::Data, Self::Error>;
124}
125
126impl<A: MemAlloc> MemAlloc for std::sync::Arc<A> {
127    type Data = A::Data;
128    type Error = A::Error;
129
130    fn device(&self) -> MemDevice {
131        (**self).device()
132    }
133
134    fn alloc(&self, data: &[u8]) -> Result<Self::Data, Self::Error> {
135        (**self).alloc(data)
136    }
137}
138
139#[derive(Debug, Clone)]
140pub struct CpuAllocator;
141impl MemAlloc for CpuAllocator {
142    type Data = Bytes;
143    type Error = MemError;
144
145    fn device(&self) -> MemDevice {
146        MemDevice::Cpu
147    }
148
149    fn alloc(&self, data: &[u8]) -> Result<Self::Data, Self::Error> {
150        Ok(Bytes::copy_from_slice(data))
151    }
152}
flowly_core/memory.rs

flowly_core/
memory.rs
