databend_common_ast/
visitor.rs

1// Copyright 2021 Datafuse Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use derive_visitor::DriveMut;
16use derive_visitor::VisitorMut;
17
18use crate::ast::Expr;
19use crate::ast::Identifier;
20use crate::ast::Statement;
21
22/// used in bendsql
23#[derive(VisitorMut)]
24#[visitor(Expr(enter), Identifier(enter))]
25pub struct StatementReplacer<F: FnMut(&mut Expr), G: FnMut(&mut Identifier)> {
26    replace_expr: F,
27    replace_ident: G,
28}
29
30impl<F: FnMut(&mut Expr), G: FnMut(&mut Identifier)> StatementReplacer<F, G> {
31    pub fn new(replace_expr: F, replace_ident: G) -> Self {
32        Self {
33            replace_expr,
34            replace_ident,
35        }
36    }
37
38    fn enter_expr(&mut self, expr: &mut Expr) {
39        (self.replace_expr)(expr);
40    }
41
42    fn enter_identifier(&mut self, ident: &mut Identifier) {
43        (self.replace_ident)(ident);
44    }
45
46    pub fn visit(&mut self, stmt: &mut Statement) {
47        stmt.drive_mut(self);
48    }
49}