Fortune for genmarkov's current commit: Blessing − 吉cli
parent
c4fc2fde1d
commit
066811444a
@ -0,0 +1,379 @@
|
|||||||
|
//! Handles the chain load/save format
|
||||||
|
use super::*;
|
||||||
|
use std::{
|
||||||
|
io::{
|
||||||
|
self,
|
||||||
|
Read, Write, BufRead,
|
||||||
|
},
|
||||||
|
fmt,
|
||||||
|
};
|
||||||
|
use bytes::{
|
||||||
|
Buf, BufMut, Bytes,
|
||||||
|
};
|
||||||
|
use zstd::{
|
||||||
|
Encoder, Decoder,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// The chain that can be saved / loaded.
|
||||||
|
pub type Chain<T = String> = crate::Chain<T>;
|
||||||
|
|
||||||
|
/// The version of the encoded format stream
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Copy)]
|
||||||
|
#[repr(packed)]
|
||||||
|
pub struct Version(pub u8,pub u8,pub u8,pub u8);
|
||||||
|
|
||||||
|
impl fmt::Display for Version
|
||||||
|
{
|
||||||
|
#[inline]
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
|
||||||
|
{
|
||||||
|
write!(f, "{}.{}.{}", self.0,self.1, self.2)?;
|
||||||
|
if self.3 != 0 {
|
||||||
|
write!(f, "r{}", self.3)
|
||||||
|
} else {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
impl Version {
|
||||||
|
/// Current save version
|
||||||
|
pub const CURRENT: Self = Version(0,0,0,0);
|
||||||
|
|
||||||
|
/// Current value as a native integer
|
||||||
|
const CURRENT_VALUE: u32 = Self::CURRENT.as_native();
|
||||||
|
|
||||||
|
pub const fn as_native(&self) -> u32 {
|
||||||
|
u32::from_be_bytes([self.0, self.1, self.2, self.3])
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
pub const fn from_native(value: u32) -> Self {
|
||||||
|
let [a,b,c,d] = u32::to_be_bytes(value);
|
||||||
|
Self(a,b,c,d)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for Version
|
||||||
|
{
|
||||||
|
#[inline]
|
||||||
|
fn default() -> Self
|
||||||
|
{
|
||||||
|
Self::CURRENT
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub unsafe trait AutoBinaryFormat: Sized {
|
||||||
|
#[inline]
|
||||||
|
fn as_raw_for_encode(&self) -> *const [u8] {
|
||||||
|
let ptr = self as *const Self;
|
||||||
|
std::ptr::slice_from_raw_parts(ptr as *const u8, std::mem::size_of::<Self>())
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn as_raw_for_decode(&mut self) -> *mut [u8] {
|
||||||
|
let ptr = self as *mut Self;
|
||||||
|
std::ptr::slice_from_raw_parts_mut(ptr as *mut u8, std::mem::size_of::<Self>())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn raw_format_read_size(&mut self) -> usize {
|
||||||
|
std::mem::size_of::<Self>()
|
||||||
|
}
|
||||||
|
fn raw_format_write_size(&self) -> usize {
|
||||||
|
std::mem::size_of::<Self>()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
unsafe impl<T, const N: usize> AutoBinaryFormat for [T; N] {}
|
||||||
|
|
||||||
|
pub trait BinaryFormat {
|
||||||
|
fn read_from<S: ?Sized>(&mut self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Read;
|
||||||
|
|
||||||
|
fn write_to<S: ?Sized>(&self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Write;
|
||||||
|
|
||||||
|
fn binary_format_read_size(&mut self) -> Option<usize>;
|
||||||
|
fn binary_format_write_size(&self) -> usize;
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<T> BinaryFormat for T
|
||||||
|
where T: AutoBinaryFormat
|
||||||
|
{
|
||||||
|
#[inline(always)]
|
||||||
|
fn read_from<S: ?Sized>(&mut self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Read {
|
||||||
|
let ptr = self.as_raw_for_decode();
|
||||||
|
// SAFETY: The read data is guaranteed to be valid here.
|
||||||
|
Ok(unsafe {
|
||||||
|
stream.read_exact(&mut *ptr)?;
|
||||||
|
(*ptr).len()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
#[inline(always)]
|
||||||
|
fn write_to<S: ?Sized>(&self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Write {
|
||||||
|
let ptr = self.as_raw_for_encode();
|
||||||
|
// SAFETY: The written data is guaranteed to be valid here.
|
||||||
|
Ok(unsafe {
|
||||||
|
stream.write_all(&*ptr)?;
|
||||||
|
(*ptr).len()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_read_size(&mut self) -> Option<usize> {
|
||||||
|
Some(self.raw_format_read_size())
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_write_size(&self) -> usize {
|
||||||
|
self.raw_format_write_size()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<T> BinaryFormat for [T]
|
||||||
|
where T: BinaryFormat
|
||||||
|
{
|
||||||
|
#[inline]
|
||||||
|
fn read_from<S: ?Sized>(&mut self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Read {
|
||||||
|
let mut sz = 0;
|
||||||
|
for i in self.iter_mut() {
|
||||||
|
sz += i.read_from(stream)?;
|
||||||
|
}
|
||||||
|
Ok(sz)
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn write_to<S: ?Sized>(&self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Write {
|
||||||
|
let mut sz =0;
|
||||||
|
for i in self.iter() {
|
||||||
|
sz += i.write_to(stream)?;
|
||||||
|
}
|
||||||
|
Ok(sz)
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_read_size(&mut self) -> Option<usize> {
|
||||||
|
self.iter_mut().map(|x| x.binary_format_read_size()).try_fold(0, |x, y| {
|
||||||
|
match (x, y) {
|
||||||
|
(x, Some(y)) => Some(x + y),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_write_size(&self) -> usize {
|
||||||
|
self.iter().map(|x| x.binary_format_write_size()).sum()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BinaryFormat for Version {
|
||||||
|
#[inline]
|
||||||
|
fn read_from<S: ?Sized>(&mut self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Read {
|
||||||
|
let mut vi = [0u8; 4];
|
||||||
|
stream.read_exact(&mut vi[..])?;
|
||||||
|
Ok(4)
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn write_to<S: ?Sized>(&self, stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Write {
|
||||||
|
let vi = [self.0,self.1,self.2,self.3];
|
||||||
|
stream.write_all(&vi[..])?;
|
||||||
|
Ok(4)
|
||||||
|
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_read_size(&mut self) -> Option<usize> {
|
||||||
|
Some(std::mem::size_of::<u8>() * 4)
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_write_size(&self) -> usize {
|
||||||
|
std::mem::size_of::<u8>() * 4
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy, Default)]
|
||||||
|
#[repr(u32)]
|
||||||
|
pub enum Compressed {
|
||||||
|
#[default]
|
||||||
|
No = 0,
|
||||||
|
Zstd = 1,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Compressed {
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
const fn to_int(&self) -> u32 {
|
||||||
|
*self as u32
|
||||||
|
}
|
||||||
|
#[inline]
|
||||||
|
fn try_from_int(val: u32) -> Option<Self> {
|
||||||
|
match val {
|
||||||
|
// SAFETY: These variants are known
|
||||||
|
0..=1 => Some(unsafe {
|
||||||
|
std::mem::transmute(val)
|
||||||
|
}),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
|
||||||
|
pub struct FormatMetadata {
|
||||||
|
pub version: Version,
|
||||||
|
|
||||||
|
pub compressed: Compressed,
|
||||||
|
pub chain_size: usize, // NOTE: Unused
|
||||||
|
pub checksum: u64, // NOTE: Unused
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FormatMetadata {
|
||||||
|
const MAGIC_NUM: &[u8; 8] = b"MARKOV\x00\xcf";
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BinaryFormat for FormatMetadata
|
||||||
|
{
|
||||||
|
fn write_to<S: ?Sized>(&self, mut stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Write {
|
||||||
|
let sz = self.version.write_to(&mut stream)?;
|
||||||
|
|
||||||
|
let mut obuf = [0u8; std::mem::size_of::<u32>() +std::mem::size_of::<u64>() + std::mem::size_of::<u64>()];
|
||||||
|
{
|
||||||
|
let mut obuf = &mut obuf[..];
|
||||||
|
use std::convert::TryInto;
|
||||||
|
|
||||||
|
obuf.put_u32(self.compressed.to_int());
|
||||||
|
obuf.put_u64(self.chain_size.try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "Chain size attribute out-of-bounds for format size"))?);
|
||||||
|
obuf.put_u64(self.checksum);
|
||||||
|
}
|
||||||
|
stream.write_all(&obuf[..])?;
|
||||||
|
|
||||||
|
Ok(sz + obuf.len())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_from<S: ?Sized>(&mut self, mut stream: &mut S) -> io::Result<usize>
|
||||||
|
where S: io::Read {
|
||||||
|
let sz = self.version.read_from(&mut stream)?;
|
||||||
|
|
||||||
|
if self.version > Version::CURRENT {
|
||||||
|
return Err(io::Error::new(io::ErrorKind::Unsupported, format!("Unknown format version {}", self.version)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut ibuf = [0u8; std::mem::size_of::<u32>() +std::mem::size_of::<u64>() + std::mem::size_of::<u64>()];
|
||||||
|
stream.read_exact(&mut ibuf[..])?;
|
||||||
|
{
|
||||||
|
let mut ibuf = &ibuf[..];
|
||||||
|
use std::convert::TryInto;
|
||||||
|
|
||||||
|
self.compressed = Compressed::try_from_int(ibuf.get_u32()).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Invalid compression attribute"))?;
|
||||||
|
self.chain_size = ibuf.get_u64().try_into().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Chain size attribute out-of-bounds for native size"))?;
|
||||||
|
self.checksum = ibuf.get_u64();
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(sz + ibuf.len())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn binary_format_read_size(&mut self) -> Option<usize> {
|
||||||
|
let szm = self.version.binary_format_read_size()?;
|
||||||
|
Some(szm + std::mem::size_of::<u32>()
|
||||||
|
+ std::mem::size_of::<u64>()
|
||||||
|
+ std::mem::size_of::<u64>())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline(always)]
|
||||||
|
fn binary_format_write_size(&self) -> usize {
|
||||||
|
self.version.binary_format_write_size()
|
||||||
|
+ std::mem::size_of::<u32>()
|
||||||
|
+ std::mem::size_of::<u64>()
|
||||||
|
+ std::mem::size_of::<u64>()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
/// Load a chain from a stream
|
||||||
|
#[inline]
|
||||||
|
pub fn load_chain_from_sync<S>(stream: &mut S) -> io::Result<Chain<String>>
|
||||||
|
where S: io::Read + ?Sized
|
||||||
|
{
|
||||||
|
let mut stream = io::BufReader::new(stream);
|
||||||
|
{
|
||||||
|
let mut magic = FormatMetadata::MAGIC_NUM.clone();
|
||||||
|
stream.read_exact(&mut magic[..])?;
|
||||||
|
|
||||||
|
if &magic != FormatMetadata::MAGIC_NUM {
|
||||||
|
return Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid file header tag magic number"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let metadata = {
|
||||||
|
let mut metadata = FormatMetadata::default();
|
||||||
|
metadata.read_from(&mut stream)?;
|
||||||
|
metadata
|
||||||
|
};
|
||||||
|
|
||||||
|
match metadata.version {
|
||||||
|
Version::CURRENT => {
|
||||||
|
let read = |read: &mut (dyn io::Read)| serde_cbor::from_reader(read).expect("Failed to read chain from input stream"); // TODO: Error type
|
||||||
|
|
||||||
|
match metadata.compressed {
|
||||||
|
Compressed::No =>
|
||||||
|
Ok(read(&mut stream)),
|
||||||
|
Compressed::Zstd => {
|
||||||
|
let mut stream = zstd::Decoder::with_buffer(stream)?;
|
||||||
|
|
||||||
|
//#[cfg(feature="threads")]
|
||||||
|
//stream.multithread(num_cpus::get() as i32);
|
||||||
|
|
||||||
|
//NOTE: Not required here: //stream.finish()?;
|
||||||
|
Ok(read(&mut stream))
|
||||||
|
},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
unsup => {
|
||||||
|
return Err(io::Error::new(io::ErrorKind::Unsupported, format!("Unsupported payload version {}", unsup)));
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Save a chain to a stream with optional compression.
|
||||||
|
#[inline]
|
||||||
|
pub fn save_chain_to_sync<S>(stream: &mut S, chain: &Chain<String>, compress: bool) -> io::Result<()>
|
||||||
|
where S: io::Write + ?Sized
|
||||||
|
{
|
||||||
|
let mut stream = io::BufWriter::new(stream);
|
||||||
|
|
||||||
|
let metadata = FormatMetadata {
|
||||||
|
compressed: compress
|
||||||
|
.then_some(Compressed::Zstd)
|
||||||
|
.unwrap_or(Compressed::No),
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
stream.write_all(FormatMetadata::MAGIC_NUM)?;
|
||||||
|
metadata.write_to(&mut stream)?;
|
||||||
|
|
||||||
|
let write = |stream: &mut (dyn io::Write)| serde_cbor::to_writer(stream, chain).expect("Failed to write chain to output stream"); // TODO: Error type
|
||||||
|
|
||||||
|
let mut stream = match metadata.compressed {
|
||||||
|
Compressed::No => {
|
||||||
|
write(&mut stream);
|
||||||
|
stream
|
||||||
|
},
|
||||||
|
Compressed::Zstd => {
|
||||||
|
let mut stream = zstd::Encoder::new(stream, 22)?;
|
||||||
|
|
||||||
|
#[cfg(feature="threads")]
|
||||||
|
stream.multithread(num_cpus::get() as u32)?;
|
||||||
|
|
||||||
|
write(&mut stream);
|
||||||
|
// XXX: Should we flush after write here..?
|
||||||
|
|
||||||
|
// NOTE: Required here.
|
||||||
|
stream.finish()?
|
||||||
|
},
|
||||||
|
};
|
||||||
|
stream.flush()
|
||||||
|
}
|
||||||
|
|
||||||
|
//TODO: Add `tokio_uring` version of `save_chain_to_file()`/`load_chain_from_file()` that spawns the `tokio_uring` runtime internally to queue reads/writes from/to a file.
|
Loading…
Reference in new issue