feat: auto-generate type alignment (#[pgrx(alignment = "on")])

pgrx-examples/custom_types/src/alignment.rs

@@ -0,0 +1,54 @@
+//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
+//LICENSE
+//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
+//LICENSE
+//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <[email protected]>
+//LICENSE
+//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::prelude::*;
+use serde::*;
+
+#[derive(PostgresType, Serialize, Deserialize)]
+#[pgrx(alignment = "on")]
+pub struct AlignedTo4Bytes {
+    v1: u32,
+    v2: [u32; 3],
+}
+
+#[derive(PostgresType, Serialize, Deserialize)]
+#[pgrx(alignment = "on")]
+pub struct AlignedTo8Bytes {
+    v1: u64,
+    v2: [u64; 3],
+}
+
+#[derive(PostgresType, Serialize, Deserialize)]
+#[pgrx(alignment = "off")]
+pub struct NotAlignedTo8Bytes {
+    v1: u64,
+    v2: [u64; 3],
+}
+
+#[cfg(any(test, feature = "pg_test"))]
+#[pg_schema]
+mod tests {
+    use pgrx::prelude::*;
+
+    #[cfg(not(feature = "no-schema-generation"))]
+    #[pg_test]
+    fn test_alignment_is_correct() {
+        let val = Spi::get_one::<String>(r#"SELECT typalign::text FROM pg_type WHERE typname = 'alignedto4bytes'"#).unwrap().unwrap();
+
+        assert!(val == "i");
+
+        let val = Spi::get_one::<String>(r#"SELECT typalign::text FROM pg_type WHERE typname = 'alignedto8bytes'"#).unwrap().unwrap();
+
+        assert!(val == "d");
+
+        let val = Spi::get_one::<String>(r#"SELECT typalign::text FROM pg_type WHERE typname = 'notalignedto8bytes'"#).unwrap().unwrap();
+
+        assert!(val == "i");
+    }
+}

pgrx-examples/custom_types/src/lib.rs

@@ -7,6 +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.
+mod alignment;
 mod complex;
 mod fixed_size;
 mod generic_enum;

pgrx-macros/src/lib.rs

@@ -781,6 +781,7 @@ Optionally accepts the following attributes:
 
 * `inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the type.
 * `pgvarlena_inoutfuncs(some_in_fn, some_out_fn)`: Define custom in/out functions for the `PgVarlena` of this type.
+* `pgrx(alignment = "<align>")`: Derive Postgres alignment from Rust type. One of `"on"`, or `"off"`.
 * `sql`: Same arguments as [`#[pgrx(sql = ..)]`](macro@pgrx).
 */
 #[proc_macro_derive(

pgrx-sql-entity-graph/src/lib.rs

@@ -94,6 +94,8 @@ pub trait SqlGraphIdentifier {
     fn line(&self) -> Option<u32>;
 }
 
+pub use postgres_type::Alignment;
+
 /// An entity corresponding to some SQL required by the extension.
 #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub enum SqlGraphEntity {

pgrx-sql-entity-graph/src/postgres_type/entity.rs

@@ -16,13 +16,82 @@
 
 */
 use crate::mapping::RustSqlMapping;
+use crate::pgrx_attribute::{ArgValue, PgrxArg, PgrxAttribute};
 use crate::pgrx_sql::PgrxSql;
 use crate::to_sql::entity::ToSqlConfigEntity;
 use crate::to_sql::ToSql;
 use crate::{SqlGraphEntity, SqlGraphIdentifier, TypeMatch};
+use eyre::eyre;
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote, ToTokens, TokenStreamExt};
 use std::collections::BTreeSet;
+use syn::spanned::Spanned;
+use syn::{AttrStyle, Attribute, Lit};
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
+pub enum Alignment {
+    On,
+    Off,
+}
+
+const INVALID_ATTR_CONTENT: &str =
+    r#"expected `#[pgrx(alignment = align)]`, where `align` is "on", or "off""#;
+
+impl ToTokens for Alignment {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        let value = match self {
+            Alignment::On => format_ident!("On"),
+            Alignment::Off => format_ident!("Off"),
+        };
+        let quoted = quote! {
+            ::pgrx::pgrx_sql_entity_graph::Alignment::#value
+        };
+        tokens.append_all(quoted);
+    }
+}
+
+impl Alignment {
+    pub fn from_attribute(attr: &Attribute) -> Result<Option<Self>, syn::Error> {
+        if attr.style != AttrStyle::Outer {
+            return Err(syn::Error::new(
+                attr.span(),
+                "#[pgrx(alignment = ..)] is only valid in an outer context",
+            ));
+        }
+
+        let attr = attr.parse_args::<PgrxAttribute>()?;
+        for arg in attr.args.iter() {
+            let PgrxArg::NameValue(ref nv) = arg;
+            if !nv.path.is_ident("alignment") {
+                continue;
+            }
+
+            return match nv.value {
+                ArgValue::Lit(Lit::Str(ref s)) => match s.value().as_ref() {
+                    "on" => Ok(Some(Self::On)),
+                    "off" => Ok(Some(Self::Off)),
+                    _ => Err(syn::Error::new(s.span(), INVALID_ATTR_CONTENT)),
+                },
+                ArgValue::Path(ref p) => Err(syn::Error::new(p.span(), INVALID_ATTR_CONTENT)),
+                ArgValue::Lit(ref l) => Err(syn::Error::new(l.span(), INVALID_ATTR_CONTENT)),
+            };
+        }
+
+        Ok(None)
+    }
+
+    pub fn from_attributes(attrs: &[Attribute]) -> Result<Self, syn::Error> {
+        for attr in attrs {
+            if attr.path().is_ident("pgrx") {
+                if let Some(v) = Self::from_attribute(attr)? {
+                    return Ok(v);
+                }
+            }
+        }
+        Ok(Self::Off)
+    }
+}
 
-use eyre::eyre;
 /// The output of a [`PostgresType`](crate::postgres_type::PostgresTypeDerive) from `quote::ToTokens::to_tokens`.
 #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub struct PostgresTypeEntity {
@@ -37,6 +106,7 @@ pub struct PostgresTypeEntity {
     pub out_fn: &'static str,
     pub out_fn_module_path: String,
     pub to_sql_config: ToSqlConfigEntity,
+    pub alignment: Option<usize>,
 }
 
 impl TypeMatch for PostgresTypeEntity {
@@ -82,6 +152,7 @@ impl ToSql for PostgresTypeEntity {
             out_fn,
             out_fn_module_path,
             in_fn,
+            alignment,
             ..
         }) = item_node
         else {
@@ -155,6 +226,24 @@ impl ToSql for PostgresTypeEntity {
             schema = context.schema_prefix_for(&self_index),
         );
 
+        let alignment = alignment
+            .map(|alignment| {
+                assert!(alignment.is_power_of_two());
+                let alignment = match alignment {
+                    1 => "char",
+                    2 => "int2",
+                    4 => "int4",
+                    8 => "double",
+                    _ => panic!("type '{name}' wants unsupported alignment '{alignment}'"),
+                };
+                format!(
+                    ",\n\
+                    \tALIGNMENT = {}",
+                    alignment
+                )
+            })
+            .unwrap_or_default();
+
         let materialized_type = format! {
             "\n\
                 -- {file}:{line}\n\
@@ -163,7 +252,7 @@ impl ToSql for PostgresTypeEntity {
                     \tINTERNALLENGTH = variable,\n\
                     \tINPUT = {schema_prefix_in_fn}{in_fn}, /* {in_fn_path} */\n\
                     \tOUTPUT = {schema_prefix_out_fn}{out_fn}, /* {out_fn_path} */\n\
-                    \tSTORAGE = extended\n\
+                    \tSTORAGE = extended{alignment}\n\
                 );\
             ",
             schema = context.schema_prefix_for(&self_index),

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

@@ -23,6 +23,7 @@ use quote::{format_ident, quote};
 use syn::parse::{Parse, ParseStream};
 use syn::{DeriveInput, Generics, ItemStruct, Lifetime, LifetimeParam};
 
+pub use crate::postgres_type::entity::Alignment;
 use crate::{CodeEnrichment, ToSqlConfig};
 
 /// A parsed `#[derive(PostgresType)]` item.
@@ -55,6 +56,7 @@ pub struct PostgresTypeDerive {
     in_fn: Ident,
     out_fn: Ident,
     to_sql_config: ToSqlConfig,
+    alignment: Alignment,
 }
 
 impl PostgresTypeDerive {
@@ -64,11 +66,12 @@ impl PostgresTypeDerive {
         in_fn: Ident,
         out_fn: Ident,
         to_sql_config: ToSqlConfig,
+        alignment: Alignment,
     ) -> Result<CodeEnrichment<Self>, syn::Error> {
         if !to_sql_config.overrides_default() {
             crate::ident_is_acceptable_to_postgres(&name)?;
         }
-        Ok(CodeEnrichment(Self { generics, name, in_fn, out_fn, to_sql_config }))
+        Ok(CodeEnrichment(Self { generics, name, in_fn, out_fn, to_sql_config, alignment }))
     }
 
     pub fn from_derive_input(
@@ -90,12 +93,14 @@ impl PostgresTypeDerive {
             &format!("{}_out", derive_input.ident).to_lowercase(),
             derive_input.ident.span(),
         );
+        let alignment = Alignment::from_attributes(derive_input.attrs.as_slice())?;
         Self::new(
             derive_input.ident,
             derive_input.generics,
             funcname_in,
             funcname_out,
             to_sql_config,
+            alignment,
         )
     }
 }
@@ -129,6 +134,11 @@ impl ToEntityGraphTokens for PostgresTypeDerive {
 
         let to_sql_config = &self.to_sql_config;
 
+        let alignment = match &self.alignment {
+            Alignment::On => quote! { Some(::std::mem::align_of::<#name>()) },
+            Alignment::Off => quote! { None },
+        };
+
         quote! {
             unsafe impl #impl_generics ::pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable for #name #ty_generics #where_clauses {
                 fn argument_sql() -> core::result::Result<::pgrx::pgrx_sql_entity_graph::metadata::SqlMapping, ::pgrx::pgrx_sql_entity_graph::metadata::ArgumentError> {
@@ -190,6 +200,7 @@ impl ToEntityGraphTokens for PostgresTypeDerive {
                         path_items.join("::")
                     },
                     to_sql_config: #to_sql_config,
+                    alignment: #alignment,
                 };
                 ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::Type(submission)
             }
@@ -205,6 +216,7 @@ impl Parse for CodeEnrichment<PostgresTypeDerive> {
         let to_sql_config = ToSqlConfig::from_attributes(attrs.as_slice())?.unwrap_or_default();
         let in_fn = Ident::new(&format!("{}_in", ident).to_lowercase(), ident.span());
         let out_fn = Ident::new(&format!("{}_out", ident).to_lowercase(), ident.span());
-        PostgresTypeDerive::new(ident, generics, in_fn, out_fn, to_sql_config)
+        let alignment = Alignment::from_attributes(attrs.as_slice())?;
+        PostgresTypeDerive::new(ident, generics, in_fn, out_fn, to_sql_config, alignment)
     }
 }