Skip to content

Commit

Permalink
make MapSlice more general
Browse files Browse the repository at this point in the history
  • Loading branch information
MarinPostma committed Sep 30, 2024
1 parent 726c40c commit 912aaf3
Showing 1 changed file with 42 additions and 14 deletions.
56 changes: 42 additions & 14 deletions libsql-wal/src/io/buf.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
// from tokio uring

use std::mem::{size_of, MaybeUninit};
use std::{
borrow::Borrow,
marker::PhantomData,
mem::{size_of, MaybeUninit},
};

use bytes::BytesMut;
use bytes::{Bytes, BytesMut};
use zerocopy::{AsBytes, FromBytes};

pub unsafe trait IoBuf: Unpin + 'static {
Expand Down Expand Up @@ -90,6 +94,20 @@ unsafe impl IoBuf for BytesMut {
}
}

unsafe impl IoBuf for Bytes {
fn stable_ptr(&self) -> *const u8 {
self.as_ptr()
}

fn bytes_init(&self) -> usize {
self.len()
}

fn bytes_total(&self) -> usize {
self.len()
}
}

unsafe impl IoBufMut for BytesMut {
fn stable_mut_ptr(&mut self) -> *mut u8 {
self.as_mut_ptr()
Expand Down Expand Up @@ -185,11 +203,11 @@ impl<T> ZeroCopyBuf<T> {
}
}

pub fn map_slice<F>(self, f: F) -> MapSlice<T, F>
pub fn map_slice<F>(self, f: F) -> MapSlice<Self, F, T>
where
for<'a> F: Fn(&'a Self) -> &'a [u8] + Unpin + 'static,
{
MapSlice { inner: self, f }
MapSlice::new(self, f)
}

#[inline]
Expand Down Expand Up @@ -220,32 +238,42 @@ impl<T> ZeroCopyBuf<T> {
}
}

pub struct MapSlice<T, F> {
inner: ZeroCopyBuf<T>,
pub struct MapSlice<T, F, U> {
inner: T,
f: F,
_p: PhantomData<U>,
}

impl<T, F> MapSlice<T, F> {
pub(crate) fn into_inner(self) -> ZeroCopyBuf<T> {
impl<T, F, U> MapSlice<T, F, U> {
pub(crate) fn into_inner(self) -> T {
self.inner
}

pub(crate) fn new(inner: T, f: F) -> Self {
Self {
inner,
f,
_p: PhantomData,
}
}
}

unsafe impl<T, F> IoBuf for MapSlice<T, F>
unsafe impl<T, F, U> IoBuf for MapSlice<T, F, U>
where
for<'a> F: Fn(&'a ZeroCopyBuf<T>) -> &'a [u8] + Unpin + 'static,
T: Unpin + 'static + AsBytes,
for<'a> F: Fn(&'a ZeroCopyBuf<U>) -> &'a [u8] + Unpin + 'static,
T: Borrow<ZeroCopyBuf<U>> + Unpin + 'static,
U: AsBytes + Unpin + 'static,
{
fn stable_ptr(&self) -> *const u8 {
(self.f)(&self.inner).as_ptr()
(self.f)(&self.inner.borrow()).as_ptr()
}

fn bytes_init(&self) -> usize {
(self.f)(&self.inner).len()
(self.f)(&self.inner.borrow()).len()
}

fn bytes_total(&self) -> usize {
(self.f)(&self.inner).len()
(self.f)(&self.inner.borrow()).len()
}
}

Expand Down

0 comments on commit 912aaf3

Please sign in to comment.