ps_datachunk/
typed.rs

1use std::{marker::PhantomData, ops::Deref};
2
3use rancor::{Error, Strategy};
4use rkyv::{
5    api::high::HighValidator,
6    bytecheck::CheckBytes,
7    ser::{allocator::ArenaHandle, sharing::Share, Serializer},
8    util::AlignedVec,
9    Archive, Serialize,
10};
11
12use crate::*;
13
14pub struct TypedDataChunk<'lt, T: rkyv::Archive> {
15    chunk: DataChunk<'lt>,
16    _p: PhantomData<T::Archived>,
17}
18
19pub fn check_byte_layout<'lt, T>(bytes: &[u8]) -> bool
20where
21    T: Archive,
22    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
23{
24    rkyv::access::<T::Archived, Error>(bytes).is_ok()
25}
26
27impl<'lt, T: Archive> TypedDataChunk<'lt, T> {
28    /// # Safety
29    ///
30    /// Called guarantees that `chunk.data_ref()` contains a valid `T::Archive`
31    pub unsafe fn from_chunk_unchecked(chunk: DataChunk<'lt>) -> Self {
32        Self {
33            chunk,
34            _p: PhantomData,
35        }
36    }
37}
38
39impl<'lt, T> TryFrom<DataChunk<'lt>> for TypedDataChunk<'lt, T>
40where
41    T: Archive,
42    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
43{
44    type Error = PsDataChunkError;
45
46    fn try_from(chunk: DataChunk<'lt>) -> Result<TypedDataChunk<'lt, T>> {
47        match check_byte_layout::<T>(chunk.data_ref()) {
48            true => Ok(unsafe { TypedDataChunk::from_chunk_unchecked(chunk) }),
49            false => Err(PsDataChunkError::TypeError),
50        }
51    }
52}
53
54impl<'lt, T: Archive> Deref for TypedDataChunk<'lt, T>
55where
56    <T as Archive>::Archived: CheckBytes<HighValidator<'lt, Error>>,
57{
58    type Target = T::Archived;
59
60    fn deref(&self) -> &Self::Target {
61        unsafe { rkyv::access_unchecked::<T::Archived>(self.chunk.data_ref()) }
62    }
63}
64
65impl<'lt, T: Archive> DataChunkTrait for TypedDataChunk<'lt, T> {
66    fn data_ref(&self) -> &[u8] {
67        self.chunk.data_ref()
68    }
69
70    fn hash_ref(&self) -> &[u8] {
71        self.chunk.hash_ref()
72    }
73
74    fn hash(&self) -> Arc<Hash> {
75        self.chunk.hash()
76    }
77}
78
79pub trait ToDataChunk {
80    fn to_datachunk(&self) -> Result<DataChunk>;
81}
82
83impl<T: Archive + ToTypedDataChunk<T>> ToDataChunk for T {
84    fn to_datachunk(&self) -> Result<DataChunk> {
85        Ok(self.to_typed_datachunk()?.chunk)
86    }
87}
88
89pub trait ToTypedDataChunk<T: Archive> {
90    fn to_typed_datachunk(&self) -> Result<TypedDataChunk<'static, T>>;
91}
92
93impl<T> ToTypedDataChunk<T> for T
94where
95    T: rkyv::Archive,
96    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
97    T: for<'a> Serialize<Strategy<Serializer<AlignedVec, ArenaHandle<'a>, Share>, Error>>,
98{
99    fn to_typed_datachunk(&self) -> Result<TypedDataChunk<'static, T>> {
100        let aligned = AlignedDataChunk::try_from::<T>(self)?;
101        let chunk = DataChunk::Aligned(aligned);
102        let typed = unsafe { TypedDataChunk::from_chunk_unchecked(chunk) };
103
104        Ok(typed)
105    }
106}