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

Signed-off-by: Andres Taylor <[email protected]>

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/integration_rewriter_test.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/test_helpers.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