Skip to main content

ps_datachunk/typed/
mod.rs

1mod implementations;
2
3use std::{marker::PhantomData, ops::Deref};
4
5use bytes::Bytes;
6use rancor::{Error, Strategy};
7use rkyv::{
8    api::high::HighValidator,
9    bytecheck::CheckBytes,
10    ser::{allocator::ArenaHandle, sharing::Share, Serializer},
11    util::AlignedVec,
12    Archive, Serialize,
13};
14
15use crate::{AlignedDataChunk, DataChunk, Hash, Result};
16
17pub struct TypedDataChunk<D: DataChunk, T: rkyv::Archive> {
18    chunk: D,
19    _p: PhantomData<T::Archived>,
20}
21
22impl<D, T> TypedDataChunk<D, T>
23where
24    D: DataChunk,
25    T: Archive,
26    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
27{
28    /// Builds a typed view after validating that the byte layout is a valid archived `T`.
29    ///
30    /// This method assumes `D` upholds [`crate::DataChunk`] invariants (stable, immutable
31    /// bytes/hash for `&self`) for the lifetime of this value.
32    pub fn from_data_chunk(chunk: D) -> Result<Self> {
33        rkyv::access::<T::Archived, Error>(chunk.data_ref())
34            .map_err(|_| crate::DataChunkError::InvalidArchive)?;
35
36        let chunk = Self {
37            _p: PhantomData,
38            chunk,
39        };
40
41        Ok(chunk)
42    }
43}
44
45impl<D, T> Deref for TypedDataChunk<D, T>
46where
47    D: DataChunk,
48    T: Archive,
49    for<'a> <T as Archive>::Archived: CheckBytes<HighValidator<'a, Error>>,
50{
51    type Target = T::Archived;
52
53    fn deref(&self) -> &Self::Target {
54        // SAFETY:
55        // - `from_data_chunk` validates that `chunk.data_ref()` contains a valid `T::Archived`.
56        // - `TypedDataChunk` only exposes shared access to `chunk`, so no mutation happens
57        //   through this type after validation.
58        // - This relies on the `DataChunk` contract that bytes/hash are stable and immutable for `&self`.
59        unsafe { rkyv::access_unchecked::<T::Archived>(self.chunk.data_ref()) }
60    }
61}
62
63impl<D, T> DataChunk for TypedDataChunk<D, T>
64where
65    D: DataChunk,
66    T: Archive,
67    for<'a> <T as Archive>::Archived: CheckBytes<HighValidator<'a, Error>>,
68{
69    fn data_ref(&self) -> &[u8] {
70        self.chunk.data_ref()
71    }
72
73    fn hash_ref(&self) -> &Hash {
74        self.chunk.hash_ref()
75    }
76
77    /// Transforms this [`DataChunk`] into [`Bytes`].
78    fn into_bytes(self) -> Bytes {
79        self.chunk.into_bytes()
80    }
81
82    /// Transforms this chunk into an [`crate::OwnedDataChunk`]
83    fn into_owned(self) -> crate::OwnedDataChunk {
84        let Self { chunk, _p } = self;
85
86        chunk.into_owned()
87    }
88}
89
90pub trait ToDataChunk {
91    fn to_datachunk(&self) -> Result<AlignedDataChunk>;
92}
93
94impl<T: Archive + ToTypedDataChunk<T>> ToDataChunk for T {
95    fn to_datachunk(&self) -> Result<AlignedDataChunk> {
96        Ok(self.to_typed_datachunk()?.chunk)
97    }
98}
99
100pub trait ToTypedDataChunk<T: Archive> {
101    fn to_typed_datachunk(&self) -> Result<TypedDataChunk<AlignedDataChunk, T>>;
102}
103
104impl<T> ToTypedDataChunk<T> for T
105where
106    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
107    T: rkyv::Archive
108        + for<'a> Serialize<Strategy<Serializer<AlignedVec, ArenaHandle<'a>, Share>, Error>>,
109{
110    fn to_typed_datachunk(&self) -> Result<TypedDataChunk<AlignedDataChunk, T>> {
111        let chunk = AlignedDataChunk::try_from::<T>(self)?;
112
113        TypedDataChunk::from_data_chunk(chunk)
114    }
115}
116
117#[allow(clippy::expect_used)]
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::{DataChunkError, OwnedDataChunk};
122
123    #[test]
124    fn deref_returns_archived_value() -> Result<()> {
125        let typed = 42_u32.to_typed_datachunk()?;
126
127        assert_eq!(*typed, 42_u32);
128
129        Ok(())
130    }
131
132    #[test]
133    fn from_data_chunk_rejects_invalid_archive() {
134        let chunk = OwnedDataChunk::from_data([1_u8, 2, 3]).expect("hashing failed");
135
136        let result = TypedDataChunk::<OwnedDataChunk, u32>::from_data_chunk(chunk);
137
138        assert!(matches!(result, Err(DataChunkError::InvalidArchive)));
139    }
140}