Make unboxing args sound with ArgAbi

pgrx-examples/composite_type/src/lib.rs

@@ -145,6 +145,7 @@ fn make_friendship(
 }
 
 /*
+FIXME: make this example no longer DOA
 Create sum the scritches received by dogs, roughly the equivalent of:
 
 ```sql
@@ -161,64 +162,7 @@ CREATE AGGREGATE sum_scritches ("value" Dog) (
 )
 ```
 */
-struct SumScritches;
-
-#[pg_aggregate]
-impl Aggregate for SumScritches {
-    type State = i32;
-    const INITIAL_CONDITION: Option<&'static str> = Some("0");
-    type Args = pgrx::name!(value, pgrx::composite_type!('static, "Dog"));
-
-    fn state(
-        current: Self::State,
-        arg: Self::Args,
-        _fcinfo: pg_sys::FunctionCallInfo,
-    ) -> Self::State {
-        let arg_scritches: i32 = arg
-            .get_by_name("scritches")
-            .unwrap() // Unwrap the result of the conversion
-            .unwrap_or_default(); // The number of scritches, or 0 if there was none set
-        current + arg_scritches
-    }
-}
-
-/*
-Create sum the scritches received by dogs, roughly the equivalent of:
-
-```sql
-CREATE FUNCTION scritch_collector_state(state Dog, new integer)
-    RETURNS Dog
-    LANGUAGE SQL
-    STRICT
-    RETURN ROW(state.name, state.scritches + new)::Dog;
-
-CREATE AGGREGATE scritch_collector ("value" integer) (
-    SFUNC = "sum_scritches_state",
-    STYPE = Dog,
-)
-```
-*/
-struct ScritchCollector;
-
-#[pg_aggregate]
-impl Aggregate for ScritchCollector {
-    type State = Option<pgrx::composite_type!('static, "Dog")>;
-    type Args = i32;
-
-    fn state(
-        current: Self::State,
-        arg: Self::Args,
-        _fcinfo: pg_sys::FunctionCallInfo,
-    ) -> Self::State {
-        let mut current = match current {
-            Some(v) => v,
-            None => PgHeapTuple::new_composite_type(DOG_COMPOSITE_TYPE).unwrap(),
-        };
-        let current_scritches: i32 = current.get_by_name("scritches").unwrap().unwrap_or_default();
-        current.set_by_name("scritches", current_scritches + arg).unwrap();
-        Some(current)
-    }
-}
+// struct SumScritches;
 
 /*
 Create an operator allowing dogs to accept scritches directly.
@@ -298,14 +242,6 @@ mod tests {
         Ok(())
     }
 
-    #[pg_test]
-    fn test_scritch_collector() {
-        let retval = Spi::get_one::<i32>(
-            "SELECT (scritchcollector(value)).scritches FROM UNNEST(ARRAY [1,2,3]) as value;",
-        );
-        assert_eq!(retval, Ok(Some(6)));
-    }
-
     #[pg_test]
     fn test_dog_add_operator() {
         let retval = Spi::get_one::<i32>("SELECT (ROW('Nami', 0)::Dog + 1).scritches;");

pgrx-examples/custom_types/src/hexint.rs

@@ -7,7 +7,7 @@
 //LICENSE All rights reserved.
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
-use pgrx::callconv::BoxRet;
+use pgrx::callconv::{ArgAbi, BoxRet};
 use pgrx::pg_sys::{Datum, Oid};
 use pgrx::pgrx_sql_entity_graph::metadata::{
     ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
@@ -94,6 +94,15 @@ impl IntoDatum for HexInt {
     }
 }
 
+unsafe impl<'fcx> ArgAbi<'fcx> for HexInt
+where
+    Self: 'fcx,
+{
+    unsafe fn unbox_arg_unchecked(arg: ::pgrx::callconv::Arg<'_, 'fcx>) -> Self {
+        unsafe { arg.unbox_arg_using_from_datum().unwrap() }
+    }
+}
+
 unsafe impl BoxRet for HexInt {
     unsafe fn box_in_fcinfo(self, _fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
         Datum::from(self.value)

pgrx-examples/strings/src/lib.rs

@@ -34,7 +34,7 @@ fn append(mut input: String, extra: &str) -> String {
 }
 
 #[pg_extern]
-fn split(input: &'static str, pattern: &str) -> Vec<&'static str> {
+fn split<'a>(input: &'a str, pattern: &str) -> Vec<&'a str> {
     input.split_terminator(pattern).collect()
 }
 

pgrx-macros/src/lib.rs

@@ -643,7 +643,7 @@ pub fn postgres_enum(input: TokenStream) -> TokenStream {
 fn impl_postgres_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
     let mut stream = proc_macro2::TokenStream::new();
     let sql_graph_entity_ast = ast.clone();
-    let generics = &ast.generics;
+    let generics = &ast.generics.clone();
     let enum_ident = &ast.ident;
     let enum_name = enum_ident.to_string();
 
@@ -666,6 +666,24 @@ fn impl_postgres_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
         into_datum.extend(quote! { #enum_ident::#label_ident => Some(::pgrx::enum_helper::lookup_enum_by_label(#enum_name, #label_string)), });
     }
 
+    // We need another variant of the params for the ArgAbi impl
+    let fcx_lt = syn::Lifetime::new("'fcx", proc_macro2::Span::mixed_site());
+    let mut generics_with_fcx = generics.clone();
+    // so that we can bound on Self: 'fcx
+    generics_with_fcx.make_where_clause().predicates.push(syn::WherePredicate::Type(
+        syn::PredicateType {
+            lifetimes: None,
+            bounded_ty: syn::parse_quote! { Self },
+            colon_token: syn::Token![:](proc_macro2::Span::mixed_site()),
+            bounds: syn::parse_quote! { #fcx_lt },
+        },
+    ));
+    let (impl_gens, ty_gens, where_clause) = generics_with_fcx.split_for_impl();
+    let mut impl_gens: syn::Generics = syn::parse_quote! { #impl_gens };
+    impl_gens
+        .params
+        .insert(0, syn::GenericParam::Lifetime(syn::LifetimeParam::new(fcx_lt.clone())));
+
     stream.extend(quote! {
         impl ::pgrx::datum::FromDatum for #enum_ident {
             #[inline]
@@ -683,6 +701,14 @@ fn impl_postgres_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
             }
         }
 
+        unsafe impl #impl_gens ::pgrx::callconv::ArgAbi<#fcx_lt> for #enum_ident #ty_gens #where_clause {
+            unsafe fn unbox_arg_unchecked(arg: ::pgrx::callconv::Arg<'_, #fcx_lt>) -> Self {
+                let index = arg.index();
+                unsafe { arg.unbox_arg_using_from_datum().unwrap_or_else(|| panic!("argument {index} must not be null")) }
+            }
+
+        }
+
         unsafe impl #generics ::pgrx::datum::UnboxDatum for #enum_ident #generics {
             type As<'dat> = #enum_ident #generics where Self: 'dat;
             #[inline]
@@ -761,7 +787,7 @@ pub fn postgres_type(input: TokenStream) -> TokenStream {
 
 fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
     let name = &ast.ident;
-    let generics = &ast.generics;
+    let generics = &ast.generics.clone();
     let has_lifetimes = generics.lifetimes().next();
     let funcname_in = Ident::new(&format!("{name}_in").to_lowercase(), name.span());
     let funcname_out = Ident::new(&format!("{name}_out").to_lowercase(), name.span());
@@ -794,6 +820,24 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
         None => quote! {'_},
     };
 
+    // We need another variant of the params for the ArgAbi impl
+    let fcx_lt = syn::Lifetime::new("'fcx", proc_macro2::Span::mixed_site());
+    let mut generics_with_fcx = generics.clone();
+    // so that we can bound on Self: 'fcx
+    generics_with_fcx.make_where_clause().predicates.push(syn::WherePredicate::Type(
+        syn::PredicateType {
+            lifetimes: None,
+            bounded_ty: syn::parse_quote! { Self },
+            colon_token: syn::Token![:](proc_macro2::Span::mixed_site()),
+            bounds: syn::parse_quote! { #fcx_lt },
+        },
+    ));
+    let (impl_gens, ty_gens, where_clause) = generics_with_fcx.split_for_impl();
+    let mut impl_gens: syn::Generics = syn::parse_quote! { #impl_gens };
+    impl_gens
+        .params
+        .insert(0, syn::GenericParam::Lifetime(syn::LifetimeParam::new(fcx_lt.clone())));
+
     // all #[derive(PostgresType)] need to implement that trait
     // and also the FromDatum and IntoDatum
     stream.extend(quote! {
@@ -861,6 +905,14 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
                         <Self as ::pgrx::datum::FromDatum>::from_datum(::core::mem::transmute(datum), false).unwrap()
                     }
                 }
+
+                unsafe impl #impl_gens ::pgrx::callconv::ArgAbi<#fcx_lt> for #name #ty_gens #where_clause
+                {
+                        unsafe fn unbox_arg_unchecked(arg: ::pgrx::callconv::Arg<'_, #fcx_lt>) -> Self {
+                        let index = arg.index();
+                        unsafe { arg.unbox_arg_using_from_datum().unwrap_or_else(|| panic!("argument {index} must not be null")) }
+                    }
+                }
             }
         )
     }

pgrx-pg-sys/src/submodules/datum.rs

@@ -65,6 +65,11 @@ impl Datum {
         sptr::Strict::addr(self.0)
     }
 
+    #[inline]
+    pub const fn null() -> Datum {
+        Datum(core::ptr::null_mut())
+    }
+
     /// True if the datum is equal to the null pointer.
     #[inline]
     pub fn is_null(self) -> bool {

pgrx-sql-entity-graph/src/pg_extern/mod.rs

@@ -44,7 +44,7 @@ use quote::{format_ident, quote, quote_spanned};
 use syn::parse::{Parse, ParseStream, Parser};
 use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
-use syn::{Meta, Token, Type};
+use syn::{Meta, Token};
 
 /// A parsed `#[pg_extern]` item.
 ///
@@ -374,72 +374,75 @@ impl PgExtern {
     }
 
     pub fn wrapper_func(&self) -> Result<syn::ItemFn, syn::Error> {
-        let func_name = &self.func.sig.ident;
-        let is_raw = self.extern_attrs().contains(&Attribute::Raw);
-        // We use a `_` prefix to make functions with no args more satisfied during linting.
-        let fcinfo_ident = syn::Ident::new("_fcinfo", self.func.sig.ident.span());
-        let lifetimes =
-            self.func.sig.generics.lifetimes().collect::<syn::punctuated::Punctuated<_, Comma>>();
-        let fc_lt = syn::Lifetime::new("'fcx", Span::mixed_site());
+        let signature = &self.func.sig;
+        let func_name = &signature.ident;
+        // we do this odd dance so we can pass the same ident to macros that don't know each other
+        let fcinfo_ident = syn::Ident::new("fcinfo", signature.ident.span());
+        let mut lifetimes = signature
+            .generics
+            .lifetimes()
+            .cloned()
+            .collect::<syn::punctuated::Punctuated<_, Comma>>();
+        // we pick an arbitrary lifetime from the provided signature of the fn, if available,
+        // so lifetime-bound fn are easier to write with pgrx
+        let fc_lt = lifetimes
+            .first()
+            .map(|lt_p| lt_p.lifetime.clone())
+            .filter(|lt| lt.ident != "static")
+            .unwrap_or(syn::Lifetime::new("'fcx", Span::mixed_site()));
+        // we need the bare lifetime later, but we also jam it into the bounds if it is new
         let fc_ltparam = syn::LifetimeParam::new(fc_lt.clone());
+        if lifetimes.first() != Some(&fc_ltparam) {
+            lifetimes.insert(0, fc_ltparam)
+        }
 
         let args = &self.inputs;
+        // for unclear reasons the linker vomits if we don't do this
         let arg_pats = args.iter().map(|v| format_ident!("{}_", &v.pat)).collect::<Vec<_>>();
-        let arg_fetches = args.iter().enumerate().map(|(idx, arg)| {
-            let pat = &arg_pats[idx];
-            let resolved_ty = &arg.used_ty.resolved_ty;
-            match resolved_ty {
-                // There's no danger of misinterpreting the type, as pointer coercions must typecheck!
-                Type::Path(ty_path) if ty_path.last_ident_is("FunctionCallInfo") => quote_spanned! { pat.span() =>
-                    let #pat = #fcinfo_ident;
-                },
-                Type::Tuple(tup) if tup.elems.is_empty() => quote_spanned! { pat.span() =>
-                    debug_assert!(unsafe { ::pgrx::fcinfo::pg_getarg::<()>(#fcinfo_ident, #idx).is_none() }, "A `()` argument should always receive `NULL`");
-                    let #pat = ();
-                },
-                _ => match (is_raw, &arg.used_ty.optional) {
-                    (true, None) | (true, Some(_)) => quote_spanned! { pat.span() =>
-                        let #pat = unsafe { ::pgrx::fcinfo::pg_getarg_datum_raw(#fcinfo_ident, #idx) as #resolved_ty };
-                    },
-                    (false, None) => quote_spanned! { pat.span() =>
-                        let #pat = unsafe { ::pgrx::fcinfo::pg_getarg::<#resolved_ty>(#fcinfo_ident, #idx).unwrap_or_else(|| panic!("{} is null", stringify!{#pat})) };
-                    },
-                    (false, Some(inner)) => quote_spanned! { pat.span() =>
-                        let #pat = unsafe { ::pgrx::fcinfo::pg_getarg::<#inner>(#fcinfo_ident, #idx) };
-                    },
+        let args_ident = proc_macro2::Ident::new("_args", Span::call_site());
+        let arg_fetches = arg_pats.iter().map(|pat| {
+                quote_spanned!{ pat.span() =>
+                    let #pat = #args_ident.next_arg_unchecked().unwrap_or_else(|| panic!("unboxing {} argument failed", stringify!(#pat)));
                 }
             }
-        });
+        );
 
         match &self.returns {
             Returning::None
             | Returning::Type(_)
             | Returning::SetOf { .. }
             | Returning::Iterated { .. } => {
-                let ret_ty = match &self.func.sig.output {
+                let ret_ty = match &signature.output {
                     syn::ReturnType::Default => syn::parse_quote! { () },
                     syn::ReturnType::Type(_, ret_ty) => ret_ty.clone(),
                 };
                 let wrapper_code = quote_spanned! { self.func.block.span() =>
-                    fn _internal_wrapper<#fc_ltparam, #lifetimes>(fcinfo: ::pgrx::callconv::Fcinfo<#fc_lt>) -> ::pgrx::datum::Datum<#fc_lt> {
-                    #[allow(unused_unsafe)]
-                    unsafe {
-                        let #fcinfo_ident = fcinfo.0;
-                        let result = match <#ret_ty as ::pgrx::callconv::RetAbi>::check_fcinfo_and_prepare(#fcinfo_ident) {
-                            ::pgrx::callconv::CallCx::WrappedFn(mcx) => {
-                                let mut mcx = ::pgrx::PgMemoryContexts::For(mcx);
-                                ::pgrx::callconv::RetAbi::to_ret(mcx.switch_to(|_| {
-                                    #(#arg_fetches)*
-                                    #func_name( #(#arg_pats),* )
-                                }))
-                            }
-                            ::pgrx::callconv::CallCx::RestoreCx => <#ret_ty as ::pgrx::callconv::RetAbi>::ret_from_fcinfo_fcx(#fcinfo_ident),
-                        };
-                        ::core::mem::transmute(unsafe { <#ret_ty as ::pgrx::callconv::RetAbi>::box_ret_in_fcinfo(#fcinfo_ident, result) })
-                    }}
-                    let fcinfo = ::pgrx::callconv::Fcinfo(#fcinfo_ident, ::core::marker::PhantomData);
+                    fn _internal_wrapper<#lifetimes>(fcinfo: &mut ::pgrx::callconv::FcInfo<#fc_lt>) -> ::pgrx::datum::Datum<#fc_lt> {
+                        #[allow(unused_unsafe)]
+                        unsafe {
+                            let call_flow = <#ret_ty as ::pgrx::callconv::RetAbi>::check_and_prepare(fcinfo);
+                            let result = match call_flow {
+                                ::pgrx::callconv::CallCx::WrappedFn(mcx) => {
+                                    let mut mcx = ::pgrx::PgMemoryContexts::For(mcx);
+                                    let #args_ident = &mut fcinfo.args();
+                                    let call_result = mcx.switch_to(|_| {
+                                        #(#arg_fetches)*
+                                        #func_name( #(#arg_pats),* )
+                                    });
+                                    ::pgrx::callconv::RetAbi::to_ret(call_result)
+                                }
+                                ::pgrx::callconv::CallCx::RestoreCx => <#ret_ty as ::pgrx::callconv::RetAbi>::ret_from_fcx(fcinfo),
+                            };
+                            unsafe { <#ret_ty as ::pgrx::callconv::RetAbi>::box_ret_in(fcinfo, result) }
+                        }
+                    }
                     // We preserve the invariants
-                    let datum = unsafe { ::pgrx::pg_sys::submodules::panic::pgrx_extern_c_guard(move || _internal_wrapper(fcinfo)) };
+                    let datum = unsafe {
+                        ::pgrx::pg_sys::submodules::panic::pgrx_extern_c_guard(|| {
+                            let mut fcinfo = ::pgrx::callconv::FcInfo::from_ptr(#fcinfo_ident);
+                            _internal_wrapper(&mut fcinfo)
+                        })
+                    };
                     datum.sans_lifetime()
                 };
                 finfo_v1_extern_c(&self.func, fcinfo_ident, wrapper_code)