github.com/vedadiyan/sqlparser@v1.0.0/pkg/sqlparser/cow.go (about)

     1  /*
     2  Copyright 2023 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package sqlparser
    18  
    19  // CopyOnRewrite traverses a syntax tree recursively, starting with root,
    20  // and calling pre and post for each node as described below.
    21  // Rewrite returns a syntax tree, where some nodes can be shared with the
    22  // original syntax tree.
    23  //
    24  // If pre is not nil, it is called for each node before the node's
    25  // children are traversed (pre-order). If pre returns false, no
    26  // children are traversed, but post is still called for that node.
    27  //
    28  // If post is not nil, and a prior call of pre didn't return false,
    29  // post is called for each node after its children are traversed
    30  // (post-order).
    31  //
    32  // In the post call, the cursor can be used to abort the current
    33  // traversal altogether.
    34  //
    35  // Also in the post call, a user can replace the current node.
    36  // When a node is replaced, all the ancestors of the node are cloned,
    37  // so that the original syntax tree remains untouched
    38  //
    39  // The `cloned` function will be called for all nodes that are cloned
    40  // or replaced, to give the user a chance to copy any metadata that
    41  // needs copying.
    42  //
    43  // Only fields that refer to AST nodes are considered children;
    44  // i.e., fields of basic types (strings, []byte, etc.) are ignored.
    45  func CopyOnRewrite(
    46  	node SQLNode,
    47  	pre func(node, parent SQLNode) bool,
    48  	post func(cursor *CopyOnWriteCursor),
    49  	cloned func(before, after SQLNode),
    50  ) SQLNode {
    51  	cow := cow{pre: pre, post: post, cursor: CopyOnWriteCursor{}, cloned: cloned}
    52  	out, _ := cow.copyOnRewriteSQLNode(node, nil)
    53  	return out
    54  }
    55  
    56  // StopTreeWalk aborts the current tree walking. No more nodes will be visited, and the rewriter will exit out early
    57  func (c *CopyOnWriteCursor) StopTreeWalk() {
    58  	c.stop = true
    59  }
    60  
    61  // Node returns the current node we are visiting
    62  func (c *CopyOnWriteCursor) Node() SQLNode {
    63  	return c.node
    64  }
    65  
    66  // Parent returns the parent of the current node.
    67  // Note: This is the parent before any changes have been done - the parent in the output might be a copy of this
    68  func (c *CopyOnWriteCursor) Parent() SQLNode {
    69  	return c.parent
    70  }
    71  
    72  // Replace replaces the current node with the given node.
    73  // Note: If you try to set an invalid type on a field, the field will end up with a nil and no error will be reported.
    74  func (c *CopyOnWriteCursor) Replace(n SQLNode) {
    75  	c.replaced = n
    76  }
    77  
    78  func (c *cow) postVisit(node, parent SQLNode, changed bool) (SQLNode, bool) {
    79  	c.cursor.node = node
    80  	c.cursor.parent = parent
    81  	c.cursor.replaced = nil
    82  	c.post(&c.cursor)
    83  	if c.cursor.replaced != nil {
    84  		if c.cloned != nil {
    85  			c.cloned(node, c.cursor.replaced)
    86  		}
    87  		return c.cursor.replaced, true
    88  	}
    89  	return node, changed
    90  }
    91  
    92  type (
    93  	CopyOnWriteCursor struct {
    94  		node     SQLNode
    95  		parent   SQLNode
    96  		replaced SQLNode
    97  		stop     bool
    98  	}
    99  	cow struct {
   100  		pre    func(node, parent SQLNode) bool
   101  		post   func(cursor *CopyOnWriteCursor)
   102  		cloned func(old, new SQLNode)
   103  		cursor CopyOnWriteCursor
   104  	}
   105  )