feat: make it possible for the rewriter to revisit any node type

Signed-off-by: Andres Taylor <andres@planetscale.com>

Author: systay

go/tools/asthelpergen/asthelpergen.go

@@ -77,9 +77,6 @@ type (
 	}
 )
 
-// exprInterfacePath is the path of the sqlparser.Expr interface.
-const exprInterfacePath = "vitess.io/vitess/go/vt/sqlparser.Expr"
-
 func (gen *astHelperGen) iface() *types.Interface {
 	return gen._iface
 }
@@ -207,19 +204,13 @@ func GenerateASTHelpers(options *Options) (map[string]*jen.File, error) {
 		return nil, err
 	}
 
-	exprType, _ := findTypeObject(exprInterfacePath, scopes)
-	var exprInterface *types.Interface
-	if exprType != nil {
-		exprInterface = exprType.Type().(*types.Named).Underlying().(*types.Interface)
-	}
-
 	nt := tt.Type().(*types.Named)
 	pName := nt.Obj().Pkg().Name()
 	generator := newGenerator(loaded[0].Module, loaded[0].TypesSizes, nt,
 		newEqualsGen(pName, &options.Equals),
 		newCloneGen(pName, &options.Clone),
 		newVisitGen(pName),
-		newRewriterGen(pName, types.TypeString(nt, noQualifier), exprInterface),
+		newRewriterGen(pName, types.TypeString(nt, noQualifier)),
 		newCOWGen(pName, nt),
 	)
 

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/integration/ast_rewrite.go

@@ -204,6 +204,46 @@ func TestRewriteAndRevisitInterfaceSlice(t *testing.T) {
 	})
 }
 
+func TestRewriteAndRevisitRefContainer(t *testing.T) {
+	leaf1 := &Leaf{1}
+	leaf2 := &Leaf{2}
+	ast1 := &RefContainer{
+		ASTType:               leaf1,
+		ASTImplementationType: leaf2,
+	}
+	ast2 := &RefContainer{
+		ASTType:               leaf2,
+		ASTImplementationType: leaf1,
+	}
+
+	tv := &rewriteTestVisitor{}
+
+	a := false
+	_ = Rewrite(ast1, func(cursor *Cursor) bool {
+		tv.pre(cursor)
+		switch cursor.node.(type) {
+		case *RefContainer:
+			if a {
+				break
+			}
+			a = true
+			cursor.ReplaceAndRevisit(ast2)
+		}
+		return true
+	}, tv.post)
+
+	tv.assertEquals(t, []step{
+		Pre{ast1}, // when we visit ast, we want to replace and revisit,
+		// which means that we don't do a post on this node, or visit the children
+		Pre{ast2},
+		Pre{leaf2},
+		Post{leaf2},
+		Pre{leaf1},
+		Post{leaf1},
+		Post{ast2},
+	})
+}
+
 func TestRewriteVisitRefContainerReplace(t *testing.T) {
 	ast := &RefContainer{
 		ASTType:               &RefContainer{NotASTType: 12},
@@ -361,29 +401,30 @@ func (tv *rewriteTestVisitor) assertEquals(t *testing.T, expected []step) {
 	error := false
 	expectedSize := len(expected)
 	for i, step := range tv.walk {
-		if expectedSize <= i {
-			t.Errorf("❌️ - Expected less elements %v", tv.walk[i:])
-			break
-		} else {
-			e := expected[i]
-			if reflect.DeepEqual(e, step) {
-				a := "✔️ - " + e.String()
-				if error {
-					fmt.Println(a)
-				} else {
-					lines = append(lines, a)
-				}
+		t.Run(fmt.Sprintf("step %d", i), func(t *testing.T) {
+			if expectedSize <= i {
+				t.Fatalf("❌️ - Expected less elements %v", tv.walk[i:])
 			} else {
-				if !error {
-					// first error we see.
-					error = true
-					for _, line := range lines {
-						fmt.Println(line)
+				e := expected[i]
+				if reflect.DeepEqual(e, step) {
+					a := "✔️ - " + e.String()
+					if error {
+						fmt.Println(a)
+					} else {
+						lines = append(lines, a)
+					}
+				} else {
+					if !error {
+						// first error we see.
+						error = true
+						for _, line := range lines {
+							fmt.Println(line)
+						}
 					}
+					t.Fatalf("❌️ - Expected: %s Got: %s\n", e.String(), step.String())
 				}
-				t.Errorf("❌️ - Expected: %s Got: %s\n", e.String(), step.String())
 			}
-		}
+		})
 	}
 	walkSize := len(tv.walk)
 	if expectedSize > walkSize {

go/tools/asthelpergen/integration/integration_rewriter_test.go

@@ -68,14 +68,6 @@ func (c *Cursor) Replace(newNode AST) {
 // When used, this will abort the visitation of the current node - no post or children visited,
 // and the new node visited.
 func (c *Cursor) ReplaceAndRevisit(newNode AST) {
-	switch newNode.(type) {
-	case InterfaceSlice:
-	default:
-		// We need to add support to the generated code for when to look at the revisit flag. At the moment it is only
-		// there for slices of AST implementations
-		panic("no support added for this type yet")
-	}
-
 	c.replacer(newNode, c.parent)
 	c.node = newNode
 	c.revisit = true

go/tools/asthelpergen/integration/test_helpers.go

@@ -30,21 +30,18 @@ const (
 type rewriteGen struct {
 	ifaceName string
 	file      *jen.File
-	// exprInterface is used to store the sqlparser.Expr interface
-	exprInterface *types.Interface
 }
 
 var _ generator = (*rewriteGen)(nil)
 
-func newRewriterGen(pkgname string, ifaceName string, exprInterface *types.Interface) *rewriteGen {
+func newRewriterGen(pkgname string, ifaceName string) *rewriteGen {
 	file := jen.NewFile(pkgname)
 	file.HeaderComment(licenseFileHeader)
 	file.HeaderComment("Code generated by ASTHelperGen. DO NOT EDIT.")
 
 	return &rewriteGen{
-		ifaceName:     ifaceName,
-		file:          file,
-		exprInterface: exprInterface,
+		ifaceName: ifaceName,
+		file:      file,
 	}
 }
 
@@ -108,7 +105,7 @@ func (r *rewriteGen) structMethod(t types.Type, strct *types.Struct, spi generat
 	}
 	fields := r.rewriteAllStructFields(t, strct, spi, true)
 
-	stmts := []jen.Code{r.executePre(t)}
+	stmts := []jen.Code{r.executePre()}
 	stmts = append(stmts, fields...)
 	stmts = append(stmts, executePost(len(fields) > 0))
 	stmts = append(stmts, returnTrue())
@@ -133,7 +130,7 @@ func (r *rewriteGen) ptrToStructMethod(t types.Type, strct *types.Struct, spi ge
 			return nil
 		}
 	*/
-	stmts = append(stmts, r.executePre(t))
+	stmts = append(stmts, r.executePre())
 	fields := r.rewriteAllStructFields(t, strct, spi, false)
 	stmts = append(stmts, fields...)
 	stmts = append(stmts, executePost(len(fields) > 0))
@@ -182,15 +179,12 @@ func (r *rewriteGen) sliceMethod(t types.Type, slice *types.Slice, spi generator
 		jen.If(jen.Id("node == nil").Block(returnTrue())),
 	}
 
-	typeString := types.TypeString(t, noQualifier)
-
 	preStmts := setupCursor()
 	preStmts = append(preStmts,
 		jen.Id("kontinue").Op(":=").Id("!a.pre(&a.cur)"),
 		jen.If(jen.Id("a.cur.revisit").Block(
-			jen.Id("node").Op("=").Id("a.cur.node.("+typeString+")"),
 			jen.Id("a.cur.revisit").Op("=").False(),
-			jen.Return(jen.Id("a.rewrite"+typeString+"(parent, node, replacer)")),
+			jen.Return(jen.Id("a.rewrite"+r.ifaceName+"(parent, a.cur.node, replacer)")),
 		)),
 		jen.If(jen.Id("kontinue").Block(jen.Return(jen.True()))),
 	)
@@ -228,19 +222,17 @@ func setupCursor() []jen.Code {
 		jen.Id("a.cur.node = node"),
 	}
 }
-func (r *rewriteGen) executePre(t types.Type) jen.Code {
+
+func (r *rewriteGen) executePre() jen.Code {
 	curStmts := setupCursor()
-	if r.exprInterface != nil && types.Implements(t, r.exprInterface) {
-		curStmts = append(curStmts, jen.Id("kontinue").Op(":=").Id("!a.pre(&a.cur)"),
-			jen.If(jen.Id("a.cur.revisit").Block(
-				jen.Id("a.cur.revisit").Op("=").False(),
-				jen.Return(jen.Id("a.rewriteExpr(parent, a.cur.node.(Expr), replacer)")),
-			)),
-			jen.If(jen.Id("kontinue").Block(jen.Return(jen.True()))),
-		)
-	} else {
-		curStmts = append(curStmts, jen.If(jen.Id("!a.pre(&a.cur)")).Block(returnTrue()))
-	}
+	name := fmt.Sprintf("a.rewrite%s(parent, a.cur.node, replacer)", r.ifaceName)
+	curStmts = append(curStmts, jen.Id("kontinue").Op(":=").Id("!a.pre(&a.cur)"),
+		jen.If(jen.Id("a.cur.revisit").Block(
+			jen.Id("a.cur.revisit").Op("=").False(),
+			jen.Return(jen.Id(name)),
+		)),
+		jen.If(jen.Id("kontinue").Block(jen.Return(jen.True()))),
+	)
 	return jen.If(jen.Id("a.pre!= nil").Block(curStmts...))
 }
 
@@ -264,7 +256,7 @@ func (r *rewriteGen) basicMethod(t types.Type, _ *types.Basic, spi generatorSPI)
 		return nil
 	}
 
-	stmts := []jen.Code{r.executePre(t), executePost(false), returnTrue()}
+	stmts := []jen.Code{r.executePre(), executePost(false), returnTrue()}
 	r.rewriteFunc(t, stmts)
 	return nil
 }