github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/sql/sem/tree/evalgen/op.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  
    18  	"golang.org/x/tools/go/ast/inspector"
    19  )
    20  
    21  const (
    22  	binaryOpsDefinitionsFile = "eval_binary_ops.go"
    23  	unaryOpsDefinitionsFile  = "eval_unary_ops.go"
    24  )
    25  
    26  func generateOpsFile(fileName string, files []*ast.File, byName map[string]*ast.File) (err error) {
    27  	return writeFile(fileName, func(f *os.File) error {
    28  		return opsTemplate.Execute(f, struct {
    29  			UnaryOps, BinaryOps []string
    30  		}{
    31  			UnaryOps: findEvalOps(
    32  				files, byName[unaryOpsDefinitionsFile], isUnaryEvalMethod,
    33  			).ordered(),
    34  			BinaryOps: findEvalOps(
    35  				files, byName[binaryOpsDefinitionsFile], isBinaryEvalMethod,
    36  			).ordered(),
    37  		})
    38  	})
    39  }
    40  
    41  func findEvalOps(
    42  	files []*ast.File, definitionFile *ast.File, existingMethodFilter func(decl *ast.FuncDecl) bool,
    43  ) stringSet {
    44  	types := stringSet{}
    45  	inspector.New([]*ast.File{definitionFile}).
    46  		Preorder([]ast.Node{(*ast.TypeSpec)(nil)}, func(node ast.Node) {
    47  			n := node.(*ast.TypeSpec)
    48  			types.add(n.Name.Name)
    49  		})
    50  	// Find the operations which already define their eval method.
    51  	// These things are syntactic sugar and should not be added to the
    52  	// visitor.
    53  	withExistingMethod := findTypesWithMethod(inspector.New(files), existingMethodFilter)
    54  	types.removeAll(withExistingMethod)
    55  	return types
    56  }
    57  
    58  func isBinaryEvalMethod(fd *ast.FuncDecl) bool {
    59  	return fd.Name.Name == "Eval" &&
    60  		fd.Recv != nil &&
    61  		len(fd.Type.Results.List) == 2 &&
    62  		isIdentWithName(fd.Type.Results.List[0].Type, "Datum") &&
    63  		isIdentWithName(fd.Type.Results.List[1].Type, "error") &&
    64  		((len(fd.Type.Params.List) == 3 &&
    65  			isIdentWithName(fd.Type.Params.List[2].Type, "Datum") &&
    66  			len(fd.Type.Params.List[2].Names) == 2) ||
    67  			(len(fd.Type.Params.List) == 4 &&
    68  				isIdentWithName(fd.Type.Params.List[2].Type, "Datum") &&
    69  				len(fd.Type.Params.List[2].Names) == 1 &&
    70  				isIdentWithName(fd.Type.Params.List[3].Type, "Datum") &&
    71  				len(fd.Type.Params.List[3].Names) == 1)) &&
    72  		isIdentWithName(fd.Type.Params.List[1].Type, "OpEvaluator")
    73  }
    74  
    75  func isUnaryEvalMethod(fd *ast.FuncDecl) bool {
    76  	return fd.Name.Name == "Eval" &&
    77  		fd.Recv != nil &&
    78  		len(fd.Type.Results.List) == 2 &&
    79  		isIdentWithName(fd.Type.Results.List[0].Type, "Datum") &&
    80  		isIdentWithName(fd.Type.Results.List[1].Type, "error") &&
    81  		len(fd.Type.Params.List) == 3 &&
    82  		isIdentWithName(fd.Type.Params.List[1].Type, "OpEvaluator") &&
    83  		isIdentWithName(fd.Type.Params.List[2].Type, "Datum") &&
    84  		len(fd.Type.Params.List[2].Names) == 1
    85  }
    86  
    87  var opsTemplate = template.Must(template.
    88  	New("t").
    89  	Parse(opsTemplateStr))
    90  
    91  const opsTemplateStr = header + `
    92  
    93  // UnaryEvalOp is a unary operation which can be evaluated.
    94  type UnaryEvalOp interface {
    95  	Eval(context.Context, OpEvaluator, Datum) (Datum, error)
    96  }
    97  
    98  // BinaryEvalOp is a binary operation which can be evaluated.
    99  type BinaryEvalOp interface {
   100  	Eval(context.Context, OpEvaluator, Datum, Datum) (Datum, error)
   101  }
   102  
   103  // OpEvaluator is an evaluator for UnaryEvalOp and BinaryEvalOp operations.
   104  type OpEvaluator interface {
   105  	UnaryOpEvaluator
   106  	BinaryOpEvaluator
   107  }
   108  
   109  // UnaryOpEvaluator knows how to evaluate UnaryEvalOps.
   110  type UnaryOpEvaluator interface {
   111  {{- range .UnaryOps }}
   112  	Eval{{.}}(context.Context, *{{ . }}, Datum) (Datum, error)
   113  {{- end}}
   114  }
   115  
   116  // UnaryOpEvaluator knows how to evaluate BinaryEvalOps.
   117  type BinaryOpEvaluator interface {
   118  {{- range .BinaryOps }}
   119  	Eval{{.}}(context.Context, *{{ . }}, Datum, Datum) (Datum, error)
   120  {{- end}}
   121  }
   122  
   123  {{ range .UnaryOps }}
   124  // Eval is part of the UnaryEvalOp interface.
   125  func (op *{{.}}) Eval(ctx context.Context, e OpEvaluator, v Datum) (Datum, error) {
   126  	return e.Eval{{.}}(ctx, op, v)
   127  }
   128  {{ end }}{{ range .BinaryOps }}
   129  // Eval is part of the BinaryEvalOp interface.
   130  func (op *{{.}}) Eval(ctx context.Context, e OpEvaluator, a, b Datum) (Datum, error) {
   131  	return e.Eval{{.}}(ctx, op, a, b)
   132  }
   133  {{ end }}
   134  `