Add MemCx<'current, T> arguments and PBox<'current, T> returns (pgcentralfoundation#2210)

tl;dr: `caller_picks_lifetime + macro_expansion == unsoundness`,
so prevent that by adding API for new, lifetime-bound allocations.

For functions exposed to PostgreSQL via `pg_extern`, signatures often
resemble these:

```rust
#[pg_extern]
fn do_something(
    floats: Array<'a, f32>,
    hms: Time,
    day: TimestampTz
) -> Vec<TimestampTz> {
```

```sql
-- pgrx will generate something like this
CREATE FUNCTION do_something(
    floats real[],
    hms time,
    day timestamptz
) RETURNS timestamptz[]
```


In SQL, we accept an array and return an array of a different type, but
in Rust, we return a `Vec`. This is because we lack API to allocate and
return a fresh `Array<'a, T>`s. Behind the scenes, we *do* allocate an
`ArrayType` at the translation boundary, unseen to a Rust programmer.
That, plus the Vec, means two allocations just to mutate an array, even
for a fixed-size type!

Unfortunately, pgrx-offered type translations like `Array<'a, T>`,
cannot have a lifetime *and* expose a new *safe* function to make it
from lifetime-free values! Such API must have the `unsafe` requirement
of access-exclusivity for its lifetime, described in functions like
`slice::from_raw_parts_mut`[^1], because the lifetime becomes
caller-selected. This makes it easy to use a `'static` lifetime with
`pg_extern` and expand `unsafe` code that trusts that lifetime.

A `'static` lifetime is always wrong for `palloc`ated types, even with
`TopMemoryContext`, since it dies whenever Postgres resets the context
and Postgres does have conditions under which it performs such a global
reset. The result is various extremely subtle soundness problems are
possible, yet the root cause lies in code that most never read. We can
solve this by denying the possibility of a caller-selected lifetime.
This will require removing some types, ultimately, but first we need an
API to migrate to.

This new API is `MemCx<'mcx, T>` and `PBox<'mcx, T>`. Importantly, there
is no *safe* route to creating `MemCx<'mcx, T>` with a caller-selected
lifetime. One can only reach `MemCx<'mcx, T>` safely as an argument:
either pass a closure to `memcx::current_context` or call a new
`pg_extern` function via SQL. Either way gives you `&MemCx<'current>`.
These let you create a lifetime-bound allocation that can live long
enough to return it.

We can already return fixed-size pass-by-reference types, so right now
`PBox` is only an optimization. It allows not putting a type on the
stack and then immediately moving it directly into a heap allocation.
More importantly, it enables writing tests, so it is important ground
work for returning de facto unsized values, as they can hide behind
`PBox` which is always `Sized`.

[^1]: https://doc.rust-lang.org/std/slice/fn.from_raw_parts_mut.html

Author: workingjubilee

pgrx-tests/src/tests/memcxt_tests.rs

@@ -14,6 +14,8 @@ mod tests {
     use crate as pgrx_tests;
 
     use pgrx::PgMemoryContexts;
+    use pgrx::memcx::MemCx;
+    use pgrx::palloc::PBox;
     use pgrx::pg_test;
     use pgrx::prelude::*;
     use std::sync::Arc;
@@ -29,6 +31,15 @@ mod tests {
         }
     }
 
+    #[pg_extern]
+    pub fn accept_mcx_return_timetz<'mcx>(
+        memcx: &'mcx MemCx<'mcx>,
+    ) -> PBox<'mcx, TimeWithTimeZone> {
+        let palloc = memcx.alloc_bytes(size_of::<TimeWithTimeZone>());
+        let timetz = TimeWithTimeZone::new(4, 20, 0.0).unwrap();
+        PBox::new_in(timetz, memcx)
+    }
+
     #[pg_test]
     fn test_leak_and_drop() {
         let did_drop = Arc::new(AtomicBool::new(false));

pgrx-tests/tests/compile-fail/no-static-memcx.rs

@@ -0,0 +1,11 @@
+use pgrx::memcx::MemCx;
+use pgrx::palloc::PBox;
+use pgrx::prelude::*;
+
+#[pg_extern]
+pub fn accept_mcx_return_timetz(memcx: &MemCx<'static>) -> PBox<'static, TimeWithTimeZone> {
+    let timetz = TimeWithTimeZone::new(4, 20, 0.0).unwrap();
+    PBox::new_in(timetz, memcx)
+}
+
+fn main() {}

pgrx-tests/tests/compile-fail/no-static-memcx.stderr

@@ -0,0 +1,20 @@
+error[E0521]: borrowed data escapes outside of function
+ --> tests/compile-fail/no-static-memcx.rs:6:92
+  |
+5 |   #[pg_extern]
+  |   ------------
+  |   |
+  |   lifetime `'fcx` defined here
+  |   in this procedural macro expansion
+6 |   pub fn accept_mcx_return_timetz(memcx: &MemCx<'static>) -> PBox<'static, TimeWithTimeZone> {
+  |  ____________________________________________________________________________________________^
+7 | |     let timetz = TimeWithTimeZone::new(4, 20, 0.0).unwrap();
+8 | |     PBox::new_in(timetz, memcx)
+9 | | }
+  | | ^
+  | | |
+  | | `fcinfo` is a reference that is only valid in the function body
+  | |_`fcinfo` escapes the function body here
+  |   argument requires that `'fcx` must outlive `'static`
+  |
+  = note: this error originates in the attribute macro `pg_extern` (in Nightly builds, run with -Z macro-backtrace for more info)

pgrx/src/callconv.rs

@@ -91,7 +91,7 @@ pub unsafe trait ArgAbi<'fcx>: Sized {
         }
     }
 
-    /// "Is this a virtual arg?"
+    /// Virtual args do not affect Datum arg iteration
     fn is_virtual_arg() -> bool {
         false
     }

pgrx/src/lib.rs

@@ -63,6 +63,7 @@ pub mod misc;
 pub mod namespace;
 pub mod nodes;
 pub mod nullable;
+pub mod palloc;
 pub mod pg_catalog;
 pub mod pgbox;
 pub mod rel;

pgrx/src/memcx.rs

@@ -1,15 +1,22 @@
 //! Memory Contexts in PostgreSQL, now with lifetimes.
+use pgrx_sql_entity_graph::metadata::{
+    ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
+};
+
 // "Why isn't this pgrx::mem or pgrx::memcxt?"
 // Postgres actually uses all of:
 // - mcxt
 // - memcxt
 // - mctx
 // Search engines will see "memc[tx]{2}" and assume you mean memcpy!
 // And it's nice-ish to have shorter lifetime names and have 'mcx consistently mean the lifetime.
+use crate::callconv::{Arg, ArgAbi};
+use crate::nullable::Nullable;
 use crate::pg_sys;
 use core::{marker::PhantomData, ptr::NonNull};
 
 /// A borrowed memory context.
+#[repr(transparent)]
 pub struct MemCx<'mcx> {
     ptr: NonNull<pg_sys::MemoryContextData>,
     _marker: PhantomData<&'mcx pg_sys::MemoryContextData>,
@@ -118,3 +125,38 @@ mod nightly {
         }
     }
 }
+
+unsafe impl<'fcx> ArgAbi<'fcx> for &MemCx<'fcx> {
+    unsafe fn unbox_arg_unchecked(_arg: Arg<'_, 'fcx>) -> Self {
+        // SAFETY: We are called to unbox an argument, which means the backend was initialized.
+        // We use this horrific expression to allow the lifetime to be extended arbitrarily
+        // and achieve an "in-place" transformation of CurrentMemoryContext's pointer.
+        // The soundness of this is riding on the lifetimes used for `unbox_arg_unchecked` in our macros,
+        // as the expanded code is designed so that `fcinfo` and each `arg` are truly "borrowed" in rustc's eyes.
+        unsafe { &*((&raw mut pg_sys::CurrentMemoryContext).cast()) }
+    }
+
+    unsafe fn unbox_nullable_arg(arg: Arg<'_, 'fcx>) -> Nullable<Self> {
+        // SAFETY: Should never happen in actuality, but as long as we're here...
+        if unsafe { pg_sys::CurrentMemoryContext.is_null() } {
+            Nullable::Null
+        } else {
+            Nullable::Valid(Self::unbox_arg_unchecked(arg))
+        }
+    }
+
+    fn is_virtual_arg() -> bool {
+        true
+    }
+}
+
+/// SAFETY: virtual argument
+unsafe impl<'mcx> SqlTranslatable for &MemCx<'mcx> {
+    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
+        Ok(SqlMapping::Skip)
+    }
+
+    fn return_sql() -> Result<Returns, ReturnsError> {
+        Ok(Returns::One(SqlMapping::Skip))
+    }
+}

pgrx/src/palloc.rs

@@ -0,0 +1,3 @@
+mod pbox;
+
+pub use pbox::PBox;

pgrx/src/palloc/pbox.rs

@@ -0,0 +1,108 @@
+use crate::callconv::{BoxRet, FcInfo};
+use crate::datum::{BorrowDatum, Datum};
+use crate::layout::PassBy;
+use crate::memcx::MemCx;
+use crate::pg_sys;
+use core::marker::PhantomData;
+use core::ops::{Deref, DerefMut};
+use core::ptr::NonNull;
+
+use pgrx_sql_entity_graph::metadata::{
+    ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
+};
+
+/** As [`Box<T, A>`][stdbox] where `A` is a [`MemCx`]
+
+
+[stdbox]: alloc::boxed::Box
+*/
+#[repr(transparent)]
+pub struct PBox<'mcx, T: ?Sized> {
+    ptr: NonNull<T>,
+    _cx: PhantomData<MemCx<'mcx>>,
+}
+
+impl<'mcx, T: ?Sized> PBox<'mcx, T> {
+    // # Safety
+    // The same constraints as [`Box::from_raw`], AND
+    // - you assert the pointer was allocated in the `MemCx`
+    // - you assert the pointer may be freed by `pfree`
+    pub unsafe fn from_raw_in(ptr: NonNull<T>, _cx: &MemCx<'mcx>) -> PBox<'mcx, T> {
+        PBox { ptr, _cx: PhantomData }
+    }
+}
+
+impl<'mcx, T: Sized> PBox<'mcx, T> {
+    pub fn new_in(val: T, memcx: &MemCx<'mcx>) -> Self {
+        const { assert!(align_of::<T>() <= size_of::<pg_sys::Datum>()) };
+        let ptr = memcx.alloc_bytes(size_of::<T>()).cast();
+        // SAFETY: We were guaranteed an appropriately sized allocation to write to,
+        // and we have asserted our alignment maximum was upheld
+        unsafe { ptr.write(val) };
+        PBox { ptr, _cx: PhantomData }
+    }
+}
+
+unsafe impl<'mcx, T> BoxRet for PBox<'mcx, T>
+where
+    T: ?Sized + BorrowDatum,
+{
+    unsafe fn box_into<'fcx>(self, fcinfo: &mut FcInfo<'fcx>) -> Datum<'fcx> {
+        let datum = match T::PASS {
+            PassBy::Value => {
+                // start with a zeroed Datum, just to minimize funny business
+                let mut datum = pg_sys::Datum::null();
+                // SAFETY: Due to BorrowDatum, this type has a definite size less than a Datum,
+                // and PBox must have an initialized pointee, so a copy is sound-by-construction.
+                unsafe {
+                    let size = size_of_val(&*self.ptr.as_ptr());
+                    debug_assert!(size <= size_of::<pg_sys::Datum>());
+                    // using `BorrowDatum::point_from` handles endianness
+                    let datum_ptr = T::point_from(NonNull::from_mut(&mut datum).cast::<u8>());
+                    datum_ptr.cast::<u8>().copy_from_nonoverlapping(self.ptr.cast(), size);
+                }
+                datum
+            }
+            PassBy::Ref => pg_sys::Datum::from(self.ptr.cast::<u8>().as_ptr()),
+        };
+        // SAFETY: by proxy, BorrowDatum is an `unsafe trait` so the above impl must be correct
+        unsafe { fcinfo.return_raw_datum(datum) }
+    }
+}
+
+/// SAFETY: SQL has no "pointers" so by-val and by-ref calling conventions are identical,
+/// and all `PBox` truly does is enable pass-by-ref returns of unsized values.
+unsafe impl<'mcx, T> SqlTranslatable for PBox<'mcx, T>
+where
+    T: SqlTranslatable + ?Sized,
+{
+    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
+        T::argument_sql()
+    }
+
+    fn return_sql() -> Result<Returns, ReturnsError> {
+        T::return_sql()
+    }
+}
+
+impl<'mcx, T> Deref for PBox<'mcx, T>
+where
+    T: BorrowDatum + ?Sized,
+{
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: by construction
+        unsafe { self.ptr.as_ref() }
+    }
+}
+
+impl<'mcx, T> DerefMut for PBox<'mcx, T>
+where
+    T: BorrowDatum + ?Sized,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        // SAFETY: by construction
+        unsafe { self.ptr.as_mut() }
+    }
+}