vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/operators/rewrite/rewriters.go (about)

     1  /*
     2  Copyright 2022 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 rewrite
    18  
    19  import (
    20  	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/ops"
    21  )
    22  
    23  type (
    24  	Func          func(ops.Operator) (ops.Operator, TreeIdentity, error)
    25  	BreakableFunc func(ops.Operator) (ops.Operator, TreeIdentity, VisitRule, error)
    26  
    27  	// TreeIdentity tracks modifications to node and expression trees.
    28  	// Only return SameTree when it is acceptable to return the original
    29  	// input and discard the returned result as a performance improvement.
    30  	TreeIdentity bool
    31  
    32  	// VisitRule signals to the rewriter if the children of this operator should be visited or not
    33  	VisitRule bool
    34  )
    35  
    36  const (
    37  	SameTree TreeIdentity = false
    38  	NewTree  TreeIdentity = true
    39  
    40  	VisitChildren VisitRule = true
    41  	SkipChildren  VisitRule = false
    42  )
    43  
    44  // Visit allows for the walking of the operator tree. If any error is returned, the walk is aborted
    45  func Visit(root ops.Operator, visitor func(ops.Operator) error) error {
    46  	_, err := TopDown(root, func(op ops.Operator) (ops.Operator, TreeIdentity, VisitRule, error) {
    47  		err := visitor(op)
    48  		if err != nil {
    49  			return nil, SameTree, SkipChildren, err
    50  		}
    51  		return op, SameTree, VisitChildren, nil
    52  	})
    53  	return err
    54  }
    55  
    56  // BottomUp rewrites an operator tree from the bottom up. BottomUp applies a transformation function to
    57  // the given operator tree from the bottom up. Each callback [f] returns a TreeIdentity that is aggregated
    58  // into a final output indicating whether the operator tree was changed.
    59  func BottomUp(root ops.Operator, f Func) (ops.Operator, error) {
    60  	op, _, err := bottomUp(root, f)
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  	return op, nil
    65  }
    66  
    67  // TopDown applies a transformation function to the given operator tree from the bottom up. =
    68  // Each callback [f] returns a TreeIdentity that is aggregated into a final output indicating whether the
    69  // operator tree was changed.
    70  // The callback also returns a VisitRule that signals whether the children of this operator should be visited or not
    71  func TopDown(in ops.Operator, rewriter BreakableFunc) (ops.Operator, error) {
    72  	op, _, err := breakableTopDown(in, rewriter)
    73  	return op, err
    74  }
    75  
    76  func bottomUp(root ops.Operator, rewriter Func) (ops.Operator, TreeIdentity, error) {
    77  	oldInputs := root.Inputs()
    78  	anythingChanged := false
    79  	newInputs := make([]ops.Operator, len(oldInputs))
    80  	for i, operator := range oldInputs {
    81  		in, changed, err := bottomUp(operator, rewriter)
    82  		if err != nil {
    83  			return nil, SameTree, err
    84  		}
    85  		if changed == NewTree {
    86  			anythingChanged = true
    87  		}
    88  		newInputs[i] = in
    89  	}
    90  
    91  	if anythingChanged {
    92  		root = root.Clone(newInputs)
    93  	}
    94  
    95  	newOp, treeIdentity, err := rewriter(root)
    96  	if err != nil {
    97  		return nil, SameTree, err
    98  	}
    99  	if anythingChanged {
   100  		treeIdentity = NewTree
   101  	}
   102  	return newOp, treeIdentity, nil
   103  }
   104  
   105  func breakableTopDown(in ops.Operator, rewriter BreakableFunc) (ops.Operator, TreeIdentity, error) {
   106  	newOp, identity, visit, err := rewriter(in)
   107  	if err != nil || visit == SkipChildren {
   108  		return newOp, identity, err
   109  	}
   110  
   111  	anythingChanged := identity == NewTree
   112  
   113  	oldInputs := newOp.Inputs()
   114  	newInputs := make([]ops.Operator, len(oldInputs))
   115  	for i, oldInput := range oldInputs {
   116  		newInputs[i], identity, err = breakableTopDown(oldInput, rewriter)
   117  		anythingChanged = anythingChanged || identity == NewTree
   118  		if err != nil {
   119  			return nil, SameTree, err
   120  		}
   121  	}
   122  
   123  	if anythingChanged {
   124  		return newOp.Clone(newInputs), NewTree, nil
   125  	}
   126  
   127  	return newOp, SameTree, nil
   128  }