Skip to content

Commit

Permalink
Visitor for function arguments expressions (apache#7)
Browse files Browse the repository at this point in the history
  • Loading branch information
Amogh-Bharadwaj authored May 30, 2023
1 parent 52482b3 commit a879adf
Showing 1 changed file with 56 additions and 3 deletions.
59 changes: 56 additions & 3 deletions src/ast/visitor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,8 @@

//! Recursive visitors for ast Nodes. See [`Visitor`] for more details.
use crate::ast::{Expr, ObjectName, Statement};
use crate::ast::{Expr, FunctionArgExpr, ObjectName, Statement};
use core::ops::ControlFlow;

/// A type that can be visited by a [`Visitor`]. See [`Visitor`] for
/// recursively visiting parsed SQL statements.
///
Expand Down Expand Up @@ -199,6 +198,15 @@ pub trait Visitor {
ControlFlow::Continue(())
}

/// Invoked for function arguments that appear in the AST before visiting children
fn pre_visit_function_arg(&mut self, _expr: &FunctionArgExpr) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
}

/// Invoked for any function arguments that appear in the AST after visiting children
fn post_visit_function_arg(&mut self, _expr: &FunctionArgExpr) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
}
/// Invoked for any statements that appear in the AST before visiting children
fn pre_visit_statement(&mut self, _statement: &Statement) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
Expand Down Expand Up @@ -277,6 +285,16 @@ pub trait VisitorMut {
ControlFlow::Continue(())
}

/// Invoked for function arguments that appear in the AST before visiting children
fn pre_visit_function_arg(&mut self, _expr: &mut FunctionArgExpr) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
}

/// Invoked for any function arguments that appear in the AST after visiting children
fn post_visit_function_arg(&mut self, _expr: &mut FunctionArgExpr) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
}

/// Invoked for any statements that appear in the AST before visiting children
fn pre_visit_statement(&mut self, _statement: &mut Statement) -> ControlFlow<Self::Break> {
ControlFlow::Continue(())
Expand Down Expand Up @@ -376,6 +394,42 @@ where
ControlFlow::Continue(())
}

struct FunctionArgExprVisitor<F>(F);

impl<E, F: FnMut(&FunctionArgExpr) -> ControlFlow<E>> Visitor for FunctionArgExprVisitor<F> {
type Break = E;

fn pre_visit_function_arg(&mut self, expr: &FunctionArgExpr) -> ControlFlow<Self::Break> {
self.0(expr)
}
}

impl<E, F: FnMut(&mut FunctionArgExpr) -> ControlFlow<E>> VisitorMut for FunctionArgExprVisitor<F> {
type Break = E;

fn post_visit_function_arg(&mut self, expr: &mut FunctionArgExpr) -> ControlFlow<Self::Break> {
self.0(expr)
}
}

pub fn visit_function_arguments<V, E, F>(v: &V, f: F) -> ControlFlow<E>
where
V: Visit,
F: FnMut(&FunctionArgExpr) -> ControlFlow<E>,
{
let mut visitor = FunctionArgExprVisitor(f);
v.visit(&mut visitor)?;
ControlFlow::Continue(())
}
pub fn visit_function_arguments_mut<V, E, F>(v: &mut V, f: F) -> ControlFlow<E>
where
V: VisitMut,
F: FnMut(&mut FunctionArgExpr) -> ControlFlow<E>,
{
v.visit(&mut FunctionArgExprVisitor(f))?;
ControlFlow::Continue(())
}

struct ExprVisitor<F>(F);

impl<E, F: FnMut(&Expr) -> ControlFlow<E>> Visitor for ExprVisitor<F> {
Expand Down Expand Up @@ -432,7 +486,6 @@ where
v.visit(&mut visitor)?;
ControlFlow::Continue(())
}

/// Invokes the provided closure iteratively with a mutable reference to all expressions
/// present in `v`.
///
Expand Down

0 comments on commit a879adf

Please sign in to comment.