github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/sql/sem/tree/evalgen/expr.go (about)

     1  // Copyright 2022 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 main
    12  
    13  import (
    14  	"go/ast"
    15  	"html/template"
    16  	"os"
    17  	"regexp"
    18  
    19  	"golang.org/x/tools/go/ast/inspector"
    20  )
    21  
    22  func generateExprEval(fileName string, files []*ast.File) (err error) {
    23  	ins := inspector.New(files)
    24  	exprs := findExprTypes(ins)
    25  	var tmpl = template.Must(template.Must(template.New("name").
    26  		Funcs(template.FuncMap{
    27  			"isDatum": func(name string) bool {
    28  				return isDatumRE.MatchString(name)
    29  			},
    30  			"isPtr": func(name string) bool {
    31  				switch name {
    32  				case "dNull", "UnqualifiedStar":
    33  					return false
    34  				default:
    35  					return true
    36  				}
    37  			},
    38  		}).Parse("{{ if isPtr . }}*{{ end }}{{ . }}")).
    39  		New("t").
    40  		Parse(visitorTemplate))
    41  
    42  	return writeFile(fileName, func(f *os.File) error {
    43  		return tmpl.Execute(f, exprs.ordered())
    44  	})
    45  }
    46  
    47  var isDatumRE = regexp.MustCompile("^D[A-Z]|^dNull$")
    48  
    49  const visitorTemplate = header +
    50  	`
    51  // ExprEvaluator is used to evaluate TypedExpr expressions.
    52  type ExprEvaluator interface {
    53  {{- range . }}{{ if isDatum . }}{{ else }}
    54  	Eval{{.}}(context.Context, {{ template "name" . }}) (Datum, error)
    55  {{- end}}{{ end }}
    56  }
    57  
    58  {{ range . }}
    59  // Eval is part of the TypedExpr interface.
    60  func (node {{ template "name" . }}) Eval(ctx context.Context, v ExprEvaluator) (Datum, error) {
    61  	{{ if isDatum . }}return node, nil{{ else }}return v.Eval{{.}}(ctx, node){{ end }}
    62  }
    63  {{ end }}
    64  `
    65  
    66  // findExprTypes finds all types in the files which have both the Walk method
    67  // and the ResolvedType method as we know TypedExprs to both have. It returns
    68  // a sorted list.
    69  func findExprTypes(ins *inspector.Inspector) (exprs stringSet) {
    70  	hasWalk := findTypesWithMethod(ins, isWalkMethod)
    71  	hasResolvedType := findTypesWithMethod(ins, isResolvedTypeMethod)
    72  	hasResolvedType.addAll(
    73  		findStructTypesWithEmbeddedType(ins, hasWalk.contains, "typeAnnotation"),
    74  	)
    75  	return hasWalk.intersection(hasResolvedType)
    76  }
    77  
    78  func isResolvedTypeMethod(fd *ast.FuncDecl) bool {
    79  	return fd.Name.Name == "ResolvedType" &&
    80  		fd.Recv != nil &&
    81  		len(fd.Type.Params.List) == 0 &&
    82  		len(fd.Type.Results.List) == 1 &&
    83  		isStarTypesT(fd.Type.Results.List[0].Type)
    84  }
    85  
    86  func isStarTypesT(expr ast.Expr) bool {
    87  	st, ok := expr.(*ast.StarExpr)
    88  	if !ok {
    89  		return false
    90  	}
    91  	sel, ok := st.X.(*ast.SelectorExpr)
    92  	if !ok {
    93  		return false
    94  	}
    95  	return isIdentWithName(sel.X, "types") && sel.Sel.Name == "T"
    96  }
    97  
    98  func isWalkMethod(fd *ast.FuncDecl) bool {
    99  	return fd.Name.Name == "Walk" &&
   100  		fd.Recv != nil &&
   101  		len(fd.Type.Params.List) == 1 &&
   102  		isIdentWithName(fd.Type.Params.List[0].Type, "Visitor") &&
   103  		len(fd.Type.Results.List) == 1 &&
   104  		isIdentWithName(fd.Type.Results.List[0].Type, "Expr")
   105  }
   106  
   107  func isIdentWithName(expr ast.Expr, name string) bool {
   108  	id, ok := expr.(*ast.Ident)
   109  	if !ok {
   110  		return false
   111  	}
   112  	return id.Name == name
   113  }