diff --git a/pg-extend/src/native/bytea.rs b/pg-extend/src/native/bytea.rs new file mode 100644 index 00000000..261a84a4 --- /dev/null +++ b/pg-extend/src/native/bytea.rs @@ -0,0 +1,46 @@ +use std::ops::Deref; +use std::ptr::NonNull; + +use crate::native::VarLenA; +use crate::pg_alloc::{PgAllocated, PgAllocator}; +use crate::pg_sys; + +/// A zero-overhead view of `bytea` data from Postgres +pub struct ByteA<'mc>(PgAllocated<'mc, NonNull>); + +// :consider would be good to make a derive(FromVarLenA) macro. +impl<'mc> ByteA<'mc> { + /// Create from the raw pointer to the Postgres data + #[allow(clippy::missing_safety_doc)] + pub unsafe fn from_raw(alloc: &'mc PgAllocator, raw_ptr: *mut pg_sys::bytea) -> Self { + ByteA(PgAllocated::from_raw(alloc, raw_ptr)) + } + + /// Convert into the underlying pointer + #[allow(clippy::missing_safety_doc)] + pub unsafe fn into_ptr(mut self) -> *mut pg_sys::bytea { + self.0.take_ptr() + } + + /// Return true if this is empty + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Return the length of the bytea data + pub fn len(&self) -> usize { + let varlena = unsafe { VarLenA::from_varlena(self.0.as_ref()) }; + varlena.len() + } +} + +impl<'mc> Deref for ByteA<'mc> { + type Target = [u8]; + + fn deref(&self) -> &[u8] { + unsafe { + let varlena = VarLenA::from_varlena(self.0.as_ref()); + &*(varlena.as_slice() as *const [i8] as *const [u8]) + } + } +} diff --git a/pg-extend/src/native/mod.rs b/pg-extend/src/native/mod.rs index f0121c41..f3bd93b3 100644 --- a/pg-extend/src/native/mod.rs +++ b/pg-extend/src/native/mod.rs @@ -9,8 +9,10 @@ //! //! These shoudl be near zero overhead types, exposed from Postgres and able to be directly used. +mod bytea; mod text; mod varlena; +pub use bytea::ByteA; pub use text::Text; pub(crate) use varlena::VarLenA; diff --git a/pg-extend/src/pg_alloc.rs b/pg-extend/src/pg_alloc.rs index 19426802..5938f874 100644 --- a/pg-extend/src/pg_alloc.rs +++ b/pg-extend/src/pg_alloc.rs @@ -203,6 +203,9 @@ impl RawPtr for std::ffi::CString { } } +// FIXME +// `pg_sys::text` aliases `varlena`. This impl covers all `varlena` types, including `bytea`. +// Given that `Target` is an associated type, we are stuck with `pg_sys::text` for everything. impl RawPtr for NonNull { type Target = pg_sys::text; diff --git a/pg-extend/src/pg_datum.rs b/pg-extend/src/pg_datum.rs index b1c30fe6..d9ce9f74 100644 --- a/pg-extend/src/pg_datum.rs +++ b/pg-extend/src/pg_datum.rs @@ -11,8 +11,9 @@ use std::ffi::{CStr, CString}; use std::marker::PhantomData; use std::os::raw::c_char; use std::ptr::NonNull; +use std::convert::TryInto; -use crate::native::Text; +use crate::native::{ByteA, Text, VarLenA}; use crate::pg_alloc::{PgAllocated, PgAllocator}; use crate::pg_bool; use crate::pg_sys::{self, Datum}; @@ -211,6 +212,111 @@ impl<'s> TryFromPgDatum<'s> for String { } } +// FIXME: this impl is copy-pasta'd from that of Text. Could dedup with macro. +impl<'s> From> for PgDatum<'s> { + fn from(value: ByteA<'s>) -> Self { + let ptr = unsafe { value.into_ptr() }; + PgDatum(Some(ptr as Datum), PhantomData) + } +} + +// FIXME: this impl is copy-pasta'd from that of Text. Could dedup with macro. +impl<'s> TryFromPgDatum<'s> for ByteA<'s> { + fn try_from<'mc>( + memory_context: &'mc PgAllocator, + datum: PgDatum<'mc>, + ) -> Result + where + Self: 's, + 'mc: 's, + { + if let Some(datum) = datum.0 { + let text_ptr = datum as *const pg_sys::bytea; + + unsafe { Ok(ByteA::from_raw(memory_context, text_ptr as *mut _)) } + } else { + Err("datum was NULL") + } + } +} + +impl<'s> TryFromPgDatum<'s> for &[u8] { + fn try_from<'mc>( + memory_context: &'mc PgAllocator, + datum: PgDatum<'mc>, + ) -> Result + where + Self: 's, + 'mc: 's, + { + let ba = as TryFromPgDatum>::try_from(memory_context, datum)?; + let v: &[u8] = &ba; + // FIXME: transmutes to extend the lifetime to ~ 'mc. + // there is probably a better way to do this? + let mv: &[u8] = unsafe{ std::mem::transmute(v) }; + Ok(mv) + } +} + +// :note could just use T: AsRef<[u8]> if specialization existed. +// as it is, too many types implement AsRef<[u8]>, so we +// manually dispatch to &[u8] impl. +impl From> for PgDatum<'_> { + fn from(value: Vec) -> Self { + let v: &[u8] = &value; + v.into() + } +} + +// maybe macro is overkill. +macro_rules! _sizeof_varattrib_4b_header { () => { 4 }; } + +impl From<&[u8]> for PgDatum<'_> { + fn from(value: &[u8]) -> Self { + let ptr = unsafe{ + // :consider should palloc be guard_pg'd? + pg_sys::palloc0(value.len() + _sizeof_varattrib_4b_header!()) as *mut u8 + }; + + if ptr.is_null() { + // :fixme no alloc; better freak out properly! + unimplemented!() + } + + // :note vis. alignment, palloc docs says always 4-word aligned. + unsafe { + std::ptr::copy_nonoverlapping( + value.as_ptr(), + ptr.offset(_sizeof_varattrib_4b_header!()), + value.len() + ); + } + + // assign header length + // :note assumes little-endian + let mut l: u32 = ( + value.len() + _sizeof_varattrib_4b_header!() + ).try_into().unwrap(); + l &= 0x3fffffff; + l <<= 2; + + unsafe { + let ptr_len = ptr as *mut u32; + *ptr_len = l; + } + + debug_assert!( + unsafe { + let varl = VarLenA::from_varlena(std::mem::transmute(ptr)); + value.len() == varl.len() + }, + "length mismatch on varlena encoded byte slice." + ); + + PgDatum(Some(ptr as Datum), PhantomData) + } +} + // FIXME: this lifetime is wrong impl From for PgDatum<'_> { fn from(value: String) -> Self { diff --git a/pg-extend/src/pg_type.rs b/pg-extend/src/pg_type.rs index 5fb4c399..2516401e 100644 --- a/pg-extend/src/pg_type.rs +++ b/pg-extend/src/pg_type.rs @@ -1,6 +1,6 @@ //! Postgres type definitions -use crate::native::Text; +use crate::native::{ByteA, Text}; /// See https://www.postgresql.org/docs/11/xfunc-c.html#XFUNC-C-TYPE-TABLE /// @@ -220,6 +220,24 @@ impl PgTypeInfo for i64 { } } +impl PgTypeInfo for Vec { + fn pg_type() -> PgType { + PgType::ByteA + } +} + +impl PgTypeInfo for &[u8] { + fn pg_type() -> PgType { + PgType::ByteA + } +} + +impl PgTypeInfo for ByteA<'_> { + fn pg_type() -> PgType { + PgType::ByteA + } +} + impl PgTypeInfo for String { fn pg_type() -> PgType { PgType::Text