Add BorrowDatum for unsizing borrows of datums

pgrx-tests/src/tests/borrow_datum.rs

@@ -0,0 +1,84 @@
+use pgrx::prelude::*;
+
+// BorrowDatum only describes something as a valid argument, but every &T has a U we CAN return,
+// going from &T as arg to T as ret type. We also don't necessarily have SPI support for each type.
+// Thus we don't *exactly* reuse the roundtrip test macro.
+
+macro_rules! clonetrip {
+    ($fname:ident, $tname:ident, &$lt:lifetime $rtype:ty, &$expected:expr) => {
+        #[pg_extern]
+        fn $fname<$lt>(i: &$lt $rtype) -> $rtype {
+            i.clone()
+        }
+
+        clonetrip_test!($fname, $tname, &'_ $rtype, $expected);
+    };
+    ($fname:ident, $tname:ident, &$lt:lifetime $ref_ty:ty => $own_ty:ty, $test:expr, $value:expr) => {
+        #[pg_extern]
+        fn $fname(i: &$ref_ty) -> $own_ty {
+            i.into()
+        }
+
+        clonetrip_test!($fname, $tname, &'_ $ref_ty => $own_ty, $test, $value);
+    };
+}
+
+macro_rules! clonetrip_test {
+    ($fname:ident, $tname:ident, &'_ $rtype:ty, $expected:expr) => {
+        #[pg_test]
+        fn $tname() -> Result<(), Box<dyn std::error::Error>> {
+            let expected: $rtype = $expected;
+            let result: $rtype = Spi::get_one_with_args(
+                &format!("SELECT {}($1)", stringify!(tests.$fname)),
+                vec![(PgOid::from(<$rtype>::type_oid()), expected.into_datum())],
+            )?
+            .unwrap();
+
+            assert_eq!(&$expected, &result);
+            Ok(())
+        }
+    };
+    ($fname:ident, $tname:ident, $ref_ty:ty => $own_ty:ty, $test:expr, $value:expr) => {
+        #[pg_test]
+        fn $tname() -> Result<(), Box<dyn std::error::Error>> {
+            let value: $own_ty = $value;
+            let result: $own_ty = Spi::get_one_with_args(
+                &format!("SELECT {}($1)", stringify!(tests.$fname)),
+                vec![(PgOid::from(<$ref_ty>::type_oid()), value.into_datum())],
+            )?
+            .unwrap();
+
+            assert_eq!($test, &*result);
+            Ok(())
+        }
+    };
+}
+
+#[cfg(any(test, feature = "pg_test"))]
+#[pg_schema]
+mod tests {
+    use super::*;
+    #[allow(unused)]
+    use crate as pgrx_tests;
+    use std::ffi::{CStr, CString};
+
+    // Exercising BorrowDatum impls
+    clonetrip!(clone_bool, test_clone_bool, &'a bool, &false);
+    clonetrip!(clone_i8, test_clone_i8, &'a i8, &i8::MIN);
+    clonetrip!(clone_i16, test_clone_i16, &'a i16, &i16::MIN);
+    clonetrip!(clone_i32, test_clone_i32, &'a i32, &-1i32);
+    clonetrip!(clone_f64, test_clone_f64, &'a f64, &f64::NEG_INFINITY);
+    clonetrip!(
+        clone_point,
+        test_clone_point,
+        &'a pg_sys::Point,
+        &pg_sys::Point { x: -1.0, y: f64::INFINITY }
+    );
+    clonetrip!(clone_str, test_clone_str, &'a CStr => CString, c"cee string", CString::from(c"cee string"));
+    clonetrip!(
+        clone_oid,
+        test_clone_oid,
+        &'a pg_sys::Oid,
+        &pg_sys::Oid::from(pg_sys::BuiltinOid::RECORDOID)
+    );
+}

pgrx-tests/src/tests/mod.rs

@@ -16,6 +16,7 @@ mod attributes_tests;
 mod bgworker_tests;
 #[cfg(feature = "cshim")]
 mod bindings_of_inline_fn_tests;
+mod borrow_datum;
 mod bytea_tests;
 mod cfg_tests;
 mod complex;

pgrx/src/callconv.rs

@@ -10,14 +10,15 @@
 #![deny(unsafe_op_in_unsafe_fn)]
 //! Helper implementations for returning sets and tables from `#[pg_extern]`-style functions
 
-use crate::datum::Datum;
 use crate::datum::{
     AnyArray, AnyElement, AnyNumeric, Date, FromDatum, Inet, Internal, Interval, IntoDatum, Json,
     JsonB, Numeric, PgVarlena, Time, TimeWithTimeZone, Timestamp, TimestampWithTimeZone,
     UnboxDatum, Uuid,
 };
+use crate::datum::{BorrowDatum, Datum};
 use crate::datum::{Range, RangeSubType};
 use crate::heap_tuple::PgHeapTuple;
+use crate::layout::PassBy;
 use crate::nullable::Nullable;
 use crate::pg_sys;
 use crate::pgbox::*;
@@ -265,7 +266,38 @@ argue_from_datum! { 'fcx; Date, Interval, Time, TimeWithTimeZone, Timestamp, Tim
 argue_from_datum! { 'fcx; AnyArray, AnyElement, AnyNumeric }
 argue_from_datum! { 'fcx; Inet, Internal, Json, JsonB, Uuid, PgRelation }
 argue_from_datum! { 'fcx; pg_sys::BOX, pg_sys::ItemPointerData, pg_sys::Oid, pg_sys::Point }
-argue_from_datum! { 'fcx; &'fcx str, &'fcx CStr, &'fcx [u8] }
+// We could use the upcoming impl of ArgAbi for `&'fcx T where T: ?Sized + BorrowDatum`
+// to support these types by implementing BorrowDatum for them also, but we reject this.
+// It would greatly complicate other users of BorrowDatum like FlatArray, which want all impls
+// of BorrowDatum to return a borrow of the entire pointee's len.
+argue_from_datum! { 'fcx; &'fcx str, &'fcx [u8] }
+
+unsafe impl<'fcx, T> ArgAbi<'fcx> for &'fcx T
+where
+    T: ?Sized + BorrowDatum,
+{
+    unsafe fn unbox_arg_unchecked(arg: Arg<'_, 'fcx>) -> &'fcx T {
+        let ptr: *mut u8 = match T::PASS {
+            PassBy::Ref => arg.2.value.cast_mut_ptr(),
+            PassBy::Value => ptr::addr_of!(arg.0.raw_args()[arg.1].value).cast_mut().cast(),
+        };
+        unsafe {
+            let ptr = ptr::NonNull::new_unchecked(ptr);
+            T::borrow_unchecked(ptr)
+        }
+    }
+
+    unsafe fn unbox_nullable_arg(arg: Arg<'_, 'fcx>) -> Nullable<Self> {
+        let ptr: Option<ptr::NonNull<u8>> = NonNull::new(match T::PASS {
+            PassBy::Ref => arg.2.value.cast_mut_ptr(),
+            PassBy::Value => ptr::addr_of!(arg.0.raw_args()[arg.1].value).cast_mut().cast(),
+        });
+        match (arg.is_null(), ptr) {
+            (true, _) | (false, None) => Nullable::Null,
+            (false, Some(ptr)) => unsafe { Nullable::Valid(T::borrow_unchecked(ptr)) },
+        }
+    }
+}
 
 /// How to return a value from Rust to Postgres
 ///

pgrx/src/datum/borrow.rs

@@ -0,0 +1,128 @@
+#![deny(unsafe_op_in_unsafe_fn)]
+use super::*;
+use crate::layout::PassBy;
+use core::{ffi, mem, ptr};
+
+/// Types which can be "borrowed from" [`&Datum<'_>`] via simple cast, deref, or slicing
+///
+/// # Safety
+/// Despite its pleasant-sounding name, this implements a fairly low-level detail.
+/// It exists to allow other code to use that nice-sounding BorrowDatum bound.
+/// Outside of the pgrx library, it is probably incorrect to call and rely on this:
+/// instead use convenience functions like `Datum::borrow_as`.
+///
+/// Its behavior is trusted for ABI details, and it should not be implemented if any doubt
+/// exists of whether the type would be suitable for passing via Postgres.
+pub unsafe trait BorrowDatum {
+    /// The "native" passing convention for this type.
+    ///
+    /// - `PassBy::Value` implies [`mem::size_of<T>()`][size_of] <= [`mem::size_of::<Datum>()`][Datum].
+    /// - `PassBy::Ref` means the pointee will occupy at least 1 byte for variable-sized types.
+    ///
+    /// Note that this means a zero-sized type is inappropriate for `BorrowDatum`.
+    const PASS: PassBy;
+
+    /// Cast a pointer to this blob of bytes to a pointer to this type.
+    ///
+    /// This is not a simple `ptr.cast()` because it may be *unsizing*, which may require
+    /// reading varlena headers. For all fixed-size types, `ptr.cast()` should be correct.
+    ///
+    /// # Safety
+    /// - This must be correctly invoked for the pointee type, as it may deref and read one or more
+    ///   bytes in its implementation in order to read the inline metadata and unsize the type.
+    /// - This must be invoked with a pointee initialized for the dynamically specified length.
+    ///
+    /// ## For Implementers
+    /// Reading the **first** byte pointed to is permitted if `T::PASS = PassBy::Ref`, assuming you
+    /// are implementing a varlena type. As the other dynamic length type, CStr also does this.
+    /// This function
+    /// - must NOT mutate the pointee
+    /// - must point to the entire datum's length (`size_of_val` must not lose bytes)
+    ///
+    /// Do not attempt to handle pass-by-value versus pass-by-ref in this fn's body!
+    /// A caller may be in a context where all types are handled by-reference, for instance.
+    unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self>;
+
+    /// Cast a pointer to aligned varlena headers to this type
+    ///
+    /// This version allows you to assume the pointer is aligned to, and readable for, 4 bytes.
+    /// This optimization is not required. When in doubt, avoid implementing it, and rely on your
+    /// `point_from` implementation alone.
+    ///
+    /// # Safety
+    /// - This must be correctly invoked for the pointee type, as it may deref.
+    /// - This must be 4-byte aligned!
+    unsafe fn point_from_align4(ptr: ptr::NonNull<u32>) -> ptr::NonNull<Self> {
+        debug_assert!(ptr.is_aligned());
+        unsafe { BorrowDatum::point_from(ptr.cast()) }
+    }
+
+    /// Optimization for borrowing the referent
+    unsafe fn borrow_unchecked<'dat>(ptr: ptr::NonNull<u8>) -> &'dat Self {
+        unsafe { BorrowDatum::point_from(ptr).as_ref() }
+    }
+}
+
+/// From a pointer to a Datum, obtain a pointer to T's bytes
+///
+/// This may be None if T is PassBy::Ref
+///
+/// # Safety
+/// Assumes the Datum is init
+pub(crate) unsafe fn datum_ptr_to_bytes<T>(ptr: ptr::NonNull<Datum<'_>>) -> Option<ptr::NonNull<u8>>
+where
+    T: BorrowDatum,
+{
+    match T::PASS {
+        // Ptr<Datum> casts to Ptr<T>
+        PassBy::Value => Some(ptr.cast()),
+        // Ptr<Datum> derefs to Datum which to Ptr
+        PassBy::Ref => unsafe {
+            let datum = ptr.read();
+            let ptr = ptr::NonNull::new(datum.sans_lifetime().cast_mut_ptr());
+            ptr
+        },
+    }
+}
+
+macro_rules! impl_borrow_fixed_len {
+    ($($value_ty:ty),*) => {
+        $(
+            unsafe impl BorrowDatum for $value_ty {
+                const PASS: PassBy = if mem::size_of::<Self>() <= mem::size_of::<Datum>() {
+                    PassBy::Value
+                } else {
+                    PassBy::Ref
+                };
+
+                unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self> {
+                    ptr.cast()
+                }
+            }
+        )*
+    }
+}
+
+impl_borrow_fixed_len! {
+    i8, i16, i32, i64, bool, f32, f64,
+    pg_sys::Oid, pg_sys::Point,
+    Date, Time, Timestamp, TimestampWithTimeZone
+}
+
+/// It is rare to pass CStr via Datums, but not unheard of
+unsafe impl BorrowDatum for ffi::CStr {
+    const PASS: PassBy = PassBy::Ref;
+
+    unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self> {
+        let char_ptr: *mut ffi::c_char = ptr.as_ptr().cast();
+        unsafe {
+            let len = ffi::CStr::from_ptr(char_ptr).to_bytes_with_nul().len();
+            ptr::NonNull::new_unchecked(ptr::slice_from_raw_parts_mut(char_ptr, len) as *mut Self)
+        }
+    }
+
+    unsafe fn borrow_unchecked<'dat>(ptr: ptr::NonNull<u8>) -> &'dat Self {
+        let char_ptr: *const ffi::c_char = ptr.as_ptr().cast();
+        unsafe { ffi::CStr::from_ptr(char_ptr) }
+    }
+}

pgrx/src/datum/mod.rs

@@ -16,6 +16,7 @@
 mod anyarray;
 mod anyelement;
 mod array;
+mod borrow;
 mod date;
 pub mod datetime_support;
 mod from;
@@ -44,6 +45,7 @@ pub use self::uuid::*;
 pub use anyarray::*;
 pub use anyelement::*;
 pub use array::*;
+pub use borrow::*;
 pub use date::*;
 pub use datetime_support::*;
 pub use from::*;
@@ -63,6 +65,7 @@ pub use varlena::*;
 use crate::memcx::MemCx;
 use crate::pg_sys;
 use core::marker::PhantomData;
+use core::ptr;
 #[doc(hidden)]
 pub use with_typeid::nonstatic_typeid;
 pub use with_typeid::{WithArrayTypeIds, WithSizedTypeIds, WithTypeIds, WithVarlenaTypeIds};
@@ -136,14 +139,22 @@ pub struct Datum<'src>(
 );
 
 impl<'src> Datum<'src> {
-    /// The Datum without its lifetime.
+    /// Strip a Datum of its lifetime for FFI purposes.
     pub fn sans_lifetime(self) -> pg_sys::Datum {
         self.0
     }
     /// Construct a Datum containing only a null pointer.
     pub fn null() -> Datum<'src> {
         Self(pg_sys::Datum::from(0), PhantomData)
     }
+
+    /// Reborrow the Datum as `T`
+    ///
+    /// If the type is `PassBy::Ref`, this may be `None`.
+    pub unsafe fn borrow_as<T: BorrowDatum>(&self) -> Option<&T> {
+        let ptr = ptr::NonNull::new_unchecked(ptr::from_ref(self).cast_mut());
+        borrow::datum_ptr_to_bytes::<T>(ptr).map(|ptr| BorrowDatum::borrow_unchecked(ptr))
+    }
 }
 
 /// A tagging trait to indicate a user type is also meant to be used by Postgres

pgrx/src/layout.rs

@@ -35,7 +35,7 @@ pub(crate) struct Layout {
 }
 
 #[derive(Clone, Copy, Debug, PartialEq)]
-pub(crate) enum PassBy {
+pub enum PassBy {
     Ref,
     Value,
 }

pgrx/src/lib.rs

@@ -55,6 +55,7 @@ pub mod htup;
 pub mod inoutfuncs;
 pub mod itemptr;
 pub mod iter;
+pub mod layout;
 pub mod list;
 pub mod lwlock;
 pub mod memcx;
@@ -79,7 +80,6 @@ pub mod wrappers;
 pub mod xid;
 
 /// Not ready for public exposure.
-mod layout;
 mod ptr;
 mod slice;
 mod toast;