pgrx/datum/
array.rs

1//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
2//LICENSE
3//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
4//LICENSE
5//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
6//LICENSE
7//LICENSE All rights reserved.
8//LICENSE
9//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
10#![allow(clippy::question_mark)]
11use super::{unbox, UnboxDatum};
12use crate::array::RawArray;
13use crate::nullable::{
14    BitSliceNulls, IntoNullableIterator, MaybeStrictNulls, NullLayout, Nullable, NullableContainer,
15};
16use crate::toast::Toast;
17use crate::{layout::*, nullable};
18use crate::{pg_sys, FromDatum, IntoDatum, PgMemoryContexts};
19use core::fmt::{Debug, Formatter};
20use core::ops::DerefMut;
21use core::ptr::NonNull;
22use pgrx_sql_entity_graph::metadata::{
23    ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
24};
25use serde::{Serialize, Serializer};
26use std::iter::FusedIterator;
27
28/** An array of some type (eg. `TEXT[]`, `int[]`)
29
30While conceptually similar to a [`Vec<T>`][std::vec::Vec], arrays are lazy.
31
32Using a [`Vec<T>`][std::vec::Vec] here means each element of the passed array will be eagerly fetched and converted into a Rust type:
33
34```rust,no_run
35use pgrx::prelude::*;
36
37#[pg_extern]
38fn with_vec(elems: Vec<String>) {
39    // Elements all already converted.
40    for elem in elems {
41        todo!()
42    }
43}
44```
45
46Using an array, elements are only fetched and converted into a Rust type on demand:
47
48```rust,no_run
49use pgrx::prelude::*;
50
51#[pg_extern]
52fn with_vec(elems: Array<String>) {
53    // Elements converted one by one
54    for maybe_elem in elems {
55        let elem = maybe_elem.unwrap();
56        todo!()
57    }
58}
59```
60*/
61// An array is a detoasted varlena type, so we reason about the lifetime of
62// the memory context that the varlena is actually detoasted into.
63pub struct Array<'mcx, T> {
64    null_slice: MaybeStrictNulls<BitSliceNulls<'mcx>>,
65    slide_impl: ChaChaSlideImpl<T>,
66    // Rust drops in FIFO order, drop this last
67    raw: Toast<RawArray>,
68}
69
70impl<T: UnboxDatum> Debug for Array<'_, T>
71where
72    for<'arr> <T as UnboxDatum>::As<'arr>: Debug,
73{
74    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
75        f.debug_list().entries(self.iter()).finish()
76    }
77}
78
79type ChaChaSlideImpl<T> = Box<dyn casper::ChaChaSlide<T>>;
80
81impl<'mcx, T: UnboxDatum> serde::Serialize for Array<'mcx, T>
82where
83    for<'arr> <T as UnboxDatum>::As<'arr>: Serialize,
84{
85    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
86    where
87        S: Serializer,
88    {
89        serializer.collect_seq(self.iter())
90    }
91}
92
93#[deny(unsafe_op_in_unsafe_fn)]
94impl<'mcx, T: UnboxDatum> Array<'mcx, T> {
95    /// # Safety
96    ///
97    /// This function requires that the RawArray was obtained in a properly-constructed form
98    /// (probably from Postgres).
99    pub(crate) unsafe fn deconstruct_from(mut raw: Toast<RawArray>) -> Array<'mcx, T> {
100        let oid = raw.oid();
101        let elem_layout = Layout::lookup_oid(oid);
102        let null_inner = raw
103            .nulls_bitslice()
104            .map(|nonnull| unsafe { nullable::BitSliceNulls(&*nonnull.as_ptr()) });
105        let null_slice = MaybeStrictNulls::new(null_inner);
106        // do a little two-step before jumping into the Cha-Cha Slide and figure out
107        // which implementation is correct for the type of element in this Array.
108        let slide_impl: ChaChaSlideImpl<T> = match elem_layout.pass {
109            PassBy::Value => match elem_layout.size {
110                // The layout is one that we know how to handle efficiently.
111                Size::Fixed(1) => Box::new(casper::FixedSizeByVal::<1>),
112                Size::Fixed(2) => Box::new(casper::FixedSizeByVal::<2>),
113                Size::Fixed(4) => Box::new(casper::FixedSizeByVal::<4>),
114                #[cfg(target_pointer_width = "64")]
115                Size::Fixed(8) => Box::new(casper::FixedSizeByVal::<8>),
116
117                _ => {
118                    panic!("unrecognized pass-by-value array element layout: {elem_layout:?}")
119                }
120            },
121
122            PassBy::Ref => match elem_layout.size {
123                // Array elements are varlenas, which are pass-by-reference and have a known alignment size
124                Size::Varlena => Box::new(casper::PassByVarlena { align: elem_layout.align }),
125
126                // Array elements are C strings, which are pass-by-reference and alignments are
127                // determined at runtime based on the length of the string
128                Size::CStr => Box::new(casper::PassByCStr),
129
130                // Array elements are fixed sizes yet the data is "pass-by-reference"
131                // Most commonly, this is because of elements larger than a Datum.
132                Size::Fixed(size) => Box::new(casper::PassByFixed {
133                    padded_size: elem_layout.align.pad(size.into()),
134                }),
135            },
136        };
137
138        Array { raw, slide_impl, null_slice }
139    }
140
141    /// Return an iterator of `Option<T>`.
142    #[inline]
143    pub fn iter(&self) -> ArrayIterator<'_, T> {
144        let ptr = self.raw.data_ptr();
145        ArrayIterator { array: self, curr: 0, ptr }
146    }
147
148    /// Return an iterator over the Array's elements.
149    ///
150    /// # Panics
151    /// This function will panic when called if the array contains any SQL NULL values.
152    #[inline]
153    pub fn iter_deny_null(&self) -> ArrayTypedIterator<'_, T> {
154        if self.null_slice.contains_nulls() {
155            panic!("array contains NULL");
156        }
157
158        let ptr = self.raw.data_ptr();
159        ArrayTypedIterator { array: self, curr: 0, ptr }
160    }
161
162    /// Retrieve a value from a known-not-null (strict) array.
163    fn get_strict_inner<'arr>(&'arr self, index: usize) -> Option<T::As<'arr>> {
164        if index >= self.raw.len() {
165            return None;
166        };
167
168        // This pointer is what's walked over the entire array's data buffer.
169        // If the array has varlena or cstr elements, we can't index into the array.
170        // If the elements are fixed size, we could, but we do not exploit that optimization yet
171        // as it would significantly complicate the code and impact debugging it.
172        // Such improvements should wait until a later version.
173        //
174        // This remains true even in a strict array, because, again,
175        // varlena and cstr are variable-length.
176        let mut at_byte = self.raw.data_ptr();
177        for _i in 0..index {
178            // SAFETY: Note this entire function has to be correct,
179            // not just this one call, for this to be correct!
180            at_byte = unsafe { self.one_hop_this_time(at_byte) };
181        }
182
183        // If this has gotten this far, it is known to be non-null,
184        // all the null values in the array up to this index were skipped,
185        // and the only offsets were via our hopping function.
186        Some(unsafe { self.bring_it_back_now(at_byte, false).expect("Null value in Strict array") })
187    }
188
189    /// Retrieve a value from a known-nullable (not strict) array.
190    fn get_nullable_inner<'arr>(
191        &'arr self,
192        nulls: &'arr BitSliceNulls,
193        index: usize,
194    ) -> Option<Option<T::As<'arr>>> {
195        // This assertion should only fail if null_slice is longer than the
196        // actual array,thanks to the check above.
197        debug_assert!(index < self.raw.len());
198
199        // This pointer is what's walked over the entire array's data buffer.
200        // If the array has varlena or cstr elements, we can't index into the array.
201        // If the elements are fixed size, we could, but we do not exploit that optimization yet
202        // as it would significantly complicate the code and impact debugging it.
203        // Such improvements should wait until a later version (today's: 0.7.4, preparing 0.8.0).
204        let mut at_byte = self.raw.data_ptr();
205        for i in 0..index {
206            match nulls.is_null(i) {
207                None => unreachable!("array was exceeded while walking to known non-null index???"),
208                // Skip nulls: the data buffer has no placeholders for them!
209                Some(true) => continue,
210                Some(false) => {
211                    // SAFETY: Note this entire function has to be correct,
212                    // not just this one call, for this to be correct!
213                    at_byte = unsafe { self.one_hop_this_time(at_byte) };
214                }
215            }
216        }
217
218        // If this has gotten this far, it is known to be non-null,
219        // all the null values in the array up to this index were skipped,
220        // and the only offsets were via our hopping function.
221        Some(unsafe { self.bring_it_back_now(at_byte, false) })
222    }
223
224    #[allow(clippy::option_option)]
225    #[inline]
226    pub fn get<'arr>(&'arr self, index: usize) -> Option<Option<T::As<'arr>>> {
227        if index >= self.len() {
228            return None;
229        }
230        match self.null_slice.get_inner().map(|v| (v, v.is_null(index))) {
231            // No null_slice, strict array.
232            None => self.get_strict_inner(index).map(Some),
233            // self.null_slice exists but index is not in bounds.
234            Some((_, None)) => None,
235            // Elem is null
236            Some((_, Some(true))) => Some(None),
237            // Elem is not null.
238            Some((nulls, Some(false))) => self.get_nullable_inner(nulls, index),
239        }
240    }
241
242    /// Extracts an element from a Postgres Array's data buffer
243    ///
244    /// # Safety
245    /// This assumes the pointer is to a valid element of that type.
246    #[inline]
247    unsafe fn bring_it_back_now<'arr>(
248        &'arr self,
249        ptr: *const u8,
250        is_null: bool,
251    ) -> Option<T::As<'arr>> {
252        match is_null {
253            true => None,
254            false => {
255                // Ensure we are not attempting to dereference a pointer
256                // outside of the array.
257                debug_assert!(self.is_within_bounds(ptr));
258                // Prevent a datum that begins inside the array but would end
259                // outside the array from being dereferenced.
260                debug_assert!(self.is_within_bounds_inclusive(
261                    ptr.wrapping_add(unsafe { self.slide_impl.hop_size(ptr) })
262                ));
263
264                unsafe { self.slide_impl.bring_it_back_now(self, ptr) }
265            }
266        }
267    }
268
269    /// Walk the data of a Postgres Array, "hopping" according to element layout.
270    ///
271    /// # Safety
272    /// For the varlena/cstring layout, data in the buffer is read.
273    /// In either case, pointer arithmetic is done, with the usual implications,
274    /// e.g. the pointer must be <= a "one past the end" pointer
275    /// This means this function must be invoked with the correct layout, and
276    /// either the array's `data_ptr` or a correctly offset pointer into it.
277    ///
278    /// Null elements will NOT be present in a Postgres Array's data buffer!
279    /// Do not cumulatively invoke this more than `len - null_count`!
280    /// Doing so will result in reading uninitialized data, which is UB!
281    #[inline]
282    unsafe fn one_hop_this_time(&self, ptr: *const u8) -> *const u8 {
283        unsafe {
284            let offset = self.slide_impl.hop_size(ptr);
285            // SAFETY: ptr stops at 1-past-end of the array's varlena
286            debug_assert!(ptr.wrapping_add(offset) <= self.raw.end_ptr());
287            ptr.add(offset)
288        }
289    }
290
291    /// Returns true if the pointer provided is within the bounds of the array.
292    /// Primarily intended for use with debug_assert!()s.
293    /// Note that this will return false for the 1-past-end, which is a useful
294    /// position for a pointer to be in, but not valid to dereference.
295    #[inline]
296    pub(crate) fn is_within_bounds(&self, ptr: *const u8) -> bool {
297        // Cast to usize, to prevent LLVM from doing counterintuitive things
298        // with pointer equality.
299        // See https://github.com/pgcentralfoundation/pgrx/pull/1514#discussion_r1480447846
300        let ptr: usize = ptr as usize;
301        let data_ptr = self.raw.data_ptr() as usize;
302        let end_ptr = self.raw.end_ptr() as usize;
303        (data_ptr <= ptr) && (ptr < end_ptr)
304    }
305    /// Similar to [`Self::is_within_bounds()`], but also returns true for the
306    /// 1-past-end position.
307    #[inline]
308    pub(crate) fn is_within_bounds_inclusive(&self, ptr: *const u8) -> bool {
309        // Cast to usize, to prevent LLVM from doing counterintuitive things
310        // with pointer equality.
311        // See https://github.com/pgcentralfoundation/pgrx/pull/1514#discussion_r1480447846
312        let ptr = ptr as usize;
313        let data_ptr = self.raw.data_ptr() as usize;
314        let end_ptr = self.raw.end_ptr() as usize;
315        (data_ptr <= ptr) && (ptr <= end_ptr)
316    }
317}
318
319#[deny(unsafe_op_in_unsafe_fn)]
320impl<T> Array<'_, T> {
321    /// Rips out the underlying `pg_sys::ArrayType` pointer.
322    /// Note that Array may have caused Postgres to allocate to unbox the datum,
323    /// and this can hypothetically cause a memory leak if so.
324    #[inline]
325    pub fn into_array_type(self) -> *const pg_sys::ArrayType {
326        // may be worth replacing this function when Toast<T> matures enough
327        // to be used as a public type with a fn(self) -> Toast<RawArray>
328
329        let Array { raw, .. } = self;
330        // Wrap the Toast<RawArray> to prevent it from deallocating itself
331        let mut raw = core::mem::ManuallyDrop::new(raw);
332        let ptr = raw.deref_mut().deref_mut() as *mut RawArray;
333        // SAFETY: Leaks are safe if they aren't use-after-frees!
334        unsafe { ptr.read() }.into_ptr().as_ptr() as _
335    }
336
337    /// Returns `true` if this [`Array`] contains one or more SQL "NULL" values
338    #[inline]
339    pub fn contains_nulls(&self) -> bool {
340        self.null_slice.get_inner().is_some_and(|slice| slice.contains_nulls())
341    }
342
343    #[inline]
344    pub fn len(&self) -> usize {
345        self.raw.len()
346    }
347
348    #[inline]
349    pub fn is_empty(&self) -> bool {
350        self.raw.len() == 0
351    }
352}
353
354/// Adapter to use `Nullable<T>` for array iteration.
355pub struct NullableArrayIterator<'mcx, T>
356where
357    T: UnboxDatum,
358{
359    inner: ArrayIterator<'mcx, T>,
360}
361
362impl<'mcx, T> Iterator for NullableArrayIterator<'mcx, T>
363where
364    T: UnboxDatum,
365{
366    type Item = Nullable<<T as unbox::UnboxDatum>::As<'mcx>>;
367
368    #[inline]
369    fn next(&mut self) -> Option<Self::Item> {
370        self.inner.next().map(|v| v.into())
371    }
372}
373
374impl<'mcx, T> IntoNullableIterator<<T as unbox::UnboxDatum>::As<'mcx>> for &'mcx Array<'mcx, T>
375where
376    T: UnboxDatum,
377{
378    type Iter = NullableArrayIterator<'mcx, T>;
379
380    fn into_nullable_iter(self) -> Self::Iter {
381        NullableArrayIterator { inner: self.iter() }
382    }
383}
384
385impl<'mcx, T: UnboxDatum> NullableContainer<'mcx, usize, <T as unbox::UnboxDatum>::As<'mcx>>
386    for Array<'mcx, T>
387{
388    type Layout = MaybeStrictNulls<BitSliceNulls<'mcx>>;
389
390    #[inline]
391    fn get_layout(&'mcx self) -> &'mcx Self::Layout {
392        &self.null_slice
393    }
394
395    #[inline]
396    unsafe fn get_raw(&'mcx self, idx: usize) -> <T as unbox::UnboxDatum>::As<'mcx> {
397        self.get_strict_inner(idx).expect(
398            "get_raw() called with an invalid index, bounds-checking\
399            *should* occur before calling this method.",
400        )
401    }
402
403    #[inline]
404    fn len(&'mcx self) -> usize {
405        self.len()
406    }
407}
408
409#[derive(thiserror::Error, Debug, Copy, Clone, Eq, PartialEq)]
410pub enum ArraySliceError {
411    #[error("Cannot create a slice of an Array that contains nulls")]
412    ContainsNulls,
413}
414
415#[cfg(target_pointer_width = "64")]
416impl Array<'_, f64> {
417    /// Returns a slice of `f64`s which comprise this [`Array`].
418    ///
419    /// # Errors
420    ///
421    /// Returns a [`ArraySliceError::ContainsNulls`] error if this [`Array`] contains one or more
422    /// SQL "NULL" values.  In this case, you'd likely want to fallback to using [`Array::iter()`].
423    #[inline]
424    pub fn as_slice(&self) -> Result<&[f64], ArraySliceError> {
425        as_slice(self)
426    }
427}
428
429impl Array<'_, f32> {
430    /// Returns a slice of `f32`s which comprise this [`Array`].
431    ///
432    /// # Errors
433    ///
434    /// Returns a [`ArraySliceError::ContainsNulls`] error if this [`Array`] contains one or more
435    /// SQL "NULL" values.  In this case, you'd likely want to fallback to using [`Array::iter()`].
436    #[inline]
437    pub fn as_slice(&self) -> Result<&[f32], ArraySliceError> {
438        as_slice(self)
439    }
440}
441
442#[cfg(target_pointer_width = "64")]
443impl Array<'_, i64> {
444    /// Returns a slice of `i64`s which comprise this [`Array`].
445    ///
446    /// # Errors
447    ///
448    /// Returns a [`ArraySliceError::ContainsNulls`] error if this [`Array`] contains one or more
449    /// SQL "NULL" values.  In this case, you'd likely want to fallback to using [`Array::iter()`].
450    #[inline]
451    pub fn as_slice(&self) -> Result<&[i64], ArraySliceError> {
452        as_slice(self)
453    }
454}
455
456impl Array<'_, i32> {
457    /// Returns a slice of `i32`s which comprise this [`Array`].
458    ///
459    /// # Errors
460    ///
461    /// Returns a [`ArraySliceError::ContainsNulls`] error if this [`Array`] contains one or more
462    /// SQL "NULL" values.  In this case, you'd likely want to fallback to using [`Array::iter()`].
463    #[inline]
464    pub fn as_slice(&self) -> Result<&[i32], ArraySliceError> {
465        as_slice(self)
466    }
467}
468
469impl Array<'_, i16> {
470    /// Returns a slice of `i16`s which comprise this [`Array`].
471    ///
472    /// # Errors
473    ///
474    /// Returns a [`ArraySliceError::ContainsNulls`] error if this [`Array`] contains one or more
475    /// SQL "NULL" values.  In this case, you'd likely want to fallback to using [`Array::iter()`].
476    #[inline]
477    pub fn as_slice(&self) -> Result<&[i16], ArraySliceError> {
478        as_slice(self)
479    }
480}
481
482impl Array<'_, i8> {
483    /// Returns a slice of `i8`s which comprise this [`Array`].
484    ///
485    /// # Errors
486    ///
487    /// Returns a [`ArraySliceError::ContainsNulls`] error if this [`Array`] contains one or more
488    /// SQL "NULL" values.  In this case, you'd likely want to fallback to using [`Array::iter()`].
489    #[inline]
490    pub fn as_slice(&self) -> Result<&[i8], ArraySliceError> {
491        as_slice(self)
492    }
493}
494
495#[inline(always)]
496fn as_slice<'a, T: Sized>(array: &'a Array<'_, T>) -> Result<&'a [T], ArraySliceError> {
497    if array.contains_nulls() {
498        return Err(ArraySliceError::ContainsNulls);
499    }
500
501    let slice =
502        unsafe { std::slice::from_raw_parts(array.raw.data_ptr() as *const _, array.len()) };
503    Ok(slice)
504}
505
506mod casper {
507    use super::UnboxDatum;
508    use crate::layout::Align;
509    use crate::{pg_sys, varlena, Array};
510
511    // it's a pop-culture reference (https://en.wikipedia.org/wiki/Cha_Cha_Slide) not some fancy crypto thing you nerd
512    /// Describes how to instantiate a value `T` from an [`Array`] and its backing byte array pointer.
513    /// It also knows how to determine the size of an [`Array`] element value.
514    pub(super) trait ChaChaSlide<T: UnboxDatum> {
515        /// Instantiate a `T` from the head of `ptr`
516        ///
517        /// # Safety
518        ///
519        /// This function is unsafe as it cannot guarantee that `ptr` points to the proper bytes
520        /// that represent a `T`, or even that it belongs to `array`.  Both of which must be true
521        unsafe fn bring_it_back_now<'arr, 'mcx>(
522            &self,
523            array: &'arr Array<'mcx, T>,
524            ptr: *const u8,
525        ) -> Option<T::As<'arr>>;
526
527        /// Determine how many bytes are used to represent `T`.  This could be fixed size or
528        /// even determined at runtime by whatever `ptr` is known to be pointing at.
529        ///
530        /// # Safety
531        ///
532        /// This function is unsafe as it cannot guarantee that `ptr` points to the bytes of a `T`,
533        /// which it must for implementations that rely on that.
534        unsafe fn hop_size(&self, ptr: *const u8) -> usize;
535    }
536
537    #[inline(always)]
538    fn is_aligned<T>(p: *const T) -> bool {
539        (p as usize) & (core::mem::align_of::<T>() - 1) == 0
540    }
541
542    /// Safety: Equivalent to a (potentially) aligned read of `ptr`, which
543    /// should be `Copy` (ideally...).
544    #[track_caller]
545    #[inline(always)]
546    pub(super) unsafe fn byval_read<T: Copy>(ptr: *const u8) -> T {
547        let ptr = ptr.cast::<T>();
548        debug_assert!(is_aligned(ptr), "not aligned to {}: {ptr:p}", std::mem::align_of::<T>());
549        ptr.read()
550    }
551
552    /// Fixed-size byval array elements. N should be 1, 2, 4, or 8. Note that
553    /// `T` (the rust type) may have a different size than `N`.
554    pub(super) struct FixedSizeByVal<const N: usize>;
555    impl<T: UnboxDatum, const N: usize> ChaChaSlide<T> for FixedSizeByVal<N> {
556        #[inline(always)]
557        unsafe fn bring_it_back_now<'arr, 'mcx>(
558            &self,
559            _array: &'arr Array<'mcx, T>,
560            ptr: *const u8,
561        ) -> Option<T::As<'arr>> {
562            // This branch is optimized away (because `N` is constant).
563            let datum = match N {
564                // for match with `Datum`, read through that directly to
565                // preserve provenance (may not be relevant but doesn't hurt).
566                1 => pg_sys::Datum::from(byval_read::<u8>(ptr)),
567                2 => pg_sys::Datum::from(byval_read::<u16>(ptr)),
568                4 => pg_sys::Datum::from(byval_read::<u32>(ptr)),
569                8 => pg_sys::Datum::from(byval_read::<u64>(ptr)),
570                _ => unreachable!("`N` must be 1, 2, 4, or 8 (got {N})"),
571            };
572            Some(T::unbox(core::mem::transmute(datum)))
573        }
574
575        #[inline(always)]
576        unsafe fn hop_size(&self, _ptr: *const u8) -> usize {
577            N
578        }
579    }
580
581    /// Array elements are [`pg_sys::varlena`] types, which are pass-by-reference
582    pub(super) struct PassByVarlena {
583        pub(super) align: Align,
584    }
585    impl<T: UnboxDatum> ChaChaSlide<T> for PassByVarlena {
586        #[inline]
587        unsafe fn bring_it_back_now<'arr, 'mcx>(
588            &self,
589            // May need this array param for MemCx reasons?
590            _array: &'arr Array<'mcx, T>,
591            ptr: *const u8,
592        ) -> Option<T::As<'arr>> {
593            let datum = pg_sys::Datum::from(ptr);
594            Some(T::unbox(core::mem::transmute(datum)))
595        }
596
597        #[inline]
598        unsafe fn hop_size(&self, ptr: *const u8) -> usize {
599            // SAFETY: This uses the varsize_any function to be safe,
600            // and the caller was informed of pointer requirements.
601            let varsize = varlena::varsize_any(ptr.cast());
602
603            // Now make sure this is aligned-up
604            self.align.pad(varsize)
605        }
606    }
607
608    /// Array elements are standard C strings (`char *`), which are pass-by-reference
609    pub(super) struct PassByCStr;
610    impl<T: UnboxDatum> ChaChaSlide<T> for PassByCStr {
611        #[inline]
612        unsafe fn bring_it_back_now<'arr, 'mcx>(
613            &self,
614            _array: &'arr Array<'mcx, T>,
615            ptr: *const u8,
616        ) -> Option<T::As<'arr>> {
617            let datum = pg_sys::Datum::from(ptr);
618            Some(T::unbox(core::mem::transmute(datum)))
619        }
620
621        #[inline]
622        unsafe fn hop_size(&self, ptr: *const u8) -> usize {
623            // SAFETY: The caller was informed of pointer requirements.
624            let strlen = core::ffi::CStr::from_ptr(ptr.cast()).to_bytes().len();
625
626            // Skip over the null which points us to the head of the next cstr
627            strlen + 1
628        }
629    }
630
631    pub(super) struct PassByFixed {
632        pub(super) padded_size: usize,
633    }
634
635    impl<T: UnboxDatum> ChaChaSlide<T> for PassByFixed {
636        #[inline]
637        unsafe fn bring_it_back_now<'arr, 'mcx>(
638            &self,
639            _array: &'arr Array<'mcx, T>,
640            ptr: *const u8,
641        ) -> Option<T::As<'arr>> {
642            let datum = pg_sys::Datum::from(ptr);
643            Some(T::unbox(core::mem::transmute(datum)))
644        }
645
646        #[inline]
647        unsafe fn hop_size(&self, _ptr: *const u8) -> usize {
648            self.padded_size
649        }
650    }
651}
652
653pub struct VariadicArray<'mcx, T>(Array<'mcx, T>);
654
655impl<'mcx, T: UnboxDatum> Serialize for VariadicArray<'mcx, T>
656where
657    for<'arr> <T as UnboxDatum>::As<'arr>: Serialize,
658{
659    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
660    where
661        S: Serializer,
662    {
663        serializer.collect_seq(self.0.iter())
664    }
665}
666
667impl<'mcx, T: UnboxDatum> VariadicArray<'mcx, T> {
668    pub(crate) fn wrap_array(arr: Array<'mcx, T>) -> Self {
669        VariadicArray(arr)
670    }
671
672    /// Return an Iterator of `Option<T>` over the contained Datums.
673    #[inline]
674    pub fn iter(&self) -> ArrayIterator<'_, T> {
675        self.0.iter()
676    }
677
678    /// Return an iterator over the Array's elements.
679    ///
680    /// # Panics
681    /// This function will panic when called if the array contains any SQL NULL values.
682    #[inline]
683    pub fn iter_deny_null(&self) -> ArrayTypedIterator<'_, T> {
684        self.0.iter_deny_null()
685    }
686
687    #[allow(clippy::option_option)]
688    #[inline]
689    pub fn get<'arr>(&'arr self, i: usize) -> Option<Option<<T as UnboxDatum>::As<'arr>>> {
690        self.0.get(i)
691    }
692}
693
694impl<T> VariadicArray<'_, T> {
695    #[inline]
696    pub fn into_array_type(self) -> *const pg_sys::ArrayType {
697        self.0.into_array_type()
698    }
699
700    #[inline]
701    pub fn len(&self) -> usize {
702        self.0.len()
703    }
704
705    #[inline]
706    pub fn is_empty(&self) -> bool {
707        self.0.is_empty()
708    }
709}
710
711pub struct ArrayTypedIterator<'arr, T> {
712    array: &'arr Array<'arr, T>,
713    curr: usize,
714    ptr: *const u8,
715}
716
717impl<'arr, T: UnboxDatum> Iterator for ArrayTypedIterator<'arr, T> {
718    type Item = T::As<'arr>;
719
720    #[inline]
721    fn next(&mut self) -> Option<Self::Item> {
722        let Self { array, curr, ptr } = self;
723        if *curr >= array.raw.len() {
724            None
725        } else {
726            // SAFETY: The constructor for this type instantly panics if any nulls are present!
727            // Thus as an invariant, this will never have to reckon with the nullbitmap.
728            let element = unsafe { array.bring_it_back_now(*ptr, false) };
729            *curr += 1;
730            *ptr = unsafe { array.one_hop_this_time(*ptr) };
731            element
732        }
733    }
734
735    #[inline]
736    fn size_hint(&self) -> (usize, Option<usize>) {
737        let len = self.array.raw.len().saturating_sub(self.curr);
738        (len, Some(len))
739    }
740}
741
742impl<'a, T: UnboxDatum> ExactSizeIterator for ArrayTypedIterator<'a, T> {}
743impl<'a, T: UnboxDatum> core::iter::FusedIterator for ArrayTypedIterator<'a, T> {}
744
745impl<'arr, T: UnboxDatum + serde::Serialize> serde::Serialize for ArrayTypedIterator<'arr, T>
746where
747    <T as UnboxDatum>::As<'arr>: serde::Serialize,
748{
749    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
750    where
751        S: Serializer,
752    {
753        serializer.collect_seq(self.array.iter())
754    }
755}
756
757pub struct ArrayIterator<'arr, T> {
758    array: &'arr Array<'arr, T>,
759    curr: usize,
760    ptr: *const u8,
761}
762
763impl<'arr, T: UnboxDatum> Iterator for ArrayIterator<'arr, T> {
764    type Item = Option<T::As<'arr>>;
765
766    #[inline]
767    fn next(&mut self) -> Option<Self::Item> {
768        let Self { array, curr, ptr } = self;
769
770        if *curr >= array.len() {
771            return None;
772        }
773        let is_null = (match array.null_slice.get_inner().map(|slice| slice.is_null(*curr)) {
774            // Null slice exists, use its logic for bounds-checking
775            // and null status.
776            Some(elem) => elem,
777            // Strict array, no nulls.
778            // Bounds-checking behavior is still needed though.
779            None => (*curr < array.len()).then_some(false),
780        })?;
781        *curr += 1;
782
783        let element = unsafe { array.bring_it_back_now(*ptr, is_null) };
784        if !is_null {
785            // SAFETY: This has to not move for nulls, as they occupy 0 data bytes,
786            // and it has to move only after unpacking a non-null varlena element,
787            // as the iterator starts by pointing to the first non-null element!
788            *ptr = unsafe { array.one_hop_this_time(*ptr) };
789        }
790        Some(element)
791    }
792
793    #[inline]
794    fn size_hint(&self) -> (usize, Option<usize>) {
795        let len = self.array.raw.len().saturating_sub(self.curr);
796        (len, Some(len))
797    }
798}
799
800impl<'arr, T: UnboxDatum> ExactSizeIterator for ArrayIterator<'arr, T> {}
801impl<'arr, T: UnboxDatum> FusedIterator for ArrayIterator<'arr, T> {}
802
803pub struct ArrayIntoIterator<'a, T> {
804    array: Array<'a, T>,
805    curr: usize,
806    ptr: *const u8,
807}
808
809// There's nowhere to name the lifetime contraction
810impl<'mcx, T> IntoIterator for Array<'mcx, T>
811where
812    for<'arr> T: UnboxDatum<As<'arr> = T> + 'static,
813{
814    type Item = Option<T::As<'mcx>>;
815    type IntoIter = ArrayIntoIterator<'mcx, T>;
816
817    #[inline]
818    fn into_iter(self) -> Self::IntoIter {
819        let ptr = self.raw.data_ptr();
820        ArrayIntoIterator { array: self, curr: 0, ptr }
821    }
822}
823
824impl<'mcx, T> IntoIterator for VariadicArray<'mcx, T>
825where
826    for<'arr> T: UnboxDatum<As<'arr> = T> + 'static,
827{
828    type Item = Option<T::As<'mcx>>;
829    type IntoIter = ArrayIntoIterator<'mcx, T>;
830
831    #[inline]
832    fn into_iter(self) -> Self::IntoIter {
833        let ptr = self.0.raw.data_ptr();
834        ArrayIntoIterator { array: self.0, curr: 0, ptr }
835    }
836}
837
838impl<'mcx, T> Iterator for ArrayIntoIterator<'mcx, T>
839where
840    for<'arr> T: UnboxDatum<As<'arr> = T> + 'static,
841{
842    type Item = Option<T::As<'static>>;
843
844    #[inline]
845    fn next(&mut self) -> Option<Self::Item> {
846        let Self { array, curr, ptr } = self;
847
848        if *curr >= array.len() {
849            return None;
850        }
851
852        let is_null = (match array.null_slice.get_inner().map(|slice| slice.is_null(*curr)) {
853            // Null slice exists, use its logic for bounds-checking
854            // and null status.
855            Some(elem) => elem,
856            // Strict array, no nulls.
857            // Bounds-checking behavior is still needed though.
858            None => (*curr < array.len()).then_some(false),
859        })?;
860
861        *curr += 1;
862        debug_assert!(array.is_within_bounds(*ptr));
863        let element = unsafe { array.bring_it_back_now(*ptr, is_null) };
864        if !is_null {
865            // SAFETY: This has to not move for nulls, as they occupy 0 data bytes,
866            // and it has to move only after unpacking a non-null varlena element,
867            // as the iterator starts by pointing to the first non-null element!
868            *ptr = unsafe { array.one_hop_this_time(*ptr) };
869        }
870        Some(element)
871    }
872
873    #[inline]
874    fn size_hint(&self) -> (usize, Option<usize>) {
875        let len = self.array.raw.len().saturating_sub(self.curr);
876        (len, Some(len))
877    }
878}
879
880impl<'mcx, T> ExactSizeIterator for ArrayIntoIterator<'mcx, T> where
881    for<'arr> T: UnboxDatum<As<'arr> = T> + 'static
882{
883}
884impl<'mcx, T: UnboxDatum> FusedIterator for ArrayIntoIterator<'mcx, T> where
885    for<'arr> T: UnboxDatum<As<'arr> = T> + 'static
886{
887}
888
889impl<'a, T: FromDatum + UnboxDatum> FromDatum for VariadicArray<'a, T> {
890    #[inline]
891    unsafe fn from_polymorphic_datum(
892        datum: pg_sys::Datum,
893        is_null: bool,
894        oid: pg_sys::Oid,
895    ) -> Option<VariadicArray<'a, T>> {
896        Array::from_polymorphic_datum(datum, is_null, oid).map(Self)
897    }
898}
899
900impl<'a, T: UnboxDatum> FromDatum for Array<'a, T> {
901    #[inline]
902    unsafe fn from_polymorphic_datum(
903        datum: pg_sys::Datum,
904        is_null: bool,
905        _typoid: pg_sys::Oid,
906    ) -> Option<Array<'a, T>> {
907        if is_null {
908            None
909        } else {
910            let Some(ptr) = NonNull::new(datum.cast_mut_ptr()) else { return None };
911            let raw = RawArray::detoast_from_varlena(ptr);
912            Some(Array::deconstruct_from(raw))
913        }
914    }
915
916    unsafe fn from_datum_in_memory_context(
917        mut memory_context: PgMemoryContexts,
918        datum: pg_sys::Datum,
919        is_null: bool,
920        typoid: pg_sys::Oid,
921    ) -> Option<Self>
922    where
923        Self: Sized,
924    {
925        if is_null {
926            None
927        } else {
928            memory_context.switch_to(|_| {
929                // copy the Datum into this MemoryContext, and then instantiate the Array wrapper
930                let copy = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr());
931                Array::<T>::from_polymorphic_datum(pg_sys::Datum::from(copy), false, typoid)
932            })
933        }
934    }
935}
936
937impl<T: IntoDatum> IntoDatum for Array<'_, T> {
938    #[inline]
939    fn into_datum(self) -> Option<pg_sys::Datum> {
940        let array_type = self.into_array_type();
941        let datum = pg_sys::Datum::from(array_type);
942        Some(datum)
943    }
944
945    #[inline]
946    fn type_oid() -> pg_sys::Oid {
947        unsafe { pg_sys::get_array_type(T::type_oid()) }
948    }
949
950    fn composite_type_oid(&self) -> Option<pg_sys::Oid> {
951        Some(unsafe { pg_sys::get_array_type(self.raw.oid()) })
952    }
953}
954
955impl<T> FromDatum for Vec<T>
956where
957    for<'arr> T: UnboxDatum<As<'arr> = T> + 'arr,
958{
959    #[inline]
960    unsafe fn from_polymorphic_datum(
961        datum: pg_sys::Datum,
962        is_null: bool,
963        typoid: pg_sys::Oid,
964    ) -> Option<Vec<T>> {
965        if is_null {
966            None
967        } else {
968            Array::<T>::from_polymorphic_datum(datum, is_null, typoid)
969                .map(|array| array.iter_deny_null().collect::<Vec<_>>())
970        }
971    }
972
973    unsafe fn from_datum_in_memory_context(
974        memory_context: PgMemoryContexts,
975        datum: pg_sys::Datum,
976        is_null: bool,
977        typoid: pg_sys::Oid,
978    ) -> Option<Self>
979    where
980        Self: Sized,
981    {
982        Array::<T>::from_datum_in_memory_context(memory_context, datum, is_null, typoid)
983            .map(|array| array.iter_deny_null().collect::<Vec<_>>())
984    }
985}
986
987impl<T> FromDatum for Vec<Option<T>>
988where
989    for<'arr> T: UnboxDatum<As<'arr> = T> + 'arr,
990{
991    #[inline]
992    unsafe fn from_polymorphic_datum(
993        datum: pg_sys::Datum,
994        is_null: bool,
995        typoid: pg_sys::Oid,
996    ) -> Option<Vec<Option<T>>> {
997        Array::<T>::from_polymorphic_datum(datum, is_null, typoid)
998            .map(|array| array.iter().collect::<Vec<_>>())
999    }
1000
1001    unsafe fn from_datum_in_memory_context(
1002        memory_context: PgMemoryContexts,
1003        datum: pg_sys::Datum,
1004        is_null: bool,
1005        typoid: pg_sys::Oid,
1006    ) -> Option<Self>
1007    where
1008        Self: Sized,
1009    {
1010        Array::<T>::from_datum_in_memory_context(memory_context, datum, is_null, typoid)
1011            .map(|array| array.iter().collect::<Vec<_>>())
1012    }
1013}
1014
1015#[inline]
1016/// Converts an iterator into an array datum
1017fn array_datum_from_iter<T: IntoDatum>(elements: impl Iterator<Item = T>) -> Option<pg_sys::Datum> {
1018    let mut state = unsafe {
1019        pg_sys::initArrayResult(
1020            T::type_oid(),
1021            PgMemoryContexts::CurrentMemoryContext.value(),
1022            // All elements use the same memory context
1023            false,
1024        )
1025    };
1026    for s in elements {
1027        let datum = s.into_datum();
1028        let isnull = datum.is_none();
1029
1030        unsafe {
1031            state = pg_sys::accumArrayResult(
1032                state,
1033                datum.unwrap_or(0.into()),
1034                isnull,
1035                T::type_oid(),
1036                PgMemoryContexts::CurrentMemoryContext.value(),
1037            );
1038        }
1039    }
1040
1041    // Should not happen: {init, accum}ArrayResult both return non-null pointers
1042    assert!(!state.is_null());
1043
1044    Some(unsafe { pg_sys::makeArrayResult(state, PgMemoryContexts::CurrentMemoryContext.value()) })
1045}
1046
1047impl<T> IntoDatum for Vec<T>
1048where
1049    T: IntoDatum,
1050{
1051    fn into_datum(self) -> Option<pg_sys::Datum> {
1052        array_datum_from_iter(self.into_iter())
1053    }
1054
1055    fn type_oid() -> pg_sys::Oid {
1056        unsafe { pg_sys::get_array_type(T::type_oid()) }
1057    }
1058
1059    #[allow(clippy::get_first)] // https://github.com/pgcentralfoundation/pgrx/issues/1363
1060    fn composite_type_oid(&self) -> Option<pg_sys::Oid> {
1061        // the composite type oid for a vec of composite types is the array type of the base composite type
1062        // the use of first() would have presented a false certainty here: it's not actually relevant that it be the first.
1063        #[allow(clippy::get_first)]
1064        self.get(0)
1065            .and_then(|v| v.composite_type_oid().map(|oid| unsafe { pg_sys::get_array_type(oid) }))
1066    }
1067
1068    #[inline]
1069    fn is_compatible_with(other: pg_sys::Oid) -> bool {
1070        Self::type_oid() == other || other == unsafe { pg_sys::get_array_type(T::type_oid()) }
1071    }
1072}
1073
1074impl<'a, T> IntoDatum for &'a [T]
1075where
1076    T: IntoDatum + Copy + 'a,
1077{
1078    fn into_datum(self) -> Option<pg_sys::Datum> {
1079        array_datum_from_iter(self.into_iter().copied())
1080    }
1081
1082    fn type_oid() -> pg_sys::Oid {
1083        unsafe { pg_sys::get_array_type(T::type_oid()) }
1084    }
1085
1086    #[inline]
1087    fn is_compatible_with(other: pg_sys::Oid) -> bool {
1088        Self::type_oid() == other || other == unsafe { pg_sys::get_array_type(T::type_oid()) }
1089    }
1090}
1091
1092unsafe impl<T> SqlTranslatable for Array<'_, T>
1093where
1094    T: SqlTranslatable,
1095{
1096    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
1097        match T::argument_sql()? {
1098            SqlMapping::As(sql) => Ok(SqlMapping::As(format!("{sql}[]"))),
1099            SqlMapping::Skip => Err(ArgumentError::SkipInArray),
1100            SqlMapping::Composite { .. } => Ok(SqlMapping::Composite { array_brackets: true }),
1101        }
1102    }
1103
1104    fn return_sql() -> Result<Returns, ReturnsError> {
1105        match T::return_sql()? {
1106            Returns::One(SqlMapping::As(sql)) => {
1107                Ok(Returns::One(SqlMapping::As(format!("{sql}[]"))))
1108            }
1109            Returns::One(SqlMapping::Composite { array_brackets: _ }) => {
1110                Ok(Returns::One(SqlMapping::Composite { array_brackets: true }))
1111            }
1112            Returns::One(SqlMapping::Skip) => Err(ReturnsError::SkipInArray),
1113            Returns::SetOf(_) => Err(ReturnsError::SetOfInArray),
1114            Returns::Table(_) => Err(ReturnsError::TableInArray),
1115        }
1116    }
1117}
1118
1119unsafe impl<T> SqlTranslatable for VariadicArray<'_, T>
1120where
1121    T: SqlTranslatable,
1122{
1123    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
1124        match T::argument_sql()? {
1125            SqlMapping::As(sql) => Ok(SqlMapping::As(format!("{sql}[]"))),
1126            SqlMapping::Skip => Err(ArgumentError::SkipInArray),
1127            SqlMapping::Composite { .. } => Ok(SqlMapping::Composite { array_brackets: true }),
1128        }
1129    }
1130
1131    fn return_sql() -> Result<Returns, ReturnsError> {
1132        match T::return_sql()? {
1133            Returns::One(SqlMapping::As(sql)) => {
1134                Ok(Returns::One(SqlMapping::As(format!("{sql}[]"))))
1135            }
1136            Returns::One(SqlMapping::Composite { array_brackets: _ }) => {
1137                Ok(Returns::One(SqlMapping::Composite { array_brackets: true }))
1138            }
1139            Returns::One(SqlMapping::Skip) => Err(ReturnsError::SkipInArray),
1140            Returns::SetOf(_) => Err(ReturnsError::SetOfInArray),
1141            Returns::Table(_) => Err(ReturnsError::TableInArray),
1142        }
1143    }
1144
1145    fn variadic() -> bool {
1146        true
1147    }
1148}
pgrx/datum/array.rs

pgrx/datum/
array.rs
