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
22#[must_use]
23pub fn check_byte_layout<'lt, T>(bytes: &[u8]) -> bool
24where
25    T: Archive,
26    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
27{
28    rkyv::access::<T::Archived, Error>(bytes).is_ok()
29}
30
31impl<D, T> TypedDataChunk<D, T>
32where
33    D: DataChunk,
34    T: Archive,
35    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
36{
37    /// Builds a typed view after validating that the byte layout is a valid archived `T`.
38    ///
39    /// This method assumes `D` upholds [`crate::DataChunk`] invariants (stable, immutable
40    /// bytes/hash for `&self`) for the lifetime of this value.
41    pub fn from_data_chunk(chunk: D) -> Result<Self> {
42        rkyv::access::<T::Archived, Error>(chunk.data_ref())
43            .map_err(|_| crate::DataChunkError::InvalidArchive)?;
44
45        let chunk = Self {
46            _p: PhantomData,
47            chunk,
48        };
49
50        Ok(chunk)
51    }
52
53    /// Returns a checked typed reference.
54    ///
55    /// Unlike [`Deref`], this method always validates the underlying bytes before
56    /// returning the archived value.
57    pub fn typed_ref(&self) -> Result<&T::Archived> {
58        rkyv::access::<T::Archived, Error>(self.chunk.data_ref())
59            .map_err(|_| crate::DataChunkError::InvalidArchive)
60    }
61}
62
63impl<D, T> Deref for TypedDataChunk<D, T>
64where
65    D: DataChunk,
66    T: Archive,
67    for<'a> <T as Archive>::Archived: CheckBytes<HighValidator<'a, Error>>,
68{
69    type Target = T::Archived;
70
71    fn deref(&self) -> &Self::Target {
72        // SAFETY:
73        // - `from_data_chunk` validates that `chunk.data_ref()` contains a valid `T::Archived`.
74        // - `TypedDataChunk` only exposes shared access to `chunk`, so no mutation happens
75        //   through this type after validation.
76        // - This relies on the `DataChunk` contract that bytes/hash are stable and immutable for `&self`.
77        unsafe { rkyv::access_unchecked::<T::Archived>(self.chunk.data_ref()) }
78    }
79}
80
81impl<D, T> DataChunk for TypedDataChunk<D, T>
82where
83    D: DataChunk,
84    T: Archive,
85    for<'a> <T as Archive>::Archived: CheckBytes<HighValidator<'a, Error>>,
86{
87    fn data_ref(&self) -> &[u8] {
88        self.chunk.data_ref()
89    }
90
91    fn hash_ref(&self) -> &Hash {
92        self.chunk.hash_ref()
93    }
94
95    /// Transforms this [`DataChunk`] into [`Bytes`].
96    fn into_bytes(self) -> Bytes {
97        self.chunk.into_bytes()
98    }
99
100    /// Transforms this chunk into an [`crate::OwnedDataChunk`]
101    fn into_owned(self) -> crate::OwnedDataChunk {
102        let Self { chunk, _p } = self;
103
104        chunk.into_owned()
105    }
106}
107
108pub trait ToDataChunk {
109    fn to_datachunk(&self) -> Result<AlignedDataChunk>;
110}
111
112impl<T: Archive + ToTypedDataChunk<T>> ToDataChunk for T {
113    fn to_datachunk(&self) -> Result<AlignedDataChunk> {
114        Ok(self.to_typed_datachunk()?.chunk)
115    }
116}
117
118pub trait ToTypedDataChunk<T: Archive> {
119    fn to_typed_datachunk(&self) -> Result<TypedDataChunk<AlignedDataChunk, T>>;
120}
121
122impl<T> ToTypedDataChunk<T> for T
123where
124    T::Archived: for<'a> CheckBytes<HighValidator<'a, Error>>,
125    T: rkyv::Archive
126        + for<'a> Serialize<Strategy<Serializer<AlignedVec, ArenaHandle<'a>, Share>, Error>>,
127{
128    fn to_typed_datachunk(&self) -> Result<TypedDataChunk<AlignedDataChunk, T>> {
129        let chunk = AlignedDataChunk::try_from::<T>(self)?;
130
131        TypedDataChunk::from_data_chunk(chunk)
132    }
133}
134
135#[allow(clippy::expect_used)]
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::{DataChunkError, OwnedDataChunk};
140
141    #[test]
142    fn typed_ref_returns_checked_ref() -> Result<()> {
143        let typed = 42_u32.to_typed_datachunk()?;
144
145        assert_eq!(*typed.typed_ref()?, 42_u32);
146        assert_eq!(*typed, *typed.typed_ref()?);
147
148        Ok(())
149    }
150
151    #[test]
152    fn from_data_chunk_rejects_invalid_archive() {
153        let chunk = OwnedDataChunk::from_data([1_u8, 2, 3]).expect("hashing failed");
154
155        let result = TypedDataChunk::<OwnedDataChunk, u32>::from_data_chunk(chunk);
156
157        assert!(matches!(result, Err(DataChunkError::InvalidArchive)));
158    }
159}
ps_datachunk/typed/mod.rs

ps_datachunk/typed/
mod.rs
