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 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}