Skip to content

Commit 600e902

Browse files
Add BorrowDatum for unsizing borrows of datums (pgcentralfoundation#1891)
Add a trait that allows borrowing Datums from various places, like arguments and arrays, instead of always taking them by-value. This is most important for unsizing borrows, as the sized kind are a simple `.cast()` away. The motivating rationale is to enable new APIs in the near future like FlatArray and Text. It also brings in some more uniformity because of the blanket-impl of ArgAbi for `&T`. This means that various types that _are_ pass-by-reference may now actually be taken by reference from their Datum, instead of copying them into the Rust function! This is more important for unsized types, but for various reasons, right now it only takes over the ArgAbi of one unsized type: CStr.
1 parent 1fe2fde commit 600e902

File tree

7 files changed

+261
-5
lines changed

7 files changed

+261
-5
lines changed
Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
use pgrx::prelude::*;
2+
3+
// BorrowDatum only describes something as a valid argument, but every &T has a U we CAN return,
4+
// going from &T as arg to T as ret type. We also don't necessarily have SPI support for each type.
5+
// Thus we don't *exactly* reuse the roundtrip test macro.
6+
7+
macro_rules! clonetrip {
8+
($fname:ident, $tname:ident, &$lt:lifetime $rtype:ty, &$expected:expr) => {
9+
#[pg_extern]
10+
fn $fname<$lt>(i: &$lt $rtype) -> $rtype {
11+
i.clone()
12+
}
13+
14+
clonetrip_test!($fname, $tname, &'_ $rtype, $expected);
15+
};
16+
($fname:ident, $tname:ident, &$lt:lifetime $ref_ty:ty => $own_ty:ty, $test:expr, $value:expr) => {
17+
#[pg_extern]
18+
fn $fname(i: &$ref_ty) -> $own_ty {
19+
i.into()
20+
}
21+
22+
clonetrip_test!($fname, $tname, &'_ $ref_ty => $own_ty, $test, $value);
23+
};
24+
}
25+
26+
macro_rules! clonetrip_test {
27+
($fname:ident, $tname:ident, &'_ $rtype:ty, $expected:expr) => {
28+
#[pg_test]
29+
fn $tname() -> Result<(), Box<dyn std::error::Error>> {
30+
let expected: $rtype = $expected;
31+
let result: $rtype = Spi::get_one_with_args(
32+
&format!("SELECT {}($1)", stringify!(tests.$fname)),
33+
vec![(PgOid::from(<$rtype>::type_oid()), expected.into_datum())],
34+
)?
35+
.unwrap();
36+
37+
assert_eq!(&$expected, &result);
38+
Ok(())
39+
}
40+
};
41+
($fname:ident, $tname:ident, $ref_ty:ty => $own_ty:ty, $test:expr, $value:expr) => {
42+
#[pg_test]
43+
fn $tname() -> Result<(), Box<dyn std::error::Error>> {
44+
let value: $own_ty = $value;
45+
let result: $own_ty = Spi::get_one_with_args(
46+
&format!("SELECT {}($1)", stringify!(tests.$fname)),
47+
vec![(PgOid::from(<$ref_ty>::type_oid()), value.into_datum())],
48+
)?
49+
.unwrap();
50+
51+
assert_eq!($test, &*result);
52+
Ok(())
53+
}
54+
};
55+
}
56+
57+
#[cfg(any(test, feature = "pg_test"))]
58+
#[pg_schema]
59+
mod tests {
60+
use super::*;
61+
#[allow(unused)]
62+
use crate as pgrx_tests;
63+
use std::ffi::{CStr, CString};
64+
65+
// Exercising BorrowDatum impls
66+
clonetrip!(clone_bool, test_clone_bool, &'a bool, &false);
67+
clonetrip!(clone_i8, test_clone_i8, &'a i8, &i8::MIN);
68+
clonetrip!(clone_i16, test_clone_i16, &'a i16, &i16::MIN);
69+
clonetrip!(clone_i32, test_clone_i32, &'a i32, &-1i32);
70+
clonetrip!(clone_f64, test_clone_f64, &'a f64, &f64::NEG_INFINITY);
71+
clonetrip!(
72+
clone_point,
73+
test_clone_point,
74+
&'a pg_sys::Point,
75+
&pg_sys::Point { x: -1.0, y: f64::INFINITY }
76+
);
77+
clonetrip!(clone_str, test_clone_str, &'a CStr => CString, c"cee string", CString::from(c"cee string"));
78+
clonetrip!(
79+
clone_oid,
80+
test_clone_oid,
81+
&'a pg_sys::Oid,
82+
&pg_sys::Oid::from(pg_sys::BuiltinOid::RECORDOID)
83+
);
84+
}

pgrx-tests/src/tests/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ mod attributes_tests;
1616
mod bgworker_tests;
1717
#[cfg(feature = "cshim")]
1818
mod bindings_of_inline_fn_tests;
19+
mod borrow_datum;
1920
mod bytea_tests;
2021
mod cfg_tests;
2122
mod complex;

pgrx/src/callconv.rs

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,15 @@
1010
#![deny(unsafe_op_in_unsafe_fn)]
1111
//! Helper implementations for returning sets and tables from `#[pg_extern]`-style functions
1212
13-
use crate::datum::Datum;
1413
use crate::datum::{
1514
AnyArray, AnyElement, AnyNumeric, Date, FromDatum, Inet, Internal, Interval, IntoDatum, Json,
1615
JsonB, Numeric, PgVarlena, Time, TimeWithTimeZone, Timestamp, TimestampWithTimeZone,
1716
UnboxDatum, Uuid,
1817
};
18+
use crate::datum::{BorrowDatum, Datum};
1919
use crate::datum::{Range, RangeSubType};
2020
use crate::heap_tuple::PgHeapTuple;
21+
use crate::layout::PassBy;
2122
use crate::nullable::Nullable;
2223
use crate::pg_sys;
2324
use crate::pgbox::*;
@@ -265,7 +266,38 @@ argue_from_datum! { 'fcx; Date, Interval, Time, TimeWithTimeZone, Timestamp, Tim
265266
argue_from_datum! { 'fcx; AnyArray, AnyElement, AnyNumeric }
266267
argue_from_datum! { 'fcx; Inet, Internal, Json, JsonB, Uuid, PgRelation }
267268
argue_from_datum! { 'fcx; pg_sys::BOX, pg_sys::ItemPointerData, pg_sys::Oid, pg_sys::Point }
268-
argue_from_datum! { 'fcx; &'fcx str, &'fcx CStr, &'fcx [u8] }
269+
// We could use the upcoming impl of ArgAbi for `&'fcx T where T: ?Sized + BorrowDatum`
270+
// to support these types by implementing BorrowDatum for them also, but we reject this.
271+
// It would greatly complicate other users of BorrowDatum like FlatArray, which want all impls
272+
// of BorrowDatum to return a borrow of the entire pointee's len.
273+
argue_from_datum! { 'fcx; &'fcx str, &'fcx [u8] }
274+
275+
unsafe impl<'fcx, T> ArgAbi<'fcx> for &'fcx T
276+
where
277+
T: ?Sized + BorrowDatum,
278+
{
279+
unsafe fn unbox_arg_unchecked(arg: Arg<'_, 'fcx>) -> &'fcx T {
280+
let ptr: *mut u8 = match T::PASS {
281+
PassBy::Ref => arg.2.value.cast_mut_ptr(),
282+
PassBy::Value => ptr::addr_of!(arg.0.raw_args()[arg.1].value).cast_mut().cast(),
283+
};
284+
unsafe {
285+
let ptr = ptr::NonNull::new_unchecked(ptr);
286+
T::borrow_unchecked(ptr)
287+
}
288+
}
289+
290+
unsafe fn unbox_nullable_arg(arg: Arg<'_, 'fcx>) -> Nullable<Self> {
291+
let ptr: Option<ptr::NonNull<u8>> = NonNull::new(match T::PASS {
292+
PassBy::Ref => arg.2.value.cast_mut_ptr(),
293+
PassBy::Value => ptr::addr_of!(arg.0.raw_args()[arg.1].value).cast_mut().cast(),
294+
});
295+
match (arg.is_null(), ptr) {
296+
(true, _) | (false, None) => Nullable::Null,
297+
(false, Some(ptr)) => unsafe { Nullable::Valid(T::borrow_unchecked(ptr)) },
298+
}
299+
}
300+
}
269301

270302
/// How to return a value from Rust to Postgres
271303
///

pgrx/src/datum/borrow.rs

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
#![deny(unsafe_op_in_unsafe_fn)]
2+
use super::*;
3+
use crate::layout::PassBy;
4+
use core::{ffi, mem, ptr};
5+
6+
/// Types which can be "borrowed from" [`&Datum<'_>`] via simple cast, deref, or slicing
7+
///
8+
/// # Safety
9+
/// Despite its pleasant-sounding name, this implements a fairly low-level detail.
10+
/// It exists to allow other code to use that nice-sounding BorrowDatum bound.
11+
/// Outside of the pgrx library, it is probably incorrect to call and rely on this:
12+
/// instead use convenience functions like `Datum::borrow_as`.
13+
///
14+
/// Its behavior is trusted for ABI details, and it should not be implemented if any doubt
15+
/// exists of whether the type would be suitable for passing via Postgres.
16+
pub unsafe trait BorrowDatum {
17+
/// The "native" passing convention for this type.
18+
///
19+
/// - `PassBy::Value` implies [`mem::size_of<T>()`][size_of] <= [`mem::size_of::<Datum>()`][Datum].
20+
/// - `PassBy::Ref` means the pointee will occupy at least 1 byte for variable-sized types.
21+
///
22+
/// Note that this means a zero-sized type is inappropriate for `BorrowDatum`.
23+
const PASS: PassBy;
24+
25+
/// Cast a pointer to this blob of bytes to a pointer to this type.
26+
///
27+
/// This is not a simple `ptr.cast()` because it may be *unsizing*, which may require
28+
/// reading varlena headers. For all fixed-size types, `ptr.cast()` should be correct.
29+
///
30+
/// # Safety
31+
/// - This must be correctly invoked for the pointee type, as it may deref and read one or more
32+
/// bytes in its implementation in order to read the inline metadata and unsize the type.
33+
/// - This must be invoked with a pointee initialized for the dynamically specified length.
34+
///
35+
/// ## For Implementers
36+
/// Reading the **first** byte pointed to is permitted if `T::PASS = PassBy::Ref`, assuming you
37+
/// are implementing a varlena type. As the other dynamic length type, CStr also does this.
38+
/// This function
39+
/// - must NOT mutate the pointee
40+
/// - must point to the entire datum's length (`size_of_val` must not lose bytes)
41+
///
42+
/// Do not attempt to handle pass-by-value versus pass-by-ref in this fn's body!
43+
/// A caller may be in a context where all types are handled by-reference, for instance.
44+
unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self>;
45+
46+
/// Cast a pointer to aligned varlena headers to this type
47+
///
48+
/// This version allows you to assume the pointer is aligned to, and readable for, 4 bytes.
49+
/// This optimization is not required. When in doubt, avoid implementing it, and rely on your
50+
/// `point_from` implementation alone.
51+
///
52+
/// # Safety
53+
/// - This must be correctly invoked for the pointee type, as it may deref.
54+
/// - This must be 4-byte aligned!
55+
unsafe fn point_from_align4(ptr: ptr::NonNull<u32>) -> ptr::NonNull<Self> {
56+
debug_assert!(ptr.is_aligned());
57+
unsafe { BorrowDatum::point_from(ptr.cast()) }
58+
}
59+
60+
/// Optimization for borrowing the referent
61+
unsafe fn borrow_unchecked<'dat>(ptr: ptr::NonNull<u8>) -> &'dat Self {
62+
unsafe { BorrowDatum::point_from(ptr).as_ref() }
63+
}
64+
}
65+
66+
/// From a pointer to a Datum, obtain a pointer to T's bytes
67+
///
68+
/// This may be None if T is PassBy::Ref
69+
///
70+
/// # Safety
71+
/// Assumes the Datum is init
72+
pub(crate) unsafe fn datum_ptr_to_bytes<T>(ptr: ptr::NonNull<Datum<'_>>) -> Option<ptr::NonNull<u8>>
73+
where
74+
T: BorrowDatum,
75+
{
76+
match T::PASS {
77+
// Ptr<Datum> casts to Ptr<T>
78+
PassBy::Value => Some(ptr.cast()),
79+
// Ptr<Datum> derefs to Datum which to Ptr
80+
PassBy::Ref => unsafe {
81+
let datum = ptr.read();
82+
let ptr = ptr::NonNull::new(datum.sans_lifetime().cast_mut_ptr());
83+
ptr
84+
},
85+
}
86+
}
87+
88+
macro_rules! impl_borrow_fixed_len {
89+
($($value_ty:ty),*) => {
90+
$(
91+
unsafe impl BorrowDatum for $value_ty {
92+
const PASS: PassBy = if mem::size_of::<Self>() <= mem::size_of::<Datum>() {
93+
PassBy::Value
94+
} else {
95+
PassBy::Ref
96+
};
97+
98+
unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self> {
99+
ptr.cast()
100+
}
101+
}
102+
)*
103+
}
104+
}
105+
106+
impl_borrow_fixed_len! {
107+
i8, i16, i32, i64, bool, f32, f64,
108+
pg_sys::Oid, pg_sys::Point,
109+
Date, Time, Timestamp, TimestampWithTimeZone
110+
}
111+
112+
/// It is rare to pass CStr via Datums, but not unheard of
113+
unsafe impl BorrowDatum for ffi::CStr {
114+
const PASS: PassBy = PassBy::Ref;
115+
116+
unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self> {
117+
let char_ptr: *mut ffi::c_char = ptr.as_ptr().cast();
118+
unsafe {
119+
let len = ffi::CStr::from_ptr(char_ptr).to_bytes_with_nul().len();
120+
ptr::NonNull::new_unchecked(ptr::slice_from_raw_parts_mut(char_ptr, len) as *mut Self)
121+
}
122+
}
123+
124+
unsafe fn borrow_unchecked<'dat>(ptr: ptr::NonNull<u8>) -> &'dat Self {
125+
let char_ptr: *const ffi::c_char = ptr.as_ptr().cast();
126+
unsafe { ffi::CStr::from_ptr(char_ptr) }
127+
}
128+
}

pgrx/src/datum/mod.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
mod anyarray;
1717
mod anyelement;
1818
mod array;
19+
mod borrow;
1920
mod date;
2021
pub mod datetime_support;
2122
mod from;
@@ -44,6 +45,7 @@ pub use self::uuid::*;
4445
pub use anyarray::*;
4546
pub use anyelement::*;
4647
pub use array::*;
48+
pub use borrow::*;
4749
pub use date::*;
4850
pub use datetime_support::*;
4951
pub use from::*;
@@ -63,6 +65,7 @@ pub use varlena::*;
6365
use crate::memcx::MemCx;
6466
use crate::pg_sys;
6567
use core::marker::PhantomData;
68+
use core::ptr;
6669
#[doc(hidden)]
6770
pub use with_typeid::nonstatic_typeid;
6871
pub use with_typeid::{WithArrayTypeIds, WithSizedTypeIds, WithTypeIds, WithVarlenaTypeIds};
@@ -136,14 +139,22 @@ pub struct Datum<'src>(
136139
);
137140

138141
impl<'src> Datum<'src> {
139-
/// The Datum without its lifetime.
142+
/// Strip a Datum of its lifetime for FFI purposes.
140143
pub fn sans_lifetime(self) -> pg_sys::Datum {
141144
self.0
142145
}
143146
/// Construct a Datum containing only a null pointer.
144147
pub fn null() -> Datum<'src> {
145148
Self(pg_sys::Datum::from(0), PhantomData)
146149
}
150+
151+
/// Reborrow the Datum as `T`
152+
///
153+
/// If the type is `PassBy::Ref`, this may be `None`.
154+
pub unsafe fn borrow_as<T: BorrowDatum>(&self) -> Option<&T> {
155+
let ptr = ptr::NonNull::new_unchecked(ptr::from_ref(self).cast_mut());
156+
borrow::datum_ptr_to_bytes::<T>(ptr).map(|ptr| BorrowDatum::borrow_unchecked(ptr))
157+
}
147158
}
148159

149160
/// A tagging trait to indicate a user type is also meant to be used by Postgres

pgrx/src/layout.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ pub(crate) struct Layout {
3535
}
3636

3737
#[derive(Clone, Copy, Debug, PartialEq)]
38-
pub(crate) enum PassBy {
38+
pub enum PassBy {
3939
Ref,
4040
Value,
4141
}

pgrx/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ pub mod htup;
5555
pub mod inoutfuncs;
5656
pub mod itemptr;
5757
pub mod iter;
58+
pub mod layout;
5859
pub mod list;
5960
pub mod lwlock;
6061
pub mod memcx;
@@ -79,7 +80,6 @@ pub mod wrappers;
7980
pub mod xid;
8081

8182
/// Not ready for public exposure.
82-
mod layout;
8383
mod ptr;
8484
mod slice;
8585
mod toast;

0 commit comments

Comments
 (0)