You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by lu...@apache.org on 2026/10/08 20:41:21 UTC
(datafusion-sqlparser-rs) 01/03: Add Visitor::pre_visit_node and post_visit_node for every AST node
This is an automated email from the ASF dual-hosted git repository.
LucaCappelletti94 pushed a commit to branch visit-node-hooks
in repository https://gitbox.apache.org/repos/asf/datafusion-sqlparser-rs.git
commit bba0133b80a150eedaa77ef23ace0f304363439e
Author: LucaCappelletti94 <ca...@gmail.com>
AuthorDate: Thu Sep 24 09:40:36 2026 +0200
Add Visitor::pre_visit_node and post_visit_node for every AST node
---
derive/src/visit.rs | 42 ++++++++++-
src/ast/visitor.rs | 195 +++++++++++++++++++++++++++++++++++++++++++++++++++-
2 files changed, 233 insertions(+), 4 deletions(-)
diff --git a/derive/src/visit.rs b/derive/src/visit.rs
index cb02733b..30e03fff 100644
--- a/derive/src/visit.rs
+++ b/derive/src/visit.rs
@@ -46,12 +46,35 @@ pub(crate) fn derive_visit(
} = visit_type;
let attributes = Attributes::parse(&input.attrs);
- // Add a bound `T: Visit` to every type parameter T.
+ // Add a bound `T: Visit` to every type parameter T, and `T: 'static` for `Visit`.
let generics = add_trait_bounds(input.generics, visit_type);
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
let (pre_visit, post_visit) = attributes.visit(quote!(self));
let children = visit_children(&input.data, visit_type);
+ let (pre_visit_node, post_visit_node) = if modifier.is_none() {
+ (
+ Some(quote! {
+ let node = {
+ use sqlparser::ast::__private::{
+ DisplayFallback as _, DisplayProbe as _, SpannedFallback as _,
+ SpannedProbe as _,
+ };
+ let probe = sqlparser::ast::__private::Probe(self);
+ sqlparser::ast::NodeRef::__new(
+ self,
+ ::core::any::type_name::<Self>(),
+ (&&probe).spanned(),
+ (&&probe).display(),
+ )
+ };
+ visitor.pre_visit_node(node)?;
+ }),
+ Some(quote!(visitor.post_visit_node(node)?;)),
+ )
+ } else {
+ (None, None)
+ };
let expanded = quote! {
// The generated impl.
@@ -63,9 +86,11 @@ pub(crate) fn derive_visit(
&#modifier self,
visitor: &mut V
) -> ::core::ops::ControlFlow<V::Break> {
+ #pre_visit_node
#pre_visit
#children
#post_visit
+ #post_visit_node
::core::ops::ControlFlow::Continue(())
}
}
@@ -137,13 +162,24 @@ impl Attributes {
}
}
-// Add a bound `T: Visit` to every type parameter T.
-fn add_trait_bounds(mut generics: Generics, VisitType { visit_trait, .. }: &VisitType) -> Generics {
+// Add a bound `T: Visit` to every type parameter T, and `T: 'static` for
+// `Visit`, whose node hook passes the node as `&dyn Any`.
+fn add_trait_bounds(
+ mut generics: Generics,
+ VisitType {
+ visit_trait,
+ modifier,
+ ..
+ }: &VisitType,
+) -> Generics {
for param in &mut generics.params {
if let GenericParam::Type(ref mut type_param) = *param {
type_param
.bounds
.push(parse_quote!(sqlparser::ast::#visit_trait));
+ if modifier.is_none() {
+ type_param.bounds.push(parse_quote!('static));
+ }
}
}
generics
diff --git a/src/ast/visitor.rs b/src/ast/visitor.rs
index c78ff2bf..cd2d3309 100644
--- a/src/ast/visitor.rs
+++ b/src/ast/visitor.rs
@@ -19,10 +19,12 @@
#[cfg(not(feature = "std"))]
use alloc::{boxed::Box, string::String, vec::Vec};
+use core::any::Any;
+use core::fmt;
use core::ops::ControlFlow;
use crate::ast::{
- Expr, GroupByExpr, Ident, ObjectName, OrderBy, OrderByExpr, Query, Select, Statement,
+ Expr, GroupByExpr, Ident, ObjectName, OrderBy, OrderByExpr, Query, Select, Spanned, Statement,
TableFactor, ValueWithSpan,
};
@@ -65,6 +67,103 @@ pub trait VisitMut {
fn visit<V: VisitorMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break>;
}
+/// Any AST node, as passed to [`Visitor::pre_visit_node`] and
+/// [`Visitor::post_visit_node`].
+#[derive(Clone, Copy)]
+pub struct NodeRef<'a> {
+ any: &'a dyn Any,
+ type_name: &'static str,
+ spanned: Option<&'a dyn Spanned>,
+ display: Option<&'a dyn fmt::Display>,
+}
+
+impl<'a> NodeRef<'a> {
+ #[doc(hidden)]
+ pub fn __new(
+ any: &'a dyn Any,
+ type_name: &'static str,
+ spanned: Option<&'a dyn Spanned>,
+ display: Option<&'a dyn fmt::Display>,
+ ) -> Self {
+ Self {
+ any,
+ type_name,
+ spanned,
+ display,
+ }
+ }
+
+ /// The node as `T`, if it is a `T`.
+ pub fn downcast_ref<T: Any>(&self) -> Option<&'a T> {
+ self.any.downcast_ref()
+ }
+
+ /// The [`core::any::type_name`] of the node.
+ pub fn type_name(&self) -> &'static str {
+ self.type_name
+ }
+
+ /// The node, if its type implements [`Spanned`].
+ pub fn spanned(&self) -> Option<&'a dyn Spanned> {
+ self.spanned
+ }
+
+ /// The node, if its type implements [`fmt::Display`].
+ pub fn display(&self) -> Option<&'a dyn fmt::Display> {
+ self.display
+ }
+}
+
+impl fmt::Debug for NodeRef<'_> {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.write_str(self.type_name)
+ }
+}
+
+/// Autoref specialization used by the `Visit` derive to fill [`NodeRef`].
+#[doc(hidden)]
+pub mod __private {
+ use super::{fmt, Spanned};
+
+ pub struct Probe<'a, T>(pub &'a T);
+
+ pub trait SpannedProbe<'a> {
+ fn spanned(&self) -> Option<&'a dyn Spanned>;
+ }
+ impl<'a, T: Spanned> SpannedProbe<'a> for &Probe<'a, T> {
+ fn spanned(&self) -> Option<&'a dyn Spanned> {
+ Some(self.0)
+ }
+ }
+
+ pub trait SpannedFallback<'a> {
+ fn spanned(&self) -> Option<&'a dyn Spanned>;
+ }
+ impl<'a, T> SpannedFallback<'a> for Probe<'a, T> {
+ fn spanned(&self) -> Option<&'a dyn Spanned> {
+ None
+ }
+ }
+
+ pub trait DisplayProbe<'a> {
+ fn display(&self) -> Option<&'a dyn fmt::Display>;
+ }
+ impl<'a, T: fmt::Display> DisplayProbe<'a> for &Probe<'a, T> {
+ fn display(&self) -> Option<&'a dyn fmt::Display> {
+ Some(self.0)
+ }
+ }
+
+ pub trait DisplayFallback<'a> {
+ fn display(&self) -> Option<&'a dyn fmt::Display>;
+ }
+ impl<'a, T> DisplayFallback<'a> for Probe<'a, T> {
+ fn display(&self) -> Option<&'a dyn fmt::Display> {
+ None
+ }
+ }
+}
+
impl<T: Visit> Visit for Option<T> {
fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
if let Some(s) = self {
@@ -318,6 +417,44 @@ pub trait Visitor {
fn post_visit_group_by(&mut self, _group_by: &GroupByExpr) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
}
+
+ /// Invoked for every AST node before visiting its children, and before
+ /// the node's typed hook such as [`Visitor::pre_visit_expr`].
+ ///
+ /// # Example
+ /// ```
+ /// # use sqlparser::parser::Parser;
+ /// # use sqlparser::dialect::GenericDialect;
+ /// # use sqlparser::ast::{Join, NodeRef, Visit, Visitor};
+ /// # use core::ops::ControlFlow;
+ /// struct Joins(Vec<String>);
+ ///
+ /// impl Visitor for Joins {
+ /// type Break = ();
+ ///
+ /// fn pre_visit_node(&mut self, node: NodeRef<'_>) -> ControlFlow<()> {
+ /// if let Some(join) = node.downcast_ref::<Join>() {
+ /// self.0.push(join.relation.to_string());
+ /// }
+ /// ControlFlow::Continue(())
+ /// }
+ /// }
+ ///
+ /// let sql = "SELECT * FROM a JOIN b ON a.id = b.id JOIN c ON b.id = c.id";
+ /// let statements = Parser::parse_sql(&GenericDialect {}, sql).unwrap();
+ /// let mut joins = Joins(Vec::new());
+ /// let _ = statements.visit(&mut joins);
+ /// assert_eq!(joins.0, ["b", "c"]);
+ /// ```
+ fn pre_visit_node(&mut self, _node: NodeRef<'_>) -> ControlFlow<Self::Break> {
+ ControlFlow::Continue(())
+ }
+
+ /// Invoked for every AST node after visiting its children, and after the
+ /// node's typed hook such as [`Visitor::post_visit_expr`].
+ fn post_visit_node(&mut self, _node: NodeRef<'_>) -> ControlFlow<Self::Break> {
+ ControlFlow::Continue(())
+ }
}
/// A visitor that can be used to mutate an AST tree.
@@ -1289,6 +1426,62 @@ mod tests {
do_visit("CREATE VIEW db1.v AS SELECT * FROM t", &mut visitor);
assert_eq!(visitor.relations, vec!["db1.v", "t"]);
}
+
+ #[derive(Default)]
+ struct NodeVisitor {
+ depth: usize,
+ nodes: Vec<(usize, &'static str, bool, Option<String>)>,
+ }
+
+ impl Visitor for NodeVisitor {
+ type Break = ();
+
+ fn pre_visit_node(&mut self, node: NodeRef<'_>) -> ControlFlow<Self::Break> {
+ let name = node.type_name().rsplit("::").next().unwrap();
+ let display = node.display().map(|d| d.to_string());
+ self.nodes
+ .push((self.depth, name, node.spanned().is_some(), display));
+ self.depth += 1;
+ ControlFlow::Continue(())
+ }
+
+ fn post_visit_node(&mut self, _node: NodeRef<'_>) -> ControlFlow<Self::Break> {
+ self.depth -= 1;
+ ControlFlow::Continue(())
+ }
+ }
+
+ #[test]
+ fn test_visit_node_reaches_every_derived_node() {
+ let mut visitor = NodeVisitor::default();
+ do_visit(
+ "SELECT CAST(a AS INT) AS x FROM t JOIN u ON t.id = u.id",
+ &mut visitor,
+ );
+ assert_eq!(visitor.depth, 0);
+
+ let find = |name: &str, display: &str| {
+ visitor
+ .nodes
+ .iter()
+ .position(|(_, n, _, d)| *n == name && d.as_deref() == Some(display))
+ .unwrap_or_else(|| panic!("{name} `{display}` not visited"))
+ };
+
+ let item = find("SelectItem", "CAST(a AS INT) AS x");
+ assert!(visitor.nodes[item].2);
+
+ let data_type = find("DataType", "INT");
+ assert!(!visitor.nodes[data_type].2);
+
+ let join = find("Join", "JOIN u ON t.id = u.id");
+ assert!(visitor.nodes[join].2);
+ let (join_depth, ..) = visitor.nodes[join];
+ assert_eq!(
+ visitor.nodes[join + 1],
+ (join_depth + 1, "TableFactor", true, Some("u".to_string()))
+ );
+ }
}
#[cfg(test)]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org