Experiment with AST paths

bench.conf

@@ -0,0 +1,5 @@
+BENCH_PCKG=./go/vt/vtgate/planbuilder
+BENCH_REGEX=Benchmark
+BENCH_DIR=../benchmarks
+BENCH_TIME=2s
+BENCH_COUNT=6

go/tools/asthelpergen/ast_path_gen.go

@@ -0,0 +1,271 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package asthelpergen
+
+import (
+	"fmt"
+	"go/types"
+
+	"github.com/dave/jennifer/jen"
+)
+
+type (
+	pathGen struct {
+		file      *jen.File
+		steps     []step
+		ifaceName string
+	}
+	step struct {
+		container types.Type // the type of the container
+		typ       types.Type // the type of the field
+		name      string     // the name of the field
+		slice     bool       // whether the field is a slice
+	}
+)
+
+const sliceMarker = " -SLICE- "
+
+var _ generator = (*pathGen)(nil)
+
+func newPathGen(pkgname, ifaceName string) *pathGen {
+	file := jen.NewFile(pkgname)
+	file.HeaderComment(licenseFileHeader)
+	file.HeaderComment("Code generated by ASTHelperGen. DO NOT EDIT.")
+
+	return &pathGen{
+		file:      file,
+		ifaceName: ifaceName,
+	}
+}
+
+func (s step) AsEnum() string {
+	if s.name == sliceMarker {
+		return printableTypeName(s.container)
+	}
+	return printableTypeName(s.container) + s.name
+}
+
+func (p *pathGen) genFile(spi generatorSPI) (string, *jen.File) {
+	p.file.ImportName("fmt", "fmt")
+
+	// Declare the ASTStep type with underlying type uint16
+	p.file.Add(jen.Type().Id("ASTStep").Uint16())
+
+	// Add the const block
+	p.file.Add(p.buildConstWithEnum())
+
+	// Add the ASTStep#DebugString() method to the file
+	p.file.Add(p.debugString())
+
+	// Add the generated WalkASTPath method
+	p.file.Add(p.generateWalkASTPath())
+
+	return "ast_path.go", p.file
+}
+
+func (p *pathGen) interfaceMethod(t types.Type, iface *types.Interface, spi generatorSPI) error {
+	return nil
+}
+
+func (p *pathGen) structMethod(t types.Type, strct *types.Struct, spi generatorSPI) error {
+	p.addStructFields(t, strct, spi)
+	return nil
+}
+
+func (p *pathGen) ptrToStructMethod(t types.Type, strct *types.Struct, spi generatorSPI) error {
+	p.addStructFields(t, strct, spi)
+	return nil
+}
+
+func (p *pathGen) addStep(
+	container types.Type, // the name of the container type
+	typ types.Type, // the type of the field
+	name string, // the name of the field
+	slice bool, // whether the field is a slice
+) {
+	s := step{
+		container: container,
+		name:      name,
+		typ:       typ,
+		slice:     slice,
+	}
+	stepName := s.AsEnum()
+	fmt.Println("Adding step:", stepName)
+	p.steps = append(p.steps, s)
+
+}
+
+func (p *pathGen) addStructFields(t types.Type, strct *types.Struct, spi generatorSPI) {
+	for i := 0; i < strct.NumFields(); i++ {
+		field := strct.Field(i)
+		// Check if the field type implements the interface
+		if types.Implements(field.Type(), spi.iface()) {
+			p.addStep(t, field.Type(), field.Name(), false)
+			continue
+		}
+		// Check if the field type is a slice
+		slice, isSlice := field.Type().(*types.Slice)
+		if isSlice {
+			// Check if the slice type implements the interface
+			if types.Implements(slice, spi.iface()) {
+				p.addStep(t, slice, field.Name(), true)
+			} else if types.Implements(slice.Elem(), spi.iface()) {
+				// Check if the element type of the slice implements the interface
+				p.addStep(t, slice.Elem(), field.Name(), true)
+			}
+		}
+	}
+}
+
+func (p *pathGen) ptrToBasicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error {
+	return nil
+}
+
+func (p *pathGen) sliceMethod(t types.Type, slice *types.Slice, spi generatorSPI) error {
+	//elemType := slice.Elem()
+	//if types.Implements(elemType, spi.iface()) {
+	//	p.addStep(t, elemType, sliceMarker, true)
+	//}
+	//
+	return nil
+}
+
+func (p *pathGen) basicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error {
+	return nil
+}
+
+func (p *pathGen) debugString() *jen.Statement {
+	var switchCases []jen.Code
+
+	for _, step := range p.steps {
+		stepName := step.AsEnum()
+
+		// Generate the debug string using the helper function
+		var debugStr string
+		if step.name == sliceMarker {
+			debugStr = fmt.Sprintf("(%s)[]", types.TypeString(step.container, noQualifier))
+		} else {
+			debugStr = fmt.Sprintf("(%s).%s", types.TypeString(step.container, noQualifier), step.name)
+		}
+
+		if !step.slice {
+			switchCases = append(switchCases, jen.Case(jen.Id(stepName)).Block(
+				jen.Return(jen.Lit(debugStr)),
+			))
+			continue
+		}
+
+		switchCases = append(switchCases, jen.Case(jen.Id(stepName+"8")).Block(
+			jen.Return(jen.Lit(debugStr+"8")),
+		))
+		switchCases = append(switchCases, jen.Case(jen.Id(stepName+"32")).Block(
+			jen.Return(jen.Lit(debugStr+"32")),
+		))
+
+	}
+
+	debugStringMethod := jen.Func().Params(jen.Id("s").Id("ASTStep")).Id("DebugString").Params().String().Block(
+		jen.Switch(jen.Id("s")).Block(switchCases...),
+		jen.Panic(jen.Lit("unknown ASTStep")),
+	)
+	return debugStringMethod
+}
+
+func (p *pathGen) buildConstWithEnum() *jen.Statement {
+	// Create the const block with all step constants
+	var constDefs []jen.Code
+	addStep := func(step string) {
+		if constDefs == nil {
+			// Use iota for the first constant
+			constDefs = append(constDefs, jen.Id(step).Id("ASTStep").Op("=").Id("iota"))
+			return
+		}
+
+		constDefs = append(constDefs, jen.Id(step))
+	}
+	for _, step := range p.steps {
+		stepName := step.AsEnum()
+		if step.slice {
+			addStep(stepName + "8")
+			addStep(stepName + "32")
+			continue
+		}
+
+		addStep(stepName)
+	}
+
+	constBlock := jen.Const().Defs(constDefs...)
+	return constBlock
+}
+
+func (p *pathGen) generateWalkASTPath() *jen.Statement {
+	method := jen.Func().Id("WalkASTPath").Params(
+		jen.Id("node").Id(p.ifaceName),
+		jen.Id("path").Id("ASTPath"),
+	).Id(p.ifaceName).Block(
+		jen.If(jen.Id("path").Op("==").Id(`""`).Block(
+			jen.Return(jen.Id("node")))),
+		jen.Id("step").Op(":=").Qual("encoding/binary", "BigEndian").Dot("Uint16").Call(jen.Index().Byte().Parens(jen.Id("path[:2]"))),
+		jen.Id("path").Op("=").Id("path[2:]"),
+
+		jen.Switch(jen.Id("ASTStep").Parens(jen.Id("step"))).Block(p.generateWalkCases()...),
+		jen.Return(jen.Nil()), // Fallback return
+	)
+	return method
+}
+
+func (p *pathGen) generateWalkCases() []jen.Code {
+	var cases []jen.Code
+
+	for _, step := range p.steps {
+		stepName := step.AsEnum()
+
+		if !step.slice {
+			// return WalkASTPath(node.(*RefContainer).ASTType, path)
+			t := types.TypeString(step.container, noQualifier)
+			n := jen.Id("node").Dot(fmt.Sprintf("(%s)", t)).Dot(step.name)
+
+			cases = append(cases, jen.Case(jen.Id(stepName)).Block(
+				//jen.Id("node").Op("=").Id("node").Dot(step.name),
+				jen.Return(jen.Id("WalkASTPath").Call(n, jen.Id("path"))),
+			))
+			continue
+		}
+
+		var assignNode jen.Code
+		t := types.TypeString(step.container, noQualifier)
+		if step.name == sliceMarker {
+			assignNode = jen.Id("node").Dot(fmt.Sprintf("(%s)", t)).Index(jen.Id("idx"))
+		} else {
+			assignNode = jen.Id("node").Dot(fmt.Sprintf("(%s)", t)).Dot(step.name).Index(jen.Id("idx"))
+		}
+
+		cases = append(cases, jen.Case(jen.Id(stepName+"8")).Block(
+			jen.Id("idx").Op(":=").Id("path[0]"),
+			jen.Id("path").Op("=").Id("path[1:]"),
+			jen.Return(jen.Id("WalkASTPath").Call(assignNode, jen.Id("path"))),
+		))
+
+		cases = append(cases, jen.Case(jen.Id(stepName+"32")).Block(
+			jen.Id("idx").Op(":=").Qual("encoding/binary", "BigEndian").Dot("Uint32").Call(jen.Index().Byte().Parens(jen.Id("path[:2]"))),
+			jen.Id("path").Op("=").Id("path[4:]"),
+			jen.Return(jen.Id("WalkASTPath").Call(assignNode, jen.Id("path"))),
+		))
+	}
+
+	return cases
+}

go/tools/asthelpergen/asthelpergen.go

@@ -32,7 +32,7 @@ import (
 	"vitess.io/vitess/go/tools/codegen"
 )
 
-const licenseFileHeader = `Copyright 2023 The Vitess Authors.
+const licenseFileHeader = `Copyright 2025 The Vitess Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -51,10 +51,10 @@ type (
 		addType(t types.Type)
 		scope() *types.Scope
 		findImplementations(iff *types.Interface, impl func(types.Type) error) error
-		iface() *types.Interface
+		iface() *types.Interface // the root interface that all nodes are expected to implement
 	}
 	generator interface {
-		genFile() (string, *jen.File)
+		genFile(generatorSPI) (string, *jen.File)
 		interfaceMethod(t types.Type, iface *types.Interface, spi generatorSPI) error
 		structMethod(t types.Type, strct *types.Struct, spi generatorSPI) error
 		ptrToStructMethod(t types.Type, strct *types.Struct, spi generatorSPI) error
@@ -215,11 +215,13 @@ func GenerateASTHelpers(options *Options) (map[string]*jen.File, error) {
 
 	nt := tt.Type().(*types.Named)
 	pName := nt.Obj().Pkg().Name()
+	ifaceName := types.TypeString(nt, noQualifier)
 	generator := newGenerator(loaded[0].Module, loaded[0].TypesSizes, nt,
 		newEqualsGen(pName, &options.Equals),
 		newCloneGen(pName, &options.Clone),
+		newPathGen(pName, ifaceName),
 		newVisitGen(pName),
-		newRewriterGen(pName, types.TypeString(nt, noQualifier), exprInterface),
+		newRewriterGen(pName, ifaceName, exprInterface),
 		newCOWGen(pName, nt),
 	)
 
@@ -307,7 +309,7 @@ func (gen *astHelperGen) createFiles() map[string]*jen.File {
 
 	result := map[string]*jen.File{}
 	for _, g := range gen.gens {
-		fName, jenFile := g.genFile()
+		fName, jenFile := g.genFile(gen)
 		result[fName] = jenFile
 	}
 	return result

go/tools/asthelpergen/asthelpergen_test.go

@@ -21,6 +21,8 @@ import (
 	"strings"
 	"testing"
 
+	"vitess.io/vitess/go/tools/codegen"
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -45,3 +47,20 @@ func TestFullGeneration(t *testing.T) {
 		require.False(t, applyIdx == 0 && cloneIdx == 0, "file doesn't contain expected contents")
 	}
 }
+
+func TestRecreateAllFiles(t *testing.T) {
+	// t.Skip("This test recreates all files in the integration directory. It should only be run when the ASTHelperGen code has changed.")
+	result, err := GenerateASTHelpers(&Options{
+		Packages:      []string{"./integration/..."},
+		RootInterface: "vitess.io/vitess/go/tools/asthelpergen/integration.AST",
+		Clone: CloneOptions{
+			Exclude: []string{"*NoCloneType"},
+		},
+	})
+	require.NoError(t, err)
+
+	for fullPath, file := range result {
+		err := codegen.SaveJenFile(fullPath, file)
+		require.NoError(t, err)
+	}
+}

go/tools/asthelpergen/clone_gen.go

@@ -56,7 +56,7 @@ func (c *cloneGen) addFunc(name string, code *jen.Statement) {
 	c.file.Add(code)
 }
 
-func (c *cloneGen) genFile() (string, *jen.File) {
+func (c *cloneGen) genFile(generatorSPI) (string, *jen.File) {
 	return "ast_clone.go", c.file
 }
 

go/tools/asthelpergen/copy_on_rewrite_gen.go

@@ -44,7 +44,7 @@ func (c *cowGen) addFunc(code *jen.Statement) {
 	c.file.Add(code)
 }
 
-func (c *cowGen) genFile() (string, *jen.File) {
+func (c *cowGen) genFile(generatorSPI) (string, *jen.File) {
 	return "ast_copy_on_rewrite.go", c.file
 }
 

go/tools/asthelpergen/equals_gen.go

@@ -62,7 +62,7 @@ func (e *equalsGen) customComparatorField(t types.Type) string {
 	return printableTypeName(t) + "_"
 }
 
-func (e *equalsGen) genFile() (string, *jen.File) {
+func (e *equalsGen) genFile(generatorSPI) (string, *jen.File) {
 	e.file.Type().Id(Comparator).StructFunc(func(g *jen.Group) {
 		for tname, t := range e.comparators {
 			if t == nil {
@@ -303,6 +303,4 @@ func (e *equalsGen) sliceMethod(t types.Type, slice *types.Slice, spi generatorS
 	return nil
 }
 
-func (e *equalsGen) basicMethod(types.Type, *types.Basic, generatorSPI) error {
-	return nil
-}
+func (*equalsGen) basicMethod(types.Type, *types.Basic, generatorSPI) error { return nil }