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