[WIP] support casting between pgnodes

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::<HashSet<_>>();
+
     // 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<NodeTag> = #node_tag;
+            }
         });
 
         // impl Rust's Display trait for all nodes

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<NodeTag>;
+
     /// Format this node
     ///
     /// # Safety
@@ -269,6 +273,66 @@ impl AsPgCStr for &Option<std::string::String> {
     }
 }
 
+pub trait PgNodeCast<T> {
+    fn cast_from(pg_node: T) -> Self
+    where
+        Self: Sized;
+}
+
+pub trait PgNodeTryCast<T> {
+    fn try_cast_from(pg_node: T) -> Option<Self>
+    where
+        Self: Sized;
+}
+
+impl<A: PgNode, B: PgNode> PgNodeCast<&A> for &B {
+    fn cast_from(pg_node: &A) -> Self {
+        Self::try_cast_from(pg_node).unwrap()
+    }
+}
+
+impl<A: PgNode, B: PgNode> PgNodeTryCast<&A> for &B {
+    fn try_cast_from(pg_node: &A) -> Option<Self> {
+        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<A: PgNode, B: PgNode> PgNodeCast<&mut A> for &mut B {
+    fn cast_from(pg_node: &mut A) -> Self {
+        Self::try_cast_from(pg_node).unwrap()
+    }
+}
+
+impl<A: PgNode, B: PgNode> PgNodeTryCast<&mut A> for &mut B {
+    fn try_cast_from(pg_node: &mut A) -> Option<Self> {
+        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;