From 8b2bee7f097e92802cb360bd1ebf962082cbafa1 Mon Sep 17 00:00:00 2001 From: Chris Bandy Date: Thu, 2 Jul 2026 01:08:51 +0000 Subject: [PATCH 1/2] Safe casting of PgNode values with bindgen tag discovery Postgres implements various node types to represent parse, plan, and executor trees. It uses tagged polymorphism to emulate object-oriented inheritance in C: every node struct begins with a `NodeTag` field that identifies its type at runtime. Pointers to these nodes are accessed generically as `Node *` and queried using the `IsA` macro. The hierarchy includes intermediate base types, such as `Expr`, which are inherited by concrete types like `Const` and `Var`. This commit introduces type-safe node casting (both upcasting and downcasting) via the `PgNode` trait. Safe casts are statically generated with the trait implementations emitted by `pgrx-bindgen`. Each node type may be cast to any of its "parent" types, all the way up to `Node`, because they inherently occupy less space in memory and on the stack. The `PgNode` trait is extended with: - A constant `CAST_TAGS` slice containing the valid tags for the type and all of its descendants. This is efficient because the hierarchy is shallow and wide. The vast majority of types have four or fewer tags in `CAST_TAGS`. - A `try_cast` associated function and a `try_as` method check `CAST_TAGS` before performing an unsafe pointer cast. The lookup is a simple linear scan and relies on compiler and CPU optimizations. - A `node_tag` method that returns the tag and an `is_a` method that behaves like the `IsA` macro. See: pgcentralfoundation/pgrx#1048 --- .gitattributes | 5 + pgrx-bindgen/src/build.rs | 202 ++++++++++++++++++++++++++------------ pgrx-pg-sys/src/node.rs | 87 +++++++++++++++- 3 files changed, 230 insertions(+), 64 deletions(-) diff --git a/.gitattributes b/.gitattributes index 5a6aef2427..a4df70f4d8 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,4 +1,9 @@ # https://git-scm.com/docs/gitattributes +# https://docs.gitlab.com/user/project/merge_requests/changes#collapse-generated-files +# https://github.com/github-linguist/linguist/blob/-/docs/overrides.md#generated-code +[attr]generated gitlab-generated linguist-generated + # Preserve LF line endings in golden files during checkout. /cargo-pgrx/tests/fixtures/**/*.toml text eol=lf +/pgrx-pg-sys/src/include/*.rs generated diff --git a/pgrx-bindgen/src/build.rs b/pgrx-bindgen/src/build.rs index 46d58f6556..0b6f84bd5f 100644 --- a/pgrx-bindgen/src/build.rs +++ b/pgrx-bindgen/src/build.rs @@ -494,71 +494,114 @@ fn format_builtin_oid_impl(oids: BTreeMap>) -> proc_m /// Implement our `PgNode` marker trait for `pg_sys::Node` and its "subclasses" fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { - let mut pgnode_impls = proc_macro2::TokenStream::new(); + let struct_graph = StructGraph::from(items); - // we scope must of the computation so we can borrow `items` and then - // extend it at the very end. - let struct_graph: StructGraph = StructGraph::from(items); - - // collect all the structs with `NodeTag` as their first member, - // these will serve as roots in our forest of `Node`s - let mut root_node_structs = Vec::new(); - for descriptor in struct_graph.descriptors.iter() { - // grab the first field, if any - let first_field = match &descriptor.struct_.fields { - syn::Fields::Named(fields) => { - if let Some(first_field) = fields.named.first() { - first_field - } else { - continue; - } + // Look through the entire file to produce a set of all variants of the Postgres `NodeTag` enum. + // Also look at type aliases of structs, as these could be node structs, too. + let mut node_tags: BTreeSet = BTreeSet::new(); + let mut possible_alias_tags = HashMap::new(); + for item in items { + match item { + // the `NodeTag` enum + syn::Item::Enum(item_enum) if item_enum.ident == "NodeTag" => { + node_tags.extend(item_enum.variants.iter().map(|v| v.ident.to_string())) } - syn::Fields::Unnamed(fields) => { - if let Some(first_field) = fields.unnamed.first() { - first_field - } else { - continue; - } + // one type alias of a struct; e.g. `pub type DistinctExpr = OpExpr` + syn::Item::Type(item_type) + if let syn::Type::Path(p) = &*item_type.ty + && let Some(last) = p.path.segments.last() + && struct_graph.name_tab.contains_key(&last.ident.to_string()) => + { + let target_struct_name = last.ident.to_string(); + let alias_name = item_type.ident.to_string(); + let tag_name = format!("T_{}", alias_name); + possible_alias_tags + .entry(target_struct_name) + .or_insert_with(BTreeSet::new) + .insert(tag_name); } _ => continue, + } + } + + // Identify the root nodes of the Postgres inheritance hierarchy (those having `NodeTag` as + // their first field) and recursively resolve the cast tags for them and their subclasses. + let mut identified_nodes = BTreeMap::new(); + for descriptor in struct_graph.descriptors.iter() { + let first_field = { + if let syn::Fields::Named(fields) = &descriptor.struct_.fields + && let Some(first) = fields.named.first() + { + first + } else if let syn::Fields::Unnamed(fields) = &descriptor.struct_.fields + && let Some(first) = fields.unnamed.first() + { + first + } else { + continue; + } }; - // grab the type name of the first field - let ty_name = if let syn::Type::Path(p) = &first_field.ty - && let Some(last_segment) = p.path.segments.last() - { - last_segment.ident.to_string() - } else { - continue; + let field_type_name = { + if let syn::Type::Path(p) = &first_field.ty + && let Some(last) = p.path.segments.last() + { + last.ident.to_string() + } else { + continue; + } }; - if ty_name == "NodeTag" { - root_node_structs.push(descriptor); + if field_type_name == "NodeTag" { + resolve_pg_node_tags( + descriptor, + &struct_graph, + &node_tags, + &possible_alias_tags, + &mut identified_nodes, + ); } } - // the set of types which subclass `Node` according to postgres' object system - let mut node_set = BTreeSet::new(); - // fill in any children of the roots with a recursive DFS - // (we are not operating on user input, so it is ok to just - // use direct recursion rather than an explicit stack). - for root in root_node_structs.into_iter() { - dfs_find_nodes(root, &struct_graph, &mut node_set); - } + // Finally, emit `PgNode` implementations for every detected Node. + let mut impls = proc_macro2::TokenStream::new(); + for (struct_name, cast_tags) in identified_nodes { + let ident_struct_name = syn::Ident::new(&struct_name, proc_macro2::Span::call_site()); + let ident_cast_tags: Vec = + cast_tags.iter().map(|t| syn::Ident::new(t, proc_macro2::Span::call_site())).collect(); + + // Seal every Node. + impls.extend(quote! { + impl pg_sys::seal::Sealed for #ident_struct_name {} + }); + + // Implement PgNode for every Node. + impls.extend(match struct_name.as_str() { + // Override the default implementation of `try_cast` for Node. + "Node" => quote! { + impl pg_sys::PgNode for #ident_struct_name { + const CAST_TAGS: &'static [pg_sys::NodeTag] = &[]; - // now we can finally iterate the Nodes and emit out Display impl - for node_struct in node_set.into_iter() { - let struct_name = &node_struct.struct_.ident; + #[inline] + fn try_cast(node: &T) -> Option<&Self> { + Some(node.as_node()) + } + } + }, - // 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 {} + // Use the default implementation of `try_as` with populated CAST_TAGS. + _ => quote! { + impl pg_sys::PgNode for #ident_struct_name { + const CAST_TAGS: &'static [pg_sys::NodeTag] = &[ + #(pg_sys::NodeTag::#ident_cast_tags),* + ]; + } + }, }); - // impl Rust's Display trait for all nodes - pgnode_impls.extend(quote! { - impl ::core::fmt::Display for #struct_name { + // Implement Display for every Node. + impls.extend(quote! { + impl ::core::fmt::Display for #ident_struct_name { fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { self.display_node().fmt(f) } @@ -566,23 +609,56 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { }); } - Ok(pgnode_impls) + Ok(impls) } -/// Given a root node, dfs_find_nodes adds all its children nodes to `node_set`. -fn dfs_find_nodes<'graph>( - node: &'graph StructDescriptor<'graph>, - graph: &'graph StructGraph<'graph>, - node_set: &mut BTreeSet>, -) { - node_set.insert(node.clone()); - - for child in node.children(graph) { - if node_set.contains(child) { - continue; +/// Recursively traverse a Node's subclasses and return the union of cast node tags. +/// At the same time, collect results into `identified_nodes`. +fn resolve_pg_node_tags<'graph>( + descriptor: &'graph StructDescriptor<'graph>, + struct_graph: &'graph StructGraph<'graph>, + node_tags: &BTreeSet, + possible_alias_tags: &HashMap>, + identified_nodes: &mut BTreeMap>, +) -> BTreeSet { + let struct_name = descriptor.struct_.ident.to_string(); + if let Some(struct_tags) = identified_nodes.get(&struct_name) { + return struct_tags.clone(); + } + + let mut cast_tags = BTreeSet::new(); + + // Start with the struct name. Any Node with this tag can cast to this struct. + let possible_tag_name = format!("T_{}", struct_name); + if node_tags.contains(&possible_tag_name) { + cast_tags.insert(possible_tag_name); + } + + // Any Node with the tag of a typedef alias can also cast to this struct. + if let Some(possible_tags) = possible_alias_tags.get(&struct_name) { + for possible_tag_name in possible_tags { + if node_tags.contains(possible_tag_name) { + cast_tags.insert(possible_tag_name.clone()); + } } - dfs_find_nodes(child, graph, node_set); } + + // Recursively collect tags from child subclasses. + for child in descriptor.children(struct_graph) { + cast_tags.extend(resolve_pg_node_tags( + child, + struct_graph, + node_tags, + possible_alias_tags, + identified_nodes, + )); + } + + // Register this Node and its resolved tags in the final result set. + identified_nodes.insert(struct_name, cast_tags.clone()); + + // Return this Nodes' resolved tags. + cast_tags } /// A graph describing the inheritance relationships between different nodes diff --git a/pgrx-pg-sys/src/node.rs b/pgrx-pg-sys/src/node.rs index dc5f3eb9cb..c88b08958e 100644 --- a/pgrx-pg-sys/src/node.rs +++ b/pgrx-pg-sys/src/node.rs @@ -2,7 +2,29 @@ use core::ffi::CStr; use core::ptr::NonNull; /// A trait applied to all Postgres `pg_sys::Node` types and subtypes -pub trait PgNode: crate::seal::Sealed { +pub trait PgNode: crate::seal::Sealed + Sized { + /// [`PgNode`] values that have one of these tags can safely be cast to `Self` + const CAST_TAGS: &'static [crate::NodeTag]; + + /// Try to safely cast node to Self + fn try_cast(node: &T) -> Option<&Self> { + // This uses `slice::contains` to perform a linear search because the vast majority (>95%) + // of PgNode have four or fewer tags in `CAST_TAGS`. For small contiguous slices of ints, + // a linear scan is faster than structured lookups (like binary search or hash sets) due + // to cache locality, register/instruction efficiency, and compiler autovectorization. + if Self::CAST_TAGS.contains(&node.node_tag()) { + Some(unsafe { &*core::ptr::from_ref(node).cast() }) + } else { + None + } + } + + /// Upcast this node to a `pg_sys::Node` + #[inline] + fn as_node(&self) -> &crate::Node { + unsafe { &*core::ptr::from_ref(self).cast() } + } + /// Format this node /// /// # Safety @@ -20,6 +42,24 @@ pub trait PgNode: crate::seal::Sealed { // this is only implemented for things known to be "Nodes" unsafe { display_node_impl(NonNull::from(self).cast()) } } + + #[doc(alias = "IsA")] + #[inline] + fn is_a(&self, tag: crate::NodeTag) -> bool { + self.node_tag() == tag + } + + #[inline] + fn node_tag(&self) -> crate::NodeTag { + self.as_node().type_ + } + + /// Try to safely cast self to T + #[inline] + fn try_as(&self) -> Option<&T> { + // Delegate to the associated function which may be specialized in T. + T::try_cast(self) + } } /// implementation function for `impl Display for $NodeType` @@ -45,3 +85,48 @@ pub(crate) unsafe fn display_node_impl(node: NonNull) -> String { result } } + +#[cfg(test)] +mod tests { + use super::PgNode; + + #[test] + fn cast_roundtrip() { + let rt = crate::RangeTblRef { type_: crate::NodeTag::T_RangeTblRef, rtindex: 9 }; + + let node = rt.try_as::().expect("upcast to Node should succeed"); + let next = node.try_as::().expect("downcast to self should succeed"); + + assert_eq!(next.type_, crate::NodeTag::T_RangeTblRef); + assert_eq!(next.rtindex, 9); + } + + #[test] + fn inheritance_downcast() { + let alt_subplan = crate::AlternativeSubPlan { + xpr: crate::Expr { type_: crate::NodeTag::T_AlternativeSubPlan }, + subplans: std::ptr::null_mut(), + }; + + let node = alt_subplan.as_node(); // infallible upcast + let expr = node.try_as::().expect("downcast to parent of self should succeed"); + assert_eq!(expr.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains"); + + expr.try_as::().expect("downcast to self should succeed"); + assert!(node.try_as::().is_none(), "downcast to an unrelated type should fail"); + } + + #[test] + fn inheritance_upcast() { + let alt_subplan = crate::AlternativeSubPlan { + xpr: crate::Expr { type_: crate::NodeTag::T_AlternativeSubPlan }, + subplans: std::ptr::null_mut(), + }; + + let expr = alt_subplan.try_as::().expect("upcast to a parent should succeed"); + assert_eq!(expr.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains"); + + let node = expr.try_as::().expect("upcast to Node should succeed"); + assert_eq!(node.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains"); + } +} From a3bffcb98173daa99248f30ea4d0ad58445f3f5f Mon Sep 17 00:00:00 2001 From: Chris Bandy Date: Sat, 11 Jul 2026 01:06:13 +0000 Subject: [PATCH 2/2] Implement PgNode for Value and ValUnion types Postgres 13 and 14 use the `Value` struct to represent multiple possible value types, including integer, float, and string. Each type has its own `NodeTag`, but they can't be automatically associated with the struct using the information in the bindgen bindings. This hard-codes the associations just before emitting all the `PgNode` implementations. Postgres 15 introduced separate structs for each value type. These fit nicely into the automatic struct-to-tag scheme, but we weren't handling the `ValUnion` union Postgres uses to represent these multiple types in a single `Node`. This extends the detection of parent-child types to implement `PgNode` for union nodes. --- pgrx-bindgen/src/build.rs | 318 ++++++++++++++++++++------------------ pgrx-pg-sys/src/node.rs | 66 ++++++++ 2 files changed, 237 insertions(+), 147 deletions(-) diff --git a/pgrx-bindgen/src/build.rs b/pgrx-bindgen/src/build.rs index 0b6f84bd5f..fdd70cf8d2 100644 --- a/pgrx-bindgen/src/build.rs +++ b/pgrx-bindgen/src/build.rs @@ -14,7 +14,6 @@ use eyre::{WrapErr, eyre}; use pgrx_pg_config::{PgConfig, PgMinorVersion, PgVersion, Pgrx, SUPPORTED_VERSIONS}; use quote::{ToTokens, quote}; use std::cell::RefCell; -use std::cmp::Ordering; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::fs; use std::path::{self, Path, PathBuf}; // disambiguate path::Path and syn::Type::Path @@ -298,7 +297,7 @@ fn generate_bindings( .wrap_err_with(|| format!("bindgen failed for pg{major_version}"))?; let oids = extract_oids(&bindgen_output); - let rewritten_items = rewrite_items(bindgen_output, &oids) + let rewritten_items = rewrite_items(major_version, bindgen_output, &oids) .wrap_err_with(|| format!("failed to rewrite items for pg{major_version}"))?; let oids = format_builtin_oid_impl(oids); @@ -405,13 +404,14 @@ fn write_rs_file( /// Given a token stream representing a file, apply a series of transformations to munge /// the bindgen generated code with some postgres specific enhancements fn rewrite_items( + major_version: u16, mut file: syn::File, oids: &BTreeMap>, ) -> eyre::Result { rewrite_c_abi_to_c_unwind(&mut file); let items_vec = rewrite_oid_consts(&file.items, oids); let mut items = apply_pg_guard(&items_vec)?; - let pgnode_impls = impl_pg_node(&items_vec)?; + let pgnode_impls = impl_pg_node(major_version, &items_vec)?; // append the pgnodes to the set of items items.extend(pgnode_impls); @@ -493,11 +493,11 @@ fn format_builtin_oid_impl(oids: BTreeMap>) -> proc_m } /// Implement our `PgNode` marker trait for `pg_sys::Node` and its "subclasses" -fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { - let struct_graph = StructGraph::from(items); +fn impl_pg_node(major_version: u16, items: &[syn::Item]) -> eyre::Result { + let type_graph = TypeGraph::from(items); // Look through the entire file to produce a set of all variants of the Postgres `NodeTag` enum. - // Also look at type aliases of structs, as these could be node structs, too. + // Also look at type aliases of structs/unions, as these could be node types, too. let mut node_tags: BTreeSet = BTreeSet::new(); let mut possible_alias_tags = HashMap::new(); for item in items { @@ -506,17 +506,17 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { syn::Item::Enum(item_enum) if item_enum.ident == "NodeTag" => { node_tags.extend(item_enum.variants.iter().map(|v| v.ident.to_string())) } - // one type alias of a struct; e.g. `pub type DistinctExpr = OpExpr` + // one type alias of a struct/union; e.g. `pub type DistinctExpr = OpExpr` syn::Item::Type(item_type) if let syn::Type::Path(p) = &*item_type.ty && let Some(last) = p.path.segments.last() - && struct_graph.name_tab.contains_key(&last.ident.to_string()) => + && type_graph.name_tab.contains_key(&last.ident.to_string()) => { - let target_struct_name = last.ident.to_string(); + let target_name = last.ident.to_string(); let alias_name = item_type.ident.to_string(); let tag_name = format!("T_{}", alias_name); possible_alias_tags - .entry(target_struct_name) + .entry(target_name) .or_insert_with(BTreeSet::new) .insert(tag_name); } @@ -524,38 +524,47 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { } } - // Identify the root nodes of the Postgres inheritance hierarchy (those having `NodeTag` as - // their first field) and recursively resolve the cast tags for them and their subclasses. + // Identify the root nodes of the Postgres inheritance hierarchy and recursively resolve the + // cast tags for them and their subclasses. The `BTreeMap` returns nodes in alphabetical order + // when we emit the trait implementations. let mut identified_nodes = BTreeMap::new(); - for descriptor in struct_graph.descriptors.iter() { - let first_field = { - if let syn::Fields::Named(fields) = &descriptor.struct_.fields - && let Some(first) = fields.named.first() - { - first - } else if let syn::Fields::Unnamed(fields) = &descriptor.struct_.fields - && let Some(first) = fields.unnamed.first() - { - first - } else { - continue; - } - }; + for descriptor in type_graph.descriptors.iter() { + let is_node = match descriptor.kind { + // a node struct has a `NodeTag` for its first field + TypeKind::Struct(struct_) => { + let first_field = if let syn::Fields::Named(fields) = &struct_.fields { + fields.named.first() + } else if let syn::Fields::Unnamed(fields) = &struct_.fields { + fields.unnamed.first() + } else { + None + }; - let field_type_name = { - if let syn::Type::Path(p) = &first_field.ty - && let Some(last) = p.path.segments.last() - { - last.ident.to_string() - } else { - continue; + if let Some(first_field) = first_field + && let syn::Type::Path(p) = &first_field.ty + && let Some(last) = p.path.segments.last() + { + last.ident == "NodeTag" + } else { + false + } } + // a node union has one member that is a `Node` + TypeKind::Union(union_) => union_.fields.named.iter().any(|field| { + if let syn::Type::Path(p) = &field.ty + && let Some(last) = p.path.segments.last() + { + last.ident == "Node" + } else { + false + } + }), }; - if field_type_name == "NodeTag" { + if is_node { resolve_pg_node_tags( descriptor, - &struct_graph, + &type_graph, &node_tags, &possible_alias_tags, &mut identified_nodes, @@ -563,23 +572,34 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { } } + // The `Value` struct of pg13 and pg14 is used with multiple node tags, but there's nothing in + // the bindgen bindings to indicate that. For these versions, directly include the tags used in + // the constructors defined in `nodes/value.c`. + if (major_version == 13 || major_version == 14) + && let Some(cast_tags) = identified_nodes.get_mut("Value") + { + cast_tags.extend( + ["T_Integer", "T_Float", "T_String", "T_BitString", "T_Null"].map(|s| s.to_string()), + ); + } + // Finally, emit `PgNode` implementations for every detected Node. let mut impls = proc_macro2::TokenStream::new(); - for (struct_name, cast_tags) in identified_nodes { - let ident_struct_name = syn::Ident::new(&struct_name, proc_macro2::Span::call_site()); + for (type_name, cast_tags) in identified_nodes { + let ident_type_name = syn::Ident::new(&type_name, proc_macro2::Span::call_site()); let ident_cast_tags: Vec = cast_tags.iter().map(|t| syn::Ident::new(t, proc_macro2::Span::call_site())).collect(); // Seal every Node. impls.extend(quote! { - impl pg_sys::seal::Sealed for #ident_struct_name {} + impl pg_sys::seal::Sealed for #ident_type_name {} }); // Implement PgNode for every Node. - impls.extend(match struct_name.as_str() { + impls.extend(match type_name.as_str() { // Override the default implementation of `try_cast` for Node. "Node" => quote! { - impl pg_sys::PgNode for #ident_struct_name { + impl pg_sys::PgNode for #ident_type_name { const CAST_TAGS: &'static [pg_sys::NodeTag] = &[]; #[inline] @@ -588,10 +608,9 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { } } }, - // Use the default implementation of `try_as` with populated CAST_TAGS. _ => quote! { - impl pg_sys::PgNode for #ident_struct_name { + impl pg_sys::PgNode for #ident_type_name { const CAST_TAGS: &'static [pg_sys::NodeTag] = &[ #(pg_sys::NodeTag::#ident_cast_tags),* ]; @@ -601,7 +620,7 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { // Implement Display for every Node. impls.extend(quote! { - impl ::core::fmt::Display for #ident_struct_name { + impl ::core::fmt::Display for #ident_type_name { fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { self.display_node().fmt(f) } @@ -615,27 +634,27 @@ fn impl_pg_node(items: &[syn::Item]) -> eyre::Result { /// Recursively traverse a Node's subclasses and return the union of cast node tags. /// At the same time, collect results into `identified_nodes`. fn resolve_pg_node_tags<'graph>( - descriptor: &'graph StructDescriptor<'graph>, - struct_graph: &'graph StructGraph<'graph>, + descriptor: &'graph TypeDescriptor<'graph>, + type_graph: &'graph TypeGraph<'graph>, node_tags: &BTreeSet, possible_alias_tags: &HashMap>, identified_nodes: &mut BTreeMap>, ) -> BTreeSet { - let struct_name = descriptor.struct_.ident.to_string(); - if let Some(struct_tags) = identified_nodes.get(&struct_name) { - return struct_tags.clone(); + let type_name = descriptor.ident.to_string(); + if let Some(tags) = identified_nodes.get(&type_name) { + return tags.clone(); } let mut cast_tags = BTreeSet::new(); - // Start with the struct name. Any Node with this tag can cast to this struct. - let possible_tag_name = format!("T_{}", struct_name); + // Start with the type name. Any Node with this tag can cast to this type. + let possible_tag_name = format!("T_{}", type_name); if node_tags.contains(&possible_tag_name) { cast_tags.insert(possible_tag_name); } - // Any Node with the tag of a typedef alias can also cast to this struct. - if let Some(possible_tags) = possible_alias_tags.get(&struct_name) { + // Any Node with the tag of a typedef alias can also cast to this type. + if let Some(possible_tags) = possible_alias_tags.get(&type_name) { for possible_tag_name in possible_tags { if node_tags.contains(possible_tag_name) { cast_tags.insert(possible_tag_name.clone()); @@ -643,126 +662,142 @@ fn resolve_pg_node_tags<'graph>( } } - // Recursively collect tags from child subclasses. - for child in descriptor.children(struct_graph) { - cast_tags.extend(resolve_pg_node_tags( - child, - struct_graph, - node_tags, - possible_alias_tags, - identified_nodes, - )); + // Unions do not inherit their member's node tags because it is not always safe to cast in that + // direction. The sizeof a union member may be smaller than the union, and casting from the + // former to the latter leads to out-of-bounds reads and UB. The CAST_TAGS of a union should + // contain, at most, its own node tag and aliases. + if let TypeKind::Struct(_) = descriptor.kind { + for child in descriptor.children(type_graph) { + cast_tags.extend(resolve_pg_node_tags( + child, + type_graph, + node_tags, + possible_alias_tags, + identified_nodes, + )); + } } // Register this Node and its resolved tags in the final result set. - identified_nodes.insert(struct_name, cast_tags.clone()); + identified_nodes.insert(type_name, cast_tags.clone()); // Return this Nodes' resolved tags. cast_tags } -/// A graph describing the inheritance relationships between different nodes +#[derive(Clone, Debug)] +enum TypeKind<'a> { + Struct(&'a syn::ItemStruct), + Union(&'a syn::ItemUnion), +} + +/// A graph describing the inheritance relationships between different types /// according to postgres' object system. /// -/// NOTE: the borrowed lifetime on a StructGraph should also ensure that the offsets +/// NOTE: the borrowed lifetime on a TypeGraph should also ensure that the offsets /// it stores into the underlying items struct are always correct. #[derive(Clone, Debug)] -struct StructGraph<'a> { +struct TypeGraph<'a> { #[allow(dead_code)] - /// A table mapping struct names to their offset in the descriptor table + /// A table mapping type names to their offset in the descriptor table name_tab: HashMap, #[allow(dead_code)] /// A table mapping offsets into the underlying items table to offsets in the descriptor table item_offset_tab: Vec>, - /// A table of struct descriptors - descriptors: Vec>, + /// A table of type descriptors + descriptors: Vec>, } -impl<'a> From<&'a [syn::Item]> for StructGraph<'a> { +impl<'a> From<&'a [syn::Item]> for TypeGraph<'a> { fn from(items: &'a [syn::Item]) -> Self { let mut descriptors = Vec::new(); - // a table mapping struct names to their offset in `descriptors` + // a table mapping type names to their offset in `descriptors` let mut name_tab: HashMap = HashMap::new(); let mut item_offset_tab: Vec> = vec![None; items.len()]; for (i, item) in items.iter().enumerate() { - if let &syn::Item::Struct(struct_) = &item { - let next_offset = descriptors.len(); - descriptors.push(StructDescriptor { - struct_, - items_offset: i, - parent: None, - children: Vec::new(), - }); - name_tab.insert(struct_.ident.to_string(), next_offset); - item_offset_tab[i] = Some(next_offset); - } + let (kind, ident) = match item { + syn::Item::Struct(struct_) => (TypeKind::Struct(struct_), struct_.ident.clone()), + syn::Item::Union(union_) => (TypeKind::Union(union_), union_.ident.clone()), + _ => continue, + }; + + let next_offset = descriptors.len(); + descriptors.push(TypeDescriptor { + kind, + ident: ident.clone(), + items_offset: i, + parent: None, + children: Vec::new(), + }); + name_tab.insert(ident.to_string(), next_offset); + item_offset_tab[i] = Some(next_offset); } for item in items.iter() { - // grab the first field if it is struct - let (id, first_field) = match &item { - syn::Item::Struct(syn::ItemStruct { - ident: id, - fields: syn::Fields::Named(fields), - .. - }) => { - if let Some(first_field) = fields.named.first() { - (id.to_string(), first_field) + match item { + // Structs represent Postgres' single-inheritance hierarchy: when the first field of + // a node struct is another node struct, the former "inherits" from the latter. The + // first field of a struct type is its parent type. + syn::Item::Struct(struct_) => { + let first_field = if let syn::Fields::Named(fields) = &struct_.fields { + fields.named.first() + } else if let syn::Fields::Unnamed(fields) = &struct_.fields { + fields.unnamed.first() } else { - continue; + None + }; + + if let Some(first_field) = first_field + && let syn::Type::Path(p) = &first_field.ty + && let Some(last_segment) = p.path.segments.last() + && let Some(parent_offset) = name_tab.get(&last_segment.ident.to_string()) + { + let child_offset = name_tab[&struct_.ident.to_string()]; + descriptors[child_offset].parent = Some(*parent_offset); + descriptors[*parent_offset].children.push(child_offset); } } - &syn::Item::Struct(syn::ItemStruct { - ident: id, - fields: syn::Fields::Unnamed(fields), - .. - }) => { - if let Some(first_field) = fields.unnamed.first() { - (id.to_string(), first_field) - } else { - continue; + // Unions represent a polymorphic container where each field is a subclass of the + // union. The union is the (abstract) parent type of the field types. + syn::Item::Union(union_) => { + let union_offset = name_tab[&union_.ident.to_string()]; + for field in &union_.fields.named { + if let syn::Type::Path(p) = &field.ty + && let Some(last_segment) = p.path.segments.last() + && let Some(child_offset) = + name_tab.get(&last_segment.ident.to_string()) + { + descriptors[*child_offset].parent = Some(union_offset); + descriptors[union_offset].children.push(*child_offset); + } } } _ => continue, - }; - - // We should be guaranteed that just extracting the last path - // segment is ok because these structs are all from the same module. - // (also, they are all generated from C code, so collisions should be - // impossible anyway thanks to C's single shared namespace). - if let syn::Type::Path(p) = &first_field.ty - && let Some(last_segment) = p.path.segments.last() - && let Some(parent_offset) = name_tab.get(&last_segment.ident.to_string()) - { - // establish the 2-way link - let child_offset = name_tab[&id]; - descriptors[child_offset].parent = Some(*parent_offset); - descriptors[*parent_offset].children.push(child_offset); } } - StructGraph { name_tab, item_offset_tab, descriptors } + TypeGraph { name_tab, item_offset_tab, descriptors } } } -impl<'a> StructDescriptor<'a> { +impl<'a> TypeDescriptor<'a> { /// children returns an iterator over the children of this node in the graph - fn children(&'a self, graph: &'a StructGraph) -> StructDescriptorChildren<'a> { - StructDescriptorChildren { offset: 0, descriptor: self, graph } + fn children(&'a self, graph: &'a TypeGraph) -> TypeDescriptorChildren<'a> { + TypeDescriptorChildren { offset: 0, descriptor: self, graph } } } -/// An iterator over a StructDescriptor's children -struct StructDescriptorChildren<'a> { +/// An iterator over a TypeDescriptor's children +struct TypeDescriptorChildren<'a> { offset: usize, - descriptor: &'a StructDescriptor<'a>, - graph: &'a StructGraph<'a>, + descriptor: &'a TypeDescriptor<'a>, + graph: &'a TypeGraph<'a>, } -impl<'a> std::iter::Iterator for StructDescriptorChildren<'a> { - type Item = &'a StructDescriptor<'a>; - fn next(&mut self) -> Option<&'a StructDescriptor<'a>> { +impl<'a> std::iter::Iterator for TypeDescriptorChildren<'a> { + type Item = &'a TypeDescriptor<'a>; + fn next(&mut self) -> Option<&'a TypeDescriptor<'a>> { if self.offset >= self.descriptor.children.len() { None } else { @@ -773,34 +808,23 @@ impl<'a> std::iter::Iterator for StructDescriptorChildren<'a> { } } -/// A node a StructGraph -#[derive(Clone, Debug, Hash, Eq, PartialEq)] -struct StructDescriptor<'a> { - /// A reference to the underlying struct syntax node - struct_: &'a syn::ItemStruct, +/// A node in a TypeGraph +#[derive(Clone, Debug)] +struct TypeDescriptor<'a> { + /// The kind of type (Struct or Union) + kind: TypeKind<'a>, + /// The identifier of the type + ident: syn::Ident, + #[allow(dead_code)] /// An offset into the items slice that was used to construct the struct graph that - /// this StructDescriptor is a part of + /// this TypeDescriptor is a part of items_offset: usize, - /// The offset of the "parent" (first member) struct (if any). + /// The offset of the "parent" struct/union (if any). parent: Option, - /// The offsets of the "children" structs (if any). + /// The offsets of the "children" structs/unions (if any). children: Vec, } -impl PartialOrd for StructDescriptor<'_> { - #[inline] - fn partial_cmp(&self, other: &StructDescriptor) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for StructDescriptor<'_> { - #[inline] - fn cmp(&self, other: &StructDescriptor) -> Ordering { - self.struct_.ident.cmp(&other.struct_.ident) - } -} - fn get_bindings( major_version: u16, pg_config: &PgConfig, diff --git a/pgrx-pg-sys/src/node.rs b/pgrx-pg-sys/src/node.rs index c88b08958e..5a0f20d0c4 100644 --- a/pgrx-pg-sys/src/node.rs +++ b/pgrx-pg-sys/src/node.rs @@ -129,4 +129,70 @@ mod tests { let node = expr.try_as::().expect("upcast to Node should succeed"); assert_eq!(node.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains"); } + + // In versions prior to pg15, the `Value` struct is used with tags T_Null, T_Integer, T_Float, + // T_String, and T_BitString. + #[cfg(any(feature = "pg13", feature = "pg14"))] + #[test] + fn value_struct_cast() { + let value = + crate::Value { type_: crate::NodeTag::T_Integer, val: crate::ValUnion { ival: 42 } }; + + let node = value.as_node(); // infallible upcast + assert_eq!(node.node_tag(), crate::NodeTag::T_Integer, "tag remains"); + + let next = node.try_as::().expect("downcast to self should succeed"); + assert_eq!(next.node_tag(), crate::NodeTag::T_Integer, "tag remains"); + assert_eq!(unsafe { next.val.ival }, 42); + + for tag in &[ + crate::NodeTag::T_Float, + crate::NodeTag::T_String, + crate::NodeTag::T_BitString, + crate::NodeTag::T_Null, + ] { + let value = crate::Value { + type_: *tag, + val: crate::ValUnion { str_: c"something".as_ptr() as *mut _ }, + }; + + let node = value.as_node(); // infallible upcast + assert_eq!(node.node_tag(), *tag, "tag remains"); + + let next = node.try_as::().expect("downcast to self should succeed"); + assert_eq!(next.node_tag(), *tag, "tag remains"); + assert_eq!(unsafe { next.val.str_ }, c"something".as_ptr() as *mut _); + } + } + + // Since pg15, Boolean, Integer, Float, String, and BitString are separate node structs. + // `ValUnion` is a single node that can hold any one of these types. That union should implement + // PgNode so we can cast to one of the above structs. + #[cfg(any( + feature = "pg15", + feature = "pg16", + feature = "pg17", + feature = "pg18", + feature = "pg19" + ))] + #[test] + fn value_union_cast() { + let vu = crate::ValUnion { + sval: crate::String { + type_: crate::NodeTag::T_String, + sval: c"something".as_ptr() as *mut _, + }, + }; + + let node = vu.try_as::().expect("upcast to Node should succeed"); + let next = node.try_as::().expect("downcast to self should succeed"); + + assert_eq!(next.type_, crate::NodeTag::T_String); + assert_eq!(next.sval, c"something".as_ptr() as *mut _); + + assert!( + node.try_as::().is_none(), + "downcast to untagged union should fail" + ); + } }