github.com/dolthub/go-mysql-server@v0.18.0/sql/transform/walk.go (about)

     1  // Copyright 2020-2021 Dolthub, Inc.
     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  
    15  package transform
    16  
    17  import (
    18  	"github.com/dolthub/go-mysql-server/sql"
    19  )
    20  
    21  // Visitor visits nodes in the plan.
    22  type Visitor interface {
    23  	// Visit method is invoked for each node encountered by Walk.
    24  	// If the result Visitor is not nil, Walk visits each of the children
    25  	// of the node with that visitor, followed by a call of Visit(nil)
    26  	// to the returned visitor.
    27  	Visit(node sql.Node) Visitor
    28  }
    29  
    30  // Walk traverses the plan tree in depth-first order. It starts by calling v.Visit(node); node must not be nil. If the
    31  // visitor returned by  v.Visit(node) is not nil, Walk is invoked recursively with the returned visitor for each
    32  // children of the node, followed by a call of v.Visit(nil) to the returned visitor. If v.Visit(node) returns non-nil,
    33  // then all children are walked, even if one of them returns nil for v.Visit().
    34  func Walk(v Visitor, node sql.Node) {
    35  	if v = v.Visit(node); v == nil {
    36  		return
    37  	}
    38  
    39  	for _, child := range node.Children() {
    40  		Walk(v, child)
    41  	}
    42  
    43  	v.Visit(nil)
    44  }
    45  
    46  type inspector func(sql.Node) bool
    47  
    48  func (f inspector) Visit(node sql.Node) Visitor {
    49  	if f(node) {
    50  		return f
    51  	}
    52  	return nil
    53  }
    54  
    55  // Inspect traverses the plan in depth-first order: It starts by calling
    56  // f(node); node must not be nil. If f returns true, Inspect invokes f
    57  // recursively for each of the children of node, followed by a call of
    58  // f(nil).
    59  func Inspect(node sql.Node, f func(sql.Node) bool) {
    60  	Walk(inspector(f), node)
    61  }
    62  
    63  // WalkExpressions traverses the plan and calls sql.Walk on any expression it finds.
    64  func WalkExpressions(v sql.Visitor, node sql.Node) {
    65  	Inspect(node, func(node sql.Node) bool {
    66  		if n, ok := node.(sql.Expressioner); ok {
    67  			for _, e := range n.Expressions() {
    68  				sql.Walk(v, e)
    69  			}
    70  		}
    71  		return true
    72  	})
    73  }
    74  
    75  // WalkExpressionsWithNode traverses the plan and calls sql.WalkWithNode on any expression it finds.
    76  func WalkExpressionsWithNode(v sql.NodeVisitor, n sql.Node) {
    77  	Inspect(n, func(n sql.Node) bool {
    78  		if expressioner, ok := n.(sql.Expressioner); ok {
    79  			for _, e := range expressioner.Expressions() {
    80  				sql.WalkWithNode(v, n, e)
    81  			}
    82  		}
    83  		return true
    84  	})
    85  }
    86  
    87  // InspectExpressions traverses the plan and calls sql.Inspect on any
    88  // expression it finds.
    89  func InspectExpressions(node sql.Node, f func(sql.Expression) bool) {
    90  	WalkExpressions(exprInspector(f), node)
    91  }
    92  
    93  type exprInspector func(sql.Expression) bool
    94  
    95  func (f exprInspector) Visit(e sql.Expression) sql.Visitor {
    96  	if f(e) {
    97  		return f
    98  	}
    99  	return nil
   100  }
   101  
   102  // InspectExpressionsWithNode traverses the plan and calls sql.Inspect on any expression it finds.
   103  func InspectExpressionsWithNode(node sql.Node, f func(sql.Node, sql.Expression) bool) {
   104  	WalkExpressionsWithNode(exprWithNodeInspector(f), node)
   105  }
   106  
   107  type exprWithNodeInspector func(sql.Node, sql.Expression) bool
   108  
   109  func (f exprWithNodeInspector) Visit(n sql.Node, e sql.Expression) sql.NodeVisitor {
   110  	if f(n, e) {
   111  		return f
   112  	}
   113  	return nil
   114  }