Use Layout::size_matches::<T> for Array cleanup (#881)

* Add `Layout.size_matches::<T>` and PassBy enum

* New Layout cleans up `Array::as_slice` innards

* Relax lifetime bounds on variadic array test

This test fn erroneously asked for overly specific lifetime bounds.

Author: workingjubilee

pgx-tests/src/tests/variadic_tests.rs

@@ -13,10 +13,7 @@ mod test {
     use pgx::VariadicArray;
 
     #[pg_extern]
-    fn func_with_variadic_array_args<'a>(
-        _field: &'a str,
-        values: VariadicArray<&'a str>,
-    ) -> String {
+    fn func_with_variadic_array_args(_field: &str, values: VariadicArray<&str>) -> String {
         values.get(0).unwrap().unwrap().to_string()
     }
 }
@@ -31,10 +28,10 @@ mod tests {
 
     #[pg_test]
     fn test_func_with_variadic_array_args() {
-        let result = Spi::get_one::<&str>(
+        let result = Spi::get_one::<String>(
             "SELECT test.func_with_variadic_array_args('test', 'a', 'b', 'c');",
         )
         .expect("didn't get SPI result");
-        assert_eq!(result, "a");
+        assert_eq!(result, String::from("a"));
     }
 }

pgx/src/datum/array.rs

@@ -60,7 +60,7 @@ pub struct Array<'a, T: FromDatum> {
     datum_palloc: Option<NonNull<pg_sys::Datum>>,
     elem_slice: &'a [pg_sys::Datum],
     null_slice: NullKind<'a>,
-    elem_layout: Option<Layout>,
+    elem_layout: Layout,
     _marker: PhantomData<T>,
 }
 
@@ -129,7 +129,7 @@ impl<'a, T: FromDatum> Array<'a, T> {
     unsafe fn deconstruct_from(
         ptr: NonNull<pg_sys::varlena>,
         raw: RawArray,
-        layout: Layout,
+        elem_layout: Layout,
     ) -> Array<'a, T> {
         let oid = raw.oid();
         let len = raw.len();
@@ -153,9 +153,9 @@ impl<'a, T: FromDatum> Array<'a, T> {
             pg_sys::deconstruct_array(
                 array,
                 oid,
-                layout.size.as_typlen().into(),
-                layout.passbyval,
-                layout.align.as_typalign(),
+                elem_layout.size.as_typlen().into(),
+                matches!(elem_layout.pass, PassBy::Value),
+                elem_layout.align.as_typalign(),
                 &mut elements,
                 &mut nulls,
                 &mut nelems,
@@ -195,7 +195,7 @@ impl<'a, T: FromDatum> Array<'a, T> {
             datum_palloc: NonNull::new(elements),
             elem_slice: /* SAFETY: &[Datum] from palloc'd [Datum] */ unsafe { slice::from_raw_parts(elements, nelems) },
             null_slice,
-            elem_layout: Some(layout),
+            elem_layout,
             _marker: PhantomData,
         }
     }
@@ -217,33 +217,16 @@ impl<'a, T: FromDatum> Array<'a, T> {
         if you are sure your usage is sound, consider RawArray"
     )]
     pub fn as_slice(&self) -> &[T] {
-        if let Some(Layout { size, passbyval, .. }) = &self.elem_layout {
-            if self.null_slice.any() {
-                panic!("null detected: can't expose potentially uninit data as a slice!")
-            }
-            const DATUM_SIZE: usize = mem::size_of::<pg_sys::Datum>();
-            let sizeof_type = match (passbyval, mem::size_of::<T>(), size.try_as_usize()) {
-                (true, rs @ (1 | 2 | 4 | 8), Some(pg @ (1 | 2 | 4 | 8))) if rs == pg => rs,
-                (true, _, _) => panic!("invalid sizes for pass-by-value datum"),
-                (false, DATUM_SIZE, _) => DATUM_SIZE,
-                (false, _, _) => panic!("invalid sizes for pass-by-reference datum"),
-            };
-            match (sizeof_type, self.raw.as_ref()) {
-                // SAFETY: Rust slice layout matches Postgres data layout and this array is "owned"
-                (1 | 2 | 4, Some(raw)) => unsafe { raw.assume_init_data_slice::<T>() },
-                (DATUM_SIZE, _) => {
-                    let sizeof_datums = mem::size_of_val(self.elem_slice);
-                    unsafe {
-                        slice::from_raw_parts(
-                            self.elem_slice.as_ptr() as *const T,
-                            sizeof_datums / sizeof_type,
-                        )
-                    }
-                }
-                (_, _) => panic!("no correctly-sized slice exists"),
-            }
-        } else {
-            panic!("not enough type information to slice correctly")
+        const DATUM_SIZE: usize = mem::size_of::<pg_sys::Datum>();
+        if self.null_slice.any() {
+            panic!("null detected: can't expose potentially uninit data as a slice!")
+        }
+        match (self.elem_layout.size_matches::<T>(), self.raw.as_ref()) {
+            // SAFETY: Rust slice layout matches Postgres data layout and this array is "owned"
+            (Some(1 | 2 | 4 | DATUM_SIZE), Some(raw)) => unsafe {
+                raw.assume_init_data_slice::<T>()
+            },
+            (_, _) => panic!("no correctly-sized slice exists"),
         }
     }
 

pgx/src/layout.rs

@@ -1,6 +1,5 @@
 /*!
 Code for interfacing with the layout in memory (and on disk?) of various data types within Postgres.
-Prefer to minimize exposing the contents of this module to public callers of the pgx library:
 This is not a mature module yet so prefer to avoid exposing the contents to public callers of pgx,
 as this is error-prone stuff to tamper with. It is easy to corrupt the database if an error is made.
 Yes, even though its main block of code duplicates htup::DatumWithTypeInfo.
@@ -17,12 +16,18 @@ use core::mem;
 
 /// Postgres type information, corresponds to part of a row in pg_type
 /// This layout describes T, not &T, even if passbyval: false, which would mean the datum array is effectively &[&T]
-#[derive(Debug, Clone)]
+#[derive(Clone, Copy, Debug)]
 pub(crate) struct Layout {
     // We could add more fields to this if we are curious enough, as the function we call pulls an entire row
-    pub align: Align,    // typalign
-    pub size: Size,      // typlen
-    pub passbyval: bool, // typbyval
+    pub align: Align, // typalign
+    pub size: Size,   // typlen
+    pub pass: PassBy, // typbyval
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub(crate) enum PassBy {
+    Ref,
+    Value,
 }
 
 impl Layout {
@@ -41,13 +46,27 @@ impl Layout {
         Layout {
             align: Align::try_from(typalign).unwrap(),
             size: Size::try_from(typlen).unwrap(),
-            passbyval,
+            pass: if passbyval { PassBy::Value } else { PassBy::Ref },
+        }
+    }
+
+    // Attempt to discern if a given Postgres and Rust layout are "matching" in some sense
+    // Some(usize) if they seem to agree on one, None if they do not
+    pub(crate) fn size_matches<T>(&self) -> Option<usize> {
+        const DATUM_SIZE: usize = mem::size_of::<pg_sys::Datum>();
+        match (self.pass, mem::size_of::<T>(), self.size.try_as_usize()) {
+            (PassBy::Value, rs @ (1 | 2 | 4 | 8), Some(pg @ (1 | 2 | 4 | 8))) if rs == pg => {
+                Some(rs)
+            }
+            (PassBy::Value, _, _) => None,
+            (PassBy::Ref, DATUM_SIZE, _) => Some(DATUM_SIZE),
+            (PassBy::Ref, _, _) => None,
         }
     }
 }
 
 #[repr(usize)]
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub(crate) enum Align {
     Byte = mem::align_of::<u8>(),
     Short = mem::align_of::<libc::c_short>(),
@@ -80,7 +99,7 @@ impl Align {
     }
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub(crate) enum Size {
     CStr,
     Varlena,