Skip to main content

pgrx_pg_sys/
node.rs

1use core::ffi::CStr;
2use core::ptr::NonNull;
3
4/// A trait applied to all Postgres `pg_sys::Node` types and subtypes
5pub trait PgNode: crate::seal::Sealed + Sized {
6    /// [`PgNode`] values that have one of these tags can safely be cast to `Self`
7    const CAST_TAGS: &'static [crate::NodeTag];
8
9    /// Try to safely cast node to Self
10    fn try_cast<T: PgNode>(node: &T) -> Option<&Self> {
11        // This uses `slice::contains` to perform a linear search because the vast majority (>95%)
12        // of PgNode have four or fewer tags in `CAST_TAGS`. For small contiguous slices of ints,
13        // a linear scan is faster than structured lookups (like binary search or hash sets) due
14        // to cache locality, register/instruction efficiency, and compiler autovectorization.
15        if Self::CAST_TAGS.contains(&node.node_tag()) {
16            Some(unsafe { &*core::ptr::from_ref(node).cast() })
17        } else {
18            None
19        }
20    }
21
22    /// Upcast this node to a `pg_sys::Node`
23    #[inline]
24    fn as_node(&self) -> &crate::Node {
25        unsafe { &*core::ptr::from_ref(self).cast() }
26    }
27
28    /// Format this node
29    ///
30    /// # Safety
31    ///
32    /// While pgrx controls the types for which [`PgNode`] is implemented and only pgrx can implement
33    /// [`PgNode`] it cannot control the members of those types and a user could assign any pointer
34    /// type member to something invalid that Postgres wouldn't understand.
35    ///
36    /// Because this function is used by `impl Display` we purposely don't mark it `unsafe`.  The
37    /// assumption here, which might be a bad one, is that only intentional misuse would actually
38    /// cause undefined behavior.
39    #[inline]
40    fn display_node(&self) -> String {
41        // SAFETY: The trait is pub but this impl is private, and
42        // this is only implemented for things known to be "Nodes"
43        unsafe { display_node_impl(NonNull::from(self).cast()) }
44    }
45
46    #[doc(alias = "IsA")]
47    #[inline]
48    fn is_a(&self, tag: crate::NodeTag) -> bool {
49        self.node_tag() == tag
50    }
51
52    #[inline]
53    fn node_tag(&self) -> crate::NodeTag {
54        self.as_node().type_
55    }
56
57    /// Try to safely cast self to T
58    #[inline]
59    fn try_as<T: PgNode>(&self) -> Option<&T> {
60        // Delegate to the associated function which may be specialized in T.
61        T::try_cast(self)
62    }
63}
64
65/// implementation function for `impl Display for $NodeType`
66///
67/// # Safety
68/// Don't use this on anything that doesn't impl PgNode, or the type may be off
69#[warn(unsafe_op_in_unsafe_fn)]
70pub(crate) unsafe fn display_node_impl(node: NonNull<crate::Node>) -> String {
71    // SAFETY: It's fine to call nodeToString with non-null well-typed pointers,
72    // and pg_sys::nodeToString() returns data via palloc, which is never null
73    // as Postgres will ERROR rather than giving us a null pointer,
74    // and Postgres starts and finishes constructing StringInfos by writing '\0'
75    unsafe {
76        let node_cstr = crate::nodeToString(node.as_ptr().cast());
77
78        let result = match CStr::from_ptr(node_cstr).to_str() {
79            Ok(cstr) => cstr.to_string(),
80            Err(e) => format!("<ffi error: {e:?}>"),
81        };
82
83        crate::pfree(node_cstr.cast());
84
85        result
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::PgNode;
92
93    #[test]
94    fn cast_roundtrip() {
95        let rt = crate::RangeTblRef { type_: crate::NodeTag::T_RangeTblRef, rtindex: 9 };
96
97        let node = rt.try_as::<crate::Node>().expect("upcast to Node should succeed");
98        let next = node.try_as::<crate::RangeTblRef>().expect("downcast to self should succeed");
99
100        assert_eq!(next.type_, crate::NodeTag::T_RangeTblRef);
101        assert_eq!(next.rtindex, 9);
102    }
103
104    #[test]
105    fn inheritance_downcast() {
106        let alt_subplan = crate::AlternativeSubPlan {
107            xpr: crate::Expr { type_: crate::NodeTag::T_AlternativeSubPlan },
108            subplans: std::ptr::null_mut(),
109        };
110
111        let node = alt_subplan.as_node(); // infallible upcast
112        let expr = node.try_as::<crate::Expr>().expect("downcast to parent of self should succeed");
113        assert_eq!(expr.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains");
114
115        expr.try_as::<crate::AlternativeSubPlan>().expect("downcast to self should succeed");
116        assert!(node.try_as::<crate::Var>().is_none(), "downcast to an unrelated type should fail");
117    }
118
119    #[test]
120    fn inheritance_upcast() {
121        let alt_subplan = crate::AlternativeSubPlan {
122            xpr: crate::Expr { type_: crate::NodeTag::T_AlternativeSubPlan },
123            subplans: std::ptr::null_mut(),
124        };
125
126        let expr = alt_subplan.try_as::<crate::Expr>().expect("upcast to a parent should succeed");
127        assert_eq!(expr.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains");
128
129        let node = expr.try_as::<crate::Node>().expect("upcast to Node should succeed");
130        assert_eq!(node.node_tag(), crate::NodeTag::T_AlternativeSubPlan, "tag remains");
131    }
132
133    // In versions prior to pg15, the `Value` struct is used with tags T_Null, T_Integer, T_Float,
134    // T_String, and T_BitString.
135    #[cfg(any(feature = "pg13", feature = "pg14"))]
136    #[test]
137    fn value_struct_cast() {
138        let value =
139            crate::Value { type_: crate::NodeTag::T_Integer, val: crate::ValUnion { ival: 42 } };
140
141        let node = value.as_node(); // infallible upcast
142        assert_eq!(node.node_tag(), crate::NodeTag::T_Integer, "tag remains");
143
144        let next = node.try_as::<crate::Value>().expect("downcast to self should succeed");
145        assert_eq!(next.node_tag(), crate::NodeTag::T_Integer, "tag remains");
146        assert_eq!(unsafe { next.val.ival }, 42);
147
148        for tag in &[
149            crate::NodeTag::T_Float,
150            crate::NodeTag::T_String,
151            crate::NodeTag::T_BitString,
152            crate::NodeTag::T_Null,
153        ] {
154            let value = crate::Value {
155                type_: *tag,
156                val: crate::ValUnion { str_: c"something".as_ptr() as *mut _ },
157            };
158
159            let node = value.as_node(); // infallible upcast
160            assert_eq!(node.node_tag(), *tag, "tag remains");
161
162            let next = node.try_as::<crate::Value>().expect("downcast to self should succeed");
163            assert_eq!(next.node_tag(), *tag, "tag remains");
164            assert_eq!(unsafe { next.val.str_ }, c"something".as_ptr() as *mut _);
165        }
166    }
167
168    // Since pg15, Boolean, Integer, Float, String, and BitString are separate node structs.
169    // `ValUnion` is a single node that can hold any one of these types. That union should implement
170    // PgNode so we can cast to one of the above structs.
171    #[cfg(any(
172        feature = "pg15",
173        feature = "pg16",
174        feature = "pg17",
175        feature = "pg18",
176        feature = "pg19"
177    ))]
178    #[test]
179    fn value_union_cast() {
180        let vu = crate::ValUnion {
181            sval: crate::String {
182                type_: crate::NodeTag::T_String,
183                sval: c"something".as_ptr() as *mut _,
184            },
185        };
186
187        let node = vu.try_as::<crate::Node>().expect("upcast to Node should succeed");
188        let next = node.try_as::<crate::String>().expect("downcast to self should succeed");
189
190        assert_eq!(next.type_, crate::NodeTag::T_String);
191        assert_eq!(next.sval, c"something".as_ptr() as *mut _);
192
193        assert!(
194            node.try_as::<crate::ValUnion>().is_none(),
195            "downcast to untagged union should fail"
196        );
197    }
198}