Merge #49939

49939: execgen: introduce // execgen:template directive for bools r=jordanlewis a=jordanlewis

Based on #49728.

This PR adds the // execgen:template<arg1, arg2> syntax for
concrete boolean template variables.

For example:

```
// execgen:inline
// execgen:template<templateParam>
func a (runtimeParam int, templateParam bool) int {
    if templateParam {
        return runtimeParam + 1
    } else {
        return runtimeParam - 1
    }
}

func main() {
    x := a(0, true)
    y := a(0, false)
    fmt.Println(x, y)
}
```

This code will print `1, -1` when eventually executed, but the
templateParam boolean will be evaluated at code generation time and not
appear in the generated code at all. Two variants of `a()` will be
inlined into their callsites - one with templateParam true, and one with
templateParam false.

Release note: None

Co-authored-by: Jordan Lewis <[email protected]>

pkg/sql/colexec/execgen/cmd/execgen/main.go

@@ -47,6 +47,7 @@ type execgenTool struct {
 
 	// cmdLine stores the set of flags used to invoke the Execgen tool.
 	cmdLine *flag.FlagSet
+	verbose bool
 }
 
 // generator is a func that, given an input file's contents as a string,
@@ -74,6 +75,7 @@ func (g *execgenTool) run(args ...string) bool {
 	g.cmdLine.SetOutput(g.stdErr)
 	g.cmdLine.Usage = g.usage
 	g.cmdLine.BoolVar(&g.useGoFmt, "useGoFmt", true, "run go fmt on generated code")
+	g.cmdLine.BoolVar(&g.verbose, "verbose", false, "print out debug information to stderr")
 	g.cmdLine.BoolVar(&printDeps, "M", false, "print the dependency list")
 	err := g.cmdLine.Parse(args)
 	if err != nil {
@@ -126,11 +128,16 @@ func (g *execgenTool) generate(path string, entry entry) error {
 		if err != nil {
 			return err
 		}
-		// Inline functions with // execgen:inline.
-		inputFileContents, err = execgen.InlineFuncs(string(inputFileBytes))
+		inputFileContents, err = execgen.Generate(string(inputFileBytes))
 		if err != nil {
 			return err
 		}
+		if g.verbose {
+			fmt.Fprintln(os.Stderr, "generated code before text/template runs")
+			fmt.Fprintln(os.Stderr, "-----------------------------------")
+			fmt.Fprintln(os.Stderr, inputFileContents)
+			fmt.Fprintln(os.Stderr, "-----------------------------------")
+		}
 	}
 
 	err = entry.fn(inputFileContents, &buf)

pkg/sql/colexec/execgen/datadriven_test.go

@@ -11,26 +11,33 @@
 package execgen
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/cockroachdb/datadriven"
+	"github.com/dave/dst/decorator"
 )
 
 // Walk walks path for datadriven files and calls RunTest on them.
 func TestExecgen(t *testing.T) {
 	datadriven.Walk(t, "testdata", func(t *testing.T, path string) {
 		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
+			f, err := decorator.Parse(d.Input)
+			if err != nil {
+				t.Fatal(err)
+			}
 			switch d.Cmd {
 			case "inline":
-				s, err := InlineFuncs(d.Input)
-				if err != nil {
-					t.Fatal(err)
-				}
-				return s
+				inlineFuncs(f)
+			case "template":
+				expandTemplates(f)
 			default:
 				t.Fatalf("unknown command: %s", d.Cmd)
 				return ""
 			}
+			var sb strings.Builder
+			_ = decorator.Fprint(&sb, f)
+			return sb.String()
 		})
 	})
 }

pkg/sql/colexec/execgen/execgen.go

@@ -0,0 +1,39 @@
+// Copyright 2020 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package execgen
+
+import (
+	"go/parser"
+	"go/token"
+	"strings"
+
+	"github.com/dave/dst/decorator"
+)
+
+// Generate transforms the string contents of an input execgen template by
+// processing all supported // execgen annotations.
+func Generate(inputFileContents string) (string, error) {
+	f, err := decorator.ParseFile(token.NewFileSet(), "", inputFileContents, parser.ParseComments)
+	if err != nil {
+		return "", err
+	}
+
+	// Generate template variants: // execgen:template
+	expandTemplates(f)
+
+	// Inline functions: // execgen:inline
+	inlineFuncs(f)
+
+	// Produce output string.
+	var sb strings.Builder
+	_ = decorator.Fprint(&sb, f)
+	return sb.String(), nil
+}