ps_datachunk/
typed.rs

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

ps_datachunk/
typed.rs
