wip

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/asthelpergen_test.go

@@ -49,7 +49,7 @@ func TestFullGeneration(t *testing.T) {
 }
 
 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.")
+	// 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",

go/tools/asthelpergen/integration/paths.go

@@ -0,0 +1,47 @@
+/*
+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 integration
+
+import "encoding/binary"
+
+// This file is a copy of the file go/vt/sqlparser/paths.go
+// We need it here to be able to test the path accumulation of the rewriter
+
+type ASTPath string
+
+func AddStep(path ASTPath, step ASTStep) ASTPath {
+	b := make([]byte, 2)
+	binary.BigEndian.PutUint16(b, uint16(step))
+	return path + ASTPath(b)
+}
+
+func AddStepWithSliceIndex(path ASTPath, step ASTStep, idx int) ASTPath {
+	if idx < 255 {
+		// 2 bytes for step code + 1 byte for index
+		b := make([]byte, 3)
+		binary.BigEndian.PutUint16(b[:2], uint16(step))
+		b[2] = byte(idx)
+		return path + ASTPath(b)
+	}
+
+	// 2 bytes for step code + 4 byte for index
+	b := make([]byte, 6)
+	longStep := step + 1
+	binary.BigEndian.PutUint16(b[:2], uint16(longStep))
+	binary.BigEndian.PutUint32(b[2:], uint32(idx))
+	return path + ASTPath(b)
+}

go/tools/asthelpergen/integration/test_helpers.go

@@ -49,6 +49,7 @@ type Cursor struct {
 	node     AST
 	// marks that the node has been replaced, and the new node should be visited
 	revisit bool
+	current ASTPath
 }
 
 // Node returns the current Node.

go/tools/asthelpergen/integration/types.go

@@ -160,8 +160,9 @@ func (r *NoCloneType) String() string {
 type Visit func(node AST) (bool, error)
 
 type application struct {
-	pre, post ApplyFunc
-	cur       Cursor
+	pre, post    ApplyFunc
+	cur          Cursor
+	collectPaths bool
 }
 
 var Equals = &Comparator{}

go/tools/asthelpergen/rewrite_gen.go

@@ -102,18 +102,30 @@ func (r *rewriteGen) interfaceMethod(t types.Type, iface *types.Interface, spi g
 	return nil
 }
 
-func (r *rewriteGen) structMethod(t types.Type, strct *types.Struct, spi generatorSPI) error {
-	if !shouldAdd(t, spi.iface()) {
-		return nil
+func (r *rewriteGen) visitStructFields(t types.Type, strct *types.Struct, spi generatorSPI, fail bool) (stmts []jen.Code) {
+	fields := r.rewriteAllStructFields(t, strct, spi, fail)
+	stmts = append(stmts, r.executePre(t))
+	hasFields := len(fields) > 0
+	if hasFields {
+		stmts = append(stmts, jen.Var().Id("path").Id("ASTPath"))
 	}
-	fields := r.rewriteAllStructFields(t, strct, spi, true)
-
-	stmts := []jen.Code{r.executePre(t)}
 	stmts = append(stmts, fields...)
+	if hasFields {
+		stmts = append(stmts, jen.If(jen.Id("a.collectPaths").Block(
+			jen.Id("a.cur.current").Op("=").Id("path"),
+		)))
+	}
 	stmts = append(stmts, executePost(len(fields) > 0))
 	stmts = append(stmts, returnTrue())
+	return
+}
 
-	r.rewriteFunc(t, stmts)
+func (r *rewriteGen) structMethod(t types.Type, strct *types.Struct, spi generatorSPI) error {
+	if !shouldAdd(t, spi.iface()) {
+		return nil
+	}
+
+	r.rewriteFunc(t, r.visitStructFields(t, strct, spi, true))
 
 	return nil
 }
@@ -127,17 +139,7 @@ func (r *rewriteGen) ptrToStructMethod(t types.Type, strct *types.Struct, spi ge
 		if node == nil { return nil }
 	*/
 	stmts := []jen.Code{jen.If(jen.Id("node == nil").Block(returnTrue()))}
-
-	/*
-		if !pre(&cur) {
-			return nil
-		}
-	*/
-	stmts = append(stmts, r.executePre(t))
-	fields := r.rewriteAllStructFields(t, strct, spi, false)
-	stmts = append(stmts, fields...)
-	stmts = append(stmts, executePost(len(fields) > 0))
-	stmts = append(stmts, returnTrue())
+	stmts = append(stmts, r.visitStructFields(t, strct, spi, false)...)
 
 	r.rewriteFunc(t, stmts)
 
@@ -300,7 +302,8 @@ func (r *rewriteGen) rewriteAllStructFields(t types.Type, strct *types.Struct, s
 		field := strct.Field(i)
 		if types.Implements(field.Type(), spi.iface()) {
 			spi.addType(field.Type())
-			output = append(output, r.rewriteChild(t, field.Type(), field.Name(), jen.Id("node").Dot(field.Name()), jen.Dot(field.Name()), fail))
+			rewriteLines := r.rewriteChild(t, field.Type(), field.Name(), jen.Id("node").Dot(field.Name()), jen.Dot(field.Name()), fail, i)
+			output = append(output, rewriteLines...)
 			continue
 		}
 		slice, isSlice := field.Type().(*types.Slice)
@@ -323,7 +326,7 @@ func failReplacer(t types.Type, f string) *jen.Statement {
 	return jen.Panic(jen.Lit(fmt.Sprintf("[BUG] tried to replace '%s' on '%s'", f, typeString)))
 }
 
-func (r *rewriteGen) rewriteChild(t, field types.Type, fieldName string, param jen.Code, replace jen.Code, fail bool) jen.Code {
+func (r *rewriteGen) rewriteChild(t, field types.Type, fieldName string, param jen.Code, replace jen.Code, fail bool, offset int) []jen.Code {
 	/*
 		if errF := rewriteAST(node, node.ASTType, func(newNode, parent AST) {
 			parent.(*RefContainer).ASTType = newNode.(AST)
@@ -359,7 +362,22 @@ func (r *rewriteGen) rewriteChild(t, field types.Type, fieldName string, param j
 			param,
 			funcBlock).Block(returnFalse()))
 
-	return rewriteField
+	block := []jen.Code{
+		jen.Id("a.cur.current").Op("=").Id("AddStep").Call(jen.Id("path"), jen.Id(printableTypeName(t)+fieldName)),
+	}
+
+	if offset == 0 {
+		// if we are on the first field, we'll save the path
+		savePath := jen.Id("path").Op("=").Id("a.cur.current")
+		block = append([]jen.Code{savePath}, block...)
+	}
+
+	ifCollectPath := jen.If(jen.Id("a.collectPaths")).Block(block...)
+
+	return []jen.Code{
+		ifCollectPath,
+		rewriteField,
+	}
 }
 
 func (r *rewriteGen) rewriteChildSlice(t, field types.Type, fieldName string, param jen.Code, replace jen.Code, fail bool) jen.Code {