Skip to content

Commit

Permalink
Macro based IntoDatum for string types (#1886)
Browse files Browse the repository at this point in the history
The pull request follows two goals:
* don't repeat yourself,
* implement `IntoDatum` for `&CString` to mirror implementation of the
trait for `&String`.
  • Loading branch information
YohDeadfall authored Sep 28, 2024
1 parent 638d18e commit 0523797
Showing 1 changed file with 59 additions and 85 deletions.
144 changes: 59 additions & 85 deletions pgrx/src/datum/into.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,12 @@
use crate::{pg_sys, rust_regtypein, set_varsize_4b, PgBox, PgOid, WhoAllocated};
use core::fmt::Display;
use pgrx_pg_sys::panic::ErrorReportable;
use std::{any::Any, ptr::addr_of_mut, str};
use std::{
any::Any,
ffi::{CStr, CString},
ptr::addr_of_mut,
str,
};

/// Convert a Rust type into a `pg_sys::Datum`.
///
Expand Down Expand Up @@ -245,54 +250,31 @@ impl IntoDatum for PgOid {
}
}

/// for text, varchar
impl<'a> IntoDatum for &'a str {
#[inline]
fn into_datum(self) -> Option<pg_sys::Datum> {
self.as_bytes().into_datum()
}

fn type_oid() -> pg_sys::Oid {
pg_sys::TEXTOID
}

#[inline]
fn is_compatible_with(other: pg_sys::Oid) -> bool {
Self::type_oid() == other || other == pg_sys::VARCHAROID
}
// for text, varchar
macro_rules! impl_into_datum_str {
($t:ty) => {
impl IntoDatum for $t {
#[inline]
fn into_datum(self) -> Option<$crate::pg_sys::Datum> {
self.as_bytes().into_datum()
}

#[inline]
fn type_oid() -> pg_sys::Oid {
pg_sys::TEXTOID
}

#[inline]
fn is_compatible_with(other: pg_sys::Oid) -> bool {
Self::type_oid() == other || other == pg_sys::VARCHAROID
}
}
};
}

impl IntoDatum for String {
#[inline]
fn into_datum(self) -> Option<pg_sys::Datum> {
self.as_str().into_datum()
}

fn type_oid() -> pg_sys::Oid {
pg_sys::TEXTOID
}

#[inline]
fn is_compatible_with(other: pg_sys::Oid) -> bool {
Self::type_oid() == other || other == pg_sys::VARCHAROID
}
}

impl IntoDatum for &String {
#[inline]
fn into_datum(self) -> Option<pg_sys::Datum> {
self.as_str().into_datum()
}

fn type_oid() -> pg_sys::Oid {
pg_sys::TEXTOID
}

#[inline]
fn is_compatible_with(other: pg_sys::Oid) -> bool {
Self::type_oid() == other || other == pg_sys::VARCHAROID
}
}
impl_into_datum_str!(String);
impl_into_datum_str!(&String);
impl_into_datum_str!(&str);

impl IntoDatum for char {
#[inline]
Expand All @@ -317,47 +299,39 @@ impl IntoDatum for char {
}
}

/// for cstring
impl<'a> IntoDatum for &'a core::ffi::CStr {
/// The [`core::ffi::CStr`] is copied to `palloc`'d memory. That memory will either be freed by
/// Postgres when [`pg_sys::CurrentMemoryContext`] is reset, or when the function you passed the
/// returned Datum to decides to free it.
#[inline]
fn into_datum(self) -> Option<pg_sys::Datum> {
unsafe {
// SAFETY: A `CStr` has already been validated to be a non-null pointer to a null-terminated
// "char *", and it won't ever overlap with a newly palloc'd block of memory. Using
// `to_bytes_with_nul()` ensures that we'll never try to palloc zero bytes -- it'll at
// least always be 1 byte to hold the null terminator for the empty string.
//
// This is akin to Postgres' `pg_sys::pstrdup` or even `pg_sys::pnstrdup` functions, but
// doing the copy ourselves allows us to elide the "strlen" or "strnlen" operations those
// functions need to do; the byte slice returned from `to_bytes_with_nul` knows its length.
let bytes = self.to_bytes_with_nul();
let copy = pg_sys::palloc(bytes.len()).cast();
core::ptr::copy_nonoverlapping(bytes.as_ptr(), copy, bytes.len());
Some(copy.into())
// for cstring
macro_rules! impl_into_datum_c_str {
($t:ty) => {
impl IntoDatum for $t {
#[inline]
fn into_datum(self) -> Option<$crate::pg_sys::Datum> {
unsafe {
// SAFETY: A `CStr` has already been validated to be a non-null pointer to a null-terminated
// "char *", and it won't ever overlap with a newly palloc'd block of memory. Using
// `to_bytes_with_nul()` ensures that we'll never try to palloc zero bytes -- it'll at
// least always be 1 byte to hold the null terminator for the empty string.
//
// This is akin to Postgres' `pg_sys::pstrdup` or even `pg_sys::pnstrdup` functions, but
// doing the copy ourselves allows us to elide the "strlen" or "strnlen" operations those
// functions need to do; the byte slice returned from `to_bytes_with_nul` knows its length.
let src = self.as_ref().to_bytes_with_nul();
let dst = $crate::pg_sys::palloc(src.len());
dst.copy_from(src.as_ptr().cast(), src.len());
Some(dst.into())
}
}

#[inline]
fn type_oid() -> pg_sys::Oid {
pg_sys::CSTRINGOID
}
}
}

fn type_oid() -> pg_sys::Oid {
pg_sys::CSTRINGOID
}
};
}

impl IntoDatum for alloc::ffi::CString {
/// The [`alloc::ffi::CString`] is copied to `palloc`'d memory. That memory will either be freed by
/// Postgres when [`pg_sys::CurrentMemoryContext`] is reset, or when the function you passed the
/// returned Datum to decides to free it.
#[inline]
fn into_datum(self) -> Option<pg_sys::Datum> {
self.as_c_str().into_datum()
}

fn type_oid() -> pg_sys::Oid {
pg_sys::CSTRINGOID
}
}
impl_into_datum_c_str!(CString);
impl_into_datum_c_str!(&CString);
impl_into_datum_c_str!(&CStr);

/// for bytea
impl<'a> IntoDatum for &'a [u8] {
Expand Down

0 comments on commit 0523797

Please sign in to comment.