github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/expr.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package colexec
    12  
    13  import (
    14  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    15  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    16  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    17  )
    18  
    19  // Remove unused warning.
    20  var _ = findIVarsInRange
    21  
    22  // findIVarsInRange searches Expr for presence of tree.IndexedVars with indices
    23  // in range [start, end). It returns a slice containing all such indices.
    24  func findIVarsInRange(expr execinfrapb.Expression, start int, end int) ([]uint32, error) {
    25  	res := make([]uint32, 0)
    26  	if start >= end {
    27  		return res, nil
    28  	}
    29  	var exprToWalk tree.Expr
    30  	if expr.LocalExpr != nil {
    31  		exprToWalk = expr.LocalExpr
    32  	} else {
    33  		e, err := parser.ParseExpr(expr.Expr)
    34  		if err != nil {
    35  			return nil, err
    36  		}
    37  		exprToWalk = e
    38  	}
    39  	visitor := ivarExpressionVisitor{ivarSeen: make([]bool, end)}
    40  	_, _ = tree.WalkExpr(visitor, exprToWalk)
    41  	for i := start; i < end; i++ {
    42  		if visitor.ivarSeen[i] {
    43  			res = append(res, uint32(i))
    44  		}
    45  	}
    46  	return res, nil
    47  }
    48  
    49  type ivarExpressionVisitor struct {
    50  	ivarSeen []bool
    51  }
    52  
    53  var _ tree.Visitor = &ivarExpressionVisitor{}
    54  
    55  // VisitPre is a part of tree.Visitor interface.
    56  func (i ivarExpressionVisitor) VisitPre(expr tree.Expr) (bool, tree.Expr) {
    57  	switch e := expr.(type) {
    58  	case *tree.IndexedVar:
    59  		if e.Idx < len(i.ivarSeen) {
    60  			i.ivarSeen[e.Idx] = true
    61  		}
    62  		return false, expr
    63  	default:
    64  		return true, expr
    65  	}
    66  }
    67  
    68  // VisitPost is a part of tree.Visitor interface.
    69  func (i ivarExpressionVisitor) VisitPost(expr tree.Expr) tree.Expr { return expr }