From a2b47b83f9f2621e7e9ef94e4492270da2784a0e Mon Sep 17 00:00:00 2001
From: lovasoa <pere.jobs@gmail.com>
Date: Wed, 4 Jan 2023 15:30:34 +0100
Subject: [PATCH] use post_* visitors for mutable visits

My initial implementation of mutable visitors closely matched the existing immutable visitor implementations,
and used pre_* visitors (a mutable expression was first passed the visitor, then its children were visited).

After some initial usage, I realize this actually impractical, and can easily lead to infinite recursion when one isn't careful.
For instance a mutable visitor would replace x by f(x),
then would be called again on the nested x and result in f(f(x)),
then again resulting in f(f(f(x))), and so on.

So, this commit changes the behavior of the visit_*_mut functions to call the visitor on an expression only
once all of its children have been visited.

It also makes the documentation more explicit and adds an example.
---
 src/ast/visitor.rs | 50 ++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 6 deletions(-)

diff --git a/src/ast/visitor.rs b/src/ast/visitor.rs
index 553e014a4..1a9eeaa24 100644
--- a/src/ast/visitor.rs
+++ b/src/ast/visitor.rs
@@ -301,7 +301,7 @@ impl<E, F: FnMut(&ObjectName) -> ControlFlow<E>> Visitor for RelationVisitor<F>
 impl<E, F: FnMut(&mut ObjectName) -> ControlFlow<E>> VisitorMut for RelationVisitor<F> {
     type Break = E;
 
-    fn pre_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow<Self::Break> {
+    fn post_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow<Self::Break> {
         self.0(relation)
     }
 }
@@ -343,7 +343,10 @@ where
     ControlFlow::Continue(())
 }
 
-/// Invokes the provided closure on all relations (e.g. table names) present in `v`
+/// Invokes the provided closure with a mutable reference to all relations (e.g. table names)
+/// present in `v`.
+///
+/// When the closure mutates its argument, the new mutated relation will not be visited again.
 ///
 /// # Example
 /// ```
@@ -386,7 +389,7 @@ impl<E, F: FnMut(&Expr) -> ControlFlow<E>> Visitor for ExprVisitor<F> {
 impl<E, F: FnMut(&mut Expr) -> ControlFlow<E>> VisitorMut for ExprVisitor<F> {
     type Break = E;
 
-    fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
+    fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
         self.0(expr)
     }
 }
@@ -430,9 +433,14 @@ where
     ControlFlow::Continue(())
 }
 
-/// Invokes the provided closure on all expressions present in `v`
+/// Invokes the provided closure iteratively with a mutable reference to all expressions
+/// present in `v`.
+///
+/// This performs a depth-first search, so if the closure mutates the expression
 ///
 /// # Example
+///
+/// ## Remove all select limits in sub-queries
 /// ```
 /// # use sqlparser::parser::Parser;
 /// # use sqlparser::dialect::GenericDialect;
@@ -451,6 +459,35 @@ where
 ///
 /// assert_eq!(statements[0].to_string(), "SELECT (SELECT y FROM z) FROM t LIMIT 3");
 /// ```
+///
+/// ## Wrap column name in function call
+///
+/// This demonstrates how to effectively replace an expression with another more complicated one
+/// that references the original. This example avoids unnecessary allocations by using the
+/// [`std::mem`](std::mem) family of functions.
+///
+/// ```
+/// # use sqlparser::parser::Parser;
+/// # use sqlparser::dialect::GenericDialect;
+/// # use sqlparser::ast::{Expr, Function, FunctionArg, FunctionArgExpr, Ident, ObjectName, Value, visit_expressions_mut, visit_statements_mut};
+/// # use core::ops::ControlFlow;
+/// let sql = "SELECT x, y FROM t";
+/// let mut statements = Parser::parse_sql(&GenericDialect{}, sql).unwrap();
+///
+/// visit_expressions_mut(&mut statements, |expr| {
+///   if matches!(expr, Expr::Identifier(col_name) if col_name.value == "x") {
+///     let old_expr = std::mem::replace(expr, Expr::Value(Value::Null));
+///     *expr = Expr::Function(Function {
+///           name: ObjectName(vec![Ident::new("f")]),
+///           args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(old_expr))],
+///           over: None, distinct: false, special: false,
+///      });
+///   }
+///   ControlFlow::<()>::Continue(())
+/// });
+///
+/// assert_eq!(statements[0].to_string(), "SELECT f(x), y FROM t");
+/// ```
 pub fn visit_expressions_mut<V, E, F>(v: &mut V, f: F) -> ControlFlow<E>
 where
     V: VisitMut,
@@ -473,12 +510,13 @@ impl<E, F: FnMut(&Statement) -> ControlFlow<E>> Visitor for StatementVisitor<F>
 impl<E, F: FnMut(&mut Statement) -> ControlFlow<E>> VisitorMut for StatementVisitor<F> {
     type Break = E;
 
-    fn pre_visit_statement(&mut self, statement: &mut Statement) -> ControlFlow<Self::Break> {
+    fn post_visit_statement(&mut self, statement: &mut Statement) -> ControlFlow<Self::Break> {
         self.0(statement)
     }
 }
 
-/// Invokes the provided closure on all statements (e.g. `SELECT`, `CREATE TABLE`, etc) present in `v`
+/// Invokes the provided closure iteratively with a mutable reference to all statements
+/// present in `v` (e.g. `SELECT`, `CREATE TABLE`, etc).
 ///
 /// # Example
 /// ```