From 3ed26c9be5420e57bad86db03d76b44cacac9d39 Mon Sep 17 00:00:00 2001 From: Alex Chi Date: Thu, 16 Feb 2023 15:29:56 -0500 Subject: [PATCH] support casting between pgnodes Signed-off-by: Alex Chi --- pgx-pg-sys/build.rs | 27 +++++++++++++++++- pgx-pg-sys/src/lib.rs | 64 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/pgx-pg-sys/build.rs b/pgx-pg-sys/build.rs index dca68af98..380bcceda 100644 --- a/pgx-pg-sys/build.rs +++ b/pgx-pg-sys/build.rs @@ -502,14 +502,39 @@ fn impl_pg_node( Box::new(node_set.into_iter()) }; + // Though these structs looks like a node, they don't seem to have a corresponding [`NodeTag`]. + // Also for `Node`, it doesn't have a node tag. + let block_list = [ + "Node", + "MemoryContextData", + "JoinPath", + "Expr", + "HeapTupleTableSlot", + "VirtualTupleTableSlot", + "PartitionPruneStep", + "BufferHeapTupleTableSlot", + "MinimalTupleTableSlot", + ] + .into_iter() + .collect::>(); + // now we can finally iterate the Nodes and emit out Display impl for node_struct in nodes { let struct_name = &node_struct.struct_.ident; + let node_tag = if !block_list.contains(struct_name.to_string().as_str()) { + let node_tag = quote::format_ident!("NodeTag_T_{}", struct_name); + quote! { Some(pg_sys::#node_tag) } + } else { + quote! { None } + }; + // impl the PgNode trait for all nodes pgnode_impls.extend(quote! { impl pg_sys::seal::Sealed for #struct_name {} - impl pg_sys::PgNode for #struct_name {} + impl pg_sys::PgNode for #struct_name { + const NODE_TAG: Option = #node_tag; + } }); // impl Rust's Display trait for all nodes diff --git a/pgx-pg-sys/src/lib.rs b/pgx-pg-sys/src/lib.rs index dfd16ae08..7c395423e 100644 --- a/pgx-pg-sys/src/lib.rs +++ b/pgx-pg-sys/src/lib.rs @@ -158,6 +158,10 @@ pub use internal::pg15::*; /// A trait applied to all Postgres `pg_sys::Node` types and subtypes pub trait PgNode: seal::Sealed { + /// The [`NodeTag`] for this node. If it is `None` then it doesn't have a corresponding + /// [`NodeTag`]. + const NODE_TAG: Option; + /// Format this node /// /// # Safety @@ -269,6 +273,66 @@ impl AsPgCStr for &Option { } } +pub trait PgNodeCast { + fn cast_from(pg_node: T) -> Self + where + Self: Sized; +} + +pub trait PgNodeTryCast { + fn try_cast_from(pg_node: T) -> Option + where + Self: Sized; +} + +impl PgNodeCast<&A> for &B { + fn cast_from(pg_node: &A) -> Self { + Self::try_cast_from(pg_node).unwrap() + } +} + +impl PgNodeTryCast<&A> for &B { + fn try_cast_from(pg_node: &A) -> Option { + unsafe { + let node = std::mem::transmute::<_, &Node>(pg_node); + match B::NODE_TAG { + Some(tag) => { + if node.type_ == tag { + Some(std::mem::transmute::<_, &B>(node)) + } else { + None + } + } + None => None, + } + } + } +} + +impl PgNodeCast<&mut A> for &mut B { + fn cast_from(pg_node: &mut A) -> Self { + Self::try_cast_from(pg_node).unwrap() + } +} + +impl PgNodeTryCast<&mut A> for &mut B { + fn try_cast_from(pg_node: &mut A) -> Option { + unsafe { + let node = std::mem::transmute::<_, &mut Node>(pg_node); + match B::NODE_TAG { + Some(tag) => { + if node.type_ == tag { + Some(std::mem::transmute::<_, &mut B>(node)) + } else { + None + } + } + None => None, + } + } + } +} + /// item declarations we want to add to all versions mod all_versions { use crate as pg_sys;