wove/src/io.rs
2024-10-23 13:54:54 -04:00

143 lines
2.8 KiB
Rust

use futures_lite::Stream;
use crate::aliases::IoResult;
use std::{future::Future, mem::MaybeUninit};
pub trait Transpose<B, T> {
fn transpose(self) -> IoResult<(B, T)>;
}
impl<B, T> Transpose<B, T> for (B, IoResult<T>) {
fn transpose(self) -> IoResult<(B, T)> {
let (buf, res) = self;
res.map(|v| (buf, v))
}
}
pub trait BufferMut {
fn as_mut_ptr(&mut self) -> *mut u8;
fn writable_bytes(&self) -> usize;
fn resize_to_fit(&mut self, written_bytes: usize);
}
impl BufferMut for Vec<u8> {
fn as_mut_ptr(&mut self) -> *mut u8 {
self.as_mut_ptr()
}
fn writable_bytes(&self) -> usize {
self.capacity()
}
fn resize_to_fit(&mut self, written_bytes: usize) {
if written_bytes > self.capacity() {
panic!();
}
unsafe { self.set_len(written_bytes) };
}
}
fn resize_box<T>(bx: &mut Box<[T]>, size: usize) {
let mut bx_tmp = MaybeUninit::uninit();
let bx_tmp_ptr: *mut Box<[T]> = bx_tmp.as_mut_ptr();
let bx_tmp = unsafe {
bx_tmp_ptr.copy_from_nonoverlapping(bx as *const _, 1);
bx_tmp.assume_init()
};
let mut vec = bx_tmp.into_vec();
vec.truncate(size);
unsafe {
(bx as *mut Box<[T]>).copy_from_nonoverlapping(&vec.into_boxed_slice(), 1);
}
}
impl BufferMut for Box<[u8]> {
fn as_mut_ptr(&mut self) -> *mut u8 {
self[..].as_mut_ptr()
}
fn writable_bytes(&self) -> usize {
self.len()
}
fn resize_to_fit(&mut self, written_bytes: usize) {
resize_box(self, written_bytes)
}
}
pub trait AsyncReadLoan {
fn read_no_resize<B: BufferMut>(
&mut self,
buf: B,
) -> impl Future<Output = (B, IoResult<usize>)>;
fn read<B: BufferMut>(&mut self, buf: B) -> impl Future<Output = (B, IoResult<usize>)> {
async {
let (mut buf, amt) = self.read_no_resize(buf).await;
let (buf, amt) = match amt {
Ok(amt) => (
{
buf.resize_to_fit(amt);
buf
},
Ok(amt),
),
Err(e) => (buf, Err(e)),
};
(buf, amt)
}
}
fn read_stream(&mut self) -> impl Stream<Item = IoResult<Vec<u8>>> {
futures_lite::stream::unfold(self, move |s| async move {
let (buf, res) = s.read(vec![0; 4096]).await;
match res {
Err(e) => Some((Err(e), s)),
Ok(0) => None,
Ok(_) => Some((Ok(buf), s)),
}
})
}
}
pub trait Buffer {
fn as_ptr(&self) -> *const u8;
fn readable_bytes(&self) -> usize;
}
impl Buffer for Vec<u8> {
fn as_ptr(&self) -> *const u8 {
self.as_ptr()
}
fn readable_bytes(&self) -> usize {
self.len()
}
}
impl Buffer for Box<[u8]> {
fn as_ptr(&self) -> *const u8 {
self[..].as_ptr()
}
fn readable_bytes(&self) -> usize {
self.len()
}
}
impl Buffer for &'static [u8] {
fn as_ptr(&self) -> *const u8 {
self[..].as_ptr()
}
fn readable_bytes(&self) -> usize {
self.len()
}
}
pub trait AsyncWriteLoan {
fn write<B: Buffer>(&mut self, buf: B) -> impl Future<Output = (B, IoResult<usize>)>;
}