use rkyv::util::AlignedVec;
use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
use std::path::Path;
use crate::machine::{AsyncFileIO, blocking::BlockingIO};
#[cfg(feature = "parallel")]
use rayon::prelude::*;
const SMALL_FILE_THRESHOLD: usize = 1024; const PARALLEL_THRESHOLD: usize = 10_000;
fn deserialize_checked<T>(bytes: &[u8]) -> Result<T, std::io::Error>
where
T: Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>>
+ RkyvDeserialize<T, rkyv::rancor::Strategy<rkyv::de::Pool, rkyv::rancor::Error>>,
{
let archived = rkyv::access::<T::Archived, rkyv::rancor::Error>(bytes)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
let mut deserializer = rkyv::de::Pool::new();
archived
.deserialize(rkyv::rancor::Strategy::wrap(&mut deserializer))
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))
}
pub struct OptimizedRkyv {
io: BlockingIO,
}
impl Default for OptimizedRkyv {
fn default() -> Self {
Self::new()
}
}
impl OptimizedRkyv {
pub fn new() -> Self {
Self {
io: BlockingIO::new(),
}
}
pub fn serialize_to_file<T>(&self, value: &T, path: &Path) -> Result<(), std::io::Error>
where
T: for<'a> RkyvSerialize<
rkyv::rancor::Strategy<
rkyv::ser::Serializer<
AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::ser::sharing::Share,
>,
rkyv::rancor::Error,
>,
>,
{
let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(value)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))?;
if bytes.len() < SMALL_FILE_THRESHOLD {
if let Some(parent) = path.parent() {
if !parent.as_os_str().is_empty() && !parent.exists() {
std::fs::create_dir_all(parent)?;
}
}
std::fs::write(path, &bytes)
} else {
self.io.write_sync(path, &bytes)
}
}
pub fn deserialize_from_file<T>(&self, path: &Path) -> Result<T, std::io::Error>
where
T: Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + RkyvDeserialize<T, rkyv::rancor::Strategy<rkyv::de::Pool, rkyv::rancor::Error>>,
{
let metadata = std::fs::metadata(path)?;
let file_size = metadata.len() as usize;
let bytes = if file_size < SMALL_FILE_THRESHOLD {
std::fs::read(path)?
} else {
self.io.read_sync(path)?
};
deserialize_checked::<T>(&bytes)
}
pub fn serialize_batch_smart<T>(&self, items: &[T]) -> Result<Vec<AlignedVec>, std::io::Error>
where
T: for<'a> RkyvSerialize<
rkyv::rancor::Strategy<
rkyv::ser::Serializer<
AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::ser::sharing::Share,
>,
rkyv::rancor::Error,
>,
> + Sync,
{
if items.is_empty() {
return Ok(Vec::new());
}
if items.len() < PARALLEL_THRESHOLD {
let mut results = Vec::with_capacity(items.len());
for item in items {
results.push(
rkyv::to_bytes::<rkyv::rancor::Error>(item).map_err(|e| {
std::io::Error::new(std::io::ErrorKind::Other, e.to_string())
})?,
);
}
return Ok(results);
}
#[cfg(feature = "parallel")]
{
items
.par_iter()
.map(|item| {
rkyv::to_bytes::<rkyv::rancor::Error>(item)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))
})
.collect()
}
#[cfg(not(feature = "parallel"))]
{
let mut results = Vec::with_capacity(items.len());
for item in items {
results.push(
rkyv::to_bytes::<rkyv::rancor::Error>(item).map_err(|e| {
std::io::Error::new(std::io::ErrorKind::Other, e.to_string())
})?,
);
}
Ok(results)
}
}
pub fn serialize_batch_to_files<T>(
&self,
items: &[(T, &Path)],
) -> Result<Vec<std::io::Result<()>>, std::io::Error>
where
T: for<'a> RkyvSerialize<
rkyv::rancor::Strategy<
rkyv::ser::Serializer<
AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::ser::sharing::Share,
>,
rkyv::rancor::Error,
>,
>,
{
let serialized: Vec<_> = items
.iter()
.map(|(item, _)| {
rkyv::to_bytes::<rkyv::rancor::Error>(item)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))
})
.collect::<Result<Vec<_>, _>>()?;
let files: Vec<_> = items
.iter()
.zip(serialized.iter())
.map(|((_, path), bytes)| (*path, bytes.as_ref()))
.collect();
self.io.write_batch_sync(&files)
}
pub fn deserialize_batch_from_files<T>(
&self,
paths: &[&Path],
) -> Result<Vec<Result<T, std::io::Error>>, std::io::Error>
where
T: Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + RkyvDeserialize<T, rkyv::rancor::Strategy<rkyv::de::Pool, rkyv::rancor::Error>>,
{
let read_results = self.io.read_batch_sync(paths)?;
Ok(read_results
.into_iter()
.map(|result| result.and_then(|bytes| deserialize_checked::<T>(&bytes)))
.collect())
}
pub fn backend_name(&self) -> &'static str {
self.io.backend_name()
}
}
#[cfg(feature = "arena")]
pub struct ArenaRkyv {
arena: crate::machine::arena::DxArena,
capacity: usize,
}
#[cfg(feature = "arena")]
impl ArenaRkyv {
pub fn new() -> Self {
let capacity = 1024 * 1024; Self {
arena: crate::machine::arena::DxArena::new(capacity),
capacity,
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
arena: crate::machine::arena::DxArena::new(capacity),
capacity,
}
}
pub fn serialize_batch<T>(&mut self, items: &[T]) -> Result<Vec<AlignedVec>, std::io::Error>
where
T: for<'a> RkyvSerialize<
rkyv::rancor::Strategy<
rkyv::ser::Serializer<
AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::ser::sharing::Share,
>,
rkyv::rancor::Error,
>,
>,
{
let mut results = Vec::with_capacity(items.len());
for item in items {
results.push(
rkyv::to_bytes::<rkyv::rancor::Error>(item)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))?,
);
}
Ok(results)
}
pub fn reset(&mut self) {
self.arena = crate::machine::arena::DxArena::new(self.capacity);
}
}
#[cfg(feature = "arena")]
impl Default for ArenaRkyv {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "compression")]
pub struct CompressedRkyv {
level: crate::machine::compress::CompressionLevel,
}
#[cfg(feature = "compression")]
impl CompressedRkyv {
pub fn new(level: crate::machine::compress::CompressionLevel) -> Self {
Self { level }
}
pub fn serialize_compressed<T>(&mut self, value: &T) -> Result<Vec<u8>, std::io::Error>
where
T: for<'a> RkyvSerialize<
rkyv::rancor::Strategy<
rkyv::ser::Serializer<
AlignedVec,
rkyv::ser::allocator::ArenaHandle<'a>,
rkyv::ser::sharing::Share,
>,
rkyv::rancor::Error,
>,
>,
{
let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(value)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))?;
let compressed = crate::machine::compress::DxCompressed::compress_level(&bytes, self.level);
Ok(compressed.to_wire())
}
pub fn deserialize_compressed<T>(&mut self, compressed: &[u8]) -> Result<T, std::io::Error>
where
T: Archive,
T::Archived: for<'a> rkyv::bytecheck::CheckBytes<
rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
> + RkyvDeserialize<T, rkyv::rancor::Strategy<rkyv::de::Pool, rkyv::rancor::Error>>,
{
let mut dx_compressed = crate::machine::compress::DxCompressed::from_wire(compressed)
.map_err(|e| {
std::io::Error::new(std::io::ErrorKind::InvalidData, format!("{:?}", e))
})?;
let bytes = dx_compressed.decompress().map_err(|e| {
std::io::Error::new(std::io::ErrorKind::InvalidData, format!("{:?}", e))
})?;
deserialize_checked::<T>(bytes)
}
}
#[cfg(feature = "mmap")]
pub struct MmapRkyv {
_phantom: std::marker::PhantomData<()>,
}
#[cfg(feature = "mmap")]
impl MmapRkyv {
pub fn new() -> Self {
Self {
_phantom: std::marker::PhantomData,
}
}
pub fn open<T>(&self, path: &Path) -> Result<crate::machine::mmap::DxMmap, std::io::Error>
where
T: Archive,
{
crate::machine::mmap::DxMmap::open(path)
}
}
#[cfg(feature = "mmap")]
impl Default for MmapRkyv {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use rkyv::{Archive, Deserialize, Serialize};
use tempfile::TempDir;
#[derive(Archive, Serialize, Deserialize, Debug, PartialEq)]
struct TestData {
id: u64,
name: String,
}
#[test]
fn test_optimized_file_io() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("test.rkyv");
let opt = OptimizedRkyv::new();
let data = TestData {
id: 42,
name: "test".to_string(),
};
opt.serialize_to_file(&data, &path).unwrap();
let loaded: TestData = opt.deserialize_from_file(&path).unwrap();
assert_eq!(data, loaded);
}
#[test]
fn test_backend_name() {
let opt = OptimizedRkyv::new();
let backend = opt.backend_name();
assert_eq!(backend, "blocking");
}
#[cfg(feature = "compression")]
#[test]
fn test_compressed_rkyv() {
use crate::machine::compress::CompressionLevel;
let mut comp = CompressedRkyv::new(CompressionLevel::Fast);
let data = TestData {
id: 42,
name: "test".to_string(),
};
let compressed = comp.serialize_compressed(&data).unwrap();
let decompressed: TestData = comp.deserialize_compressed(&compressed).unwrap();
assert_eq!(data, decompressed);
}
#[cfg(feature = "arena")]
#[test]
fn test_arena_rkyv() {
let mut arena = ArenaRkyv::new();
let items = vec![
TestData {
id: 1,
name: "one".to_string(),
},
TestData {
id: 2,
name: "two".to_string(),
},
TestData {
id: 3,
name: "three".to_string(),
},
];
let serialized = arena.serialize_batch(&items).unwrap();
assert_eq!(serialized.len(), 3);
}
}