Merge pull request #7716 from planetscale/rewrite-gen-2

Speed up AST rewriting

.github/workflows/check_make_visitor.yml

@@ -31,5 +31,5 @@ jobs:
 
     - name: check_make_visitor
       run: |
-        misc/git/hooks/visitorgen
+        misc/git/hooks/asthelpers
 

Makefile

@@ -103,6 +103,10 @@ parser:
 	make -C go/vt/sqlparser
 
 visitor:
+	>&2 echo "make visitor has been replaced by make asthelpers"
+	exit 1
+
+asthelpers:
 	go run ./go/tools/asthelpergen/main -in ./go/vt/sqlparser -iface vitess.io/vitess/go/vt/sqlparser.SQLNode -except "*ColName"
 
 sizegen:
@@ -123,7 +127,6 @@ clean:
 	go clean -i ./go/...
 	rm -rf third_party/acolyte
 	rm -rf go/vt/.proto.tmp
-	rm -rf ./visitorgen
 
 # Remove everything including stuff pulled down by bootstrap.sh
 cleanall: clean

go/tools/asthelpergen/asthelpergen.go

@@ -23,6 +23,7 @@ import (
 	"io/ioutil"
 	"log"
 	"path"
+	"sort"
 	"strings"
 
 	"github.com/dave/jennifer/jen"
@@ -43,52 +44,65 @@ 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.`
 
-type generator interface {
-	visitStruct(t types.Type, stroct *types.Struct) error
-	visitInterface(t types.Type, iface *types.Interface) error
-	visitSlice(t types.Type, slice *types.Slice) error
-	createFile(pkgName string) (string, *jen.File)
-}
+type (
+	generatorSPI interface {
+		addType(t types.Type)
+		addFunc(name string, t methodType, code jen.Code)
+		scope() *types.Scope
+		findImplementations(iff *types.Interface, impl func(types.Type) error) error
+		iface() *types.Interface
+	}
+	generator2 interface {
+		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
+		ptrToBasicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error
+		sliceMethod(t types.Type, slice *types.Slice, spi generatorSPI) error
+		basicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error
+	}
+	// astHelperGen finds implementations of the given interface,
+	// and uses the supplied `generator`s to produce the output code
+	astHelperGen struct {
+		DebugTypes bool
+		mod        *packages.Module
+		sizes      types.Sizes
+		namedIface *types.Named
+		_iface     *types.Interface
+		gens       []generator2
+
+		functions methods
+		_scope    *types.Scope
+		todo      []types.Type
+	}
 
-type generatorSPI interface {
-	addType(t types.Type)
-	addFunc(name string, t methodType, code jen.Code)
-	scope() *types.Scope
-	findImplementations(iff *types.Interface, impl func(types.Type) error) error
-	iface() *types.Interface
+	method struct {
+		name string
+		code jen.Code
+		typ  methodType
+	}
+
+	methods []method
+)
+
+func (m methods) Len() int {
+	return len(m)
 }
 
-type generator2 interface {
-	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
-	ptrToBasicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error
-	ptrToOtherMethod(t types.Type, ptr *types.Pointer, spi generatorSPI) error
-	sliceMethod(t types.Type, slice *types.Slice, spi generatorSPI) error
-	basicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error
+func (m methods) Less(i, j int) bool {
+	return m[i].name < m[j].name
 }
 
-// astHelperGen finds implementations of the given interface,
-// and uses the supplied `generator`s to produce the output code
-type astHelperGen struct {
-	DebugTypes bool
-	mod        *packages.Module
-	sizes      types.Sizes
-	namedIface *types.Named
-	_iface     *types.Interface
-	gens       []generator
-	gens2      []generator2
-
-	methods []jen.Code
-	_scope  *types.Scope
-	todo    []types.Type
+func (m methods) Swap(i, j int) {
+	m[i], m[j] = m[j], m[i]
 }
 
+var _ sort.Interface = (methods)(nil)
+
 func (gen *astHelperGen) iface() *types.Interface {
 	return gen._iface
 }
 
-func newGenerator(mod *packages.Module, sizes types.Sizes, named *types.Named, generators ...generator) *astHelperGen {
+func newGenerator(mod *packages.Module, sizes types.Sizes, named *types.Named, generators ...generator2) *astHelperGen {
 	return &astHelperGen{
 		DebugTypes: true,
 		mod:        mod,
@@ -150,73 +164,11 @@ func (gen *astHelperGen) findImplementations(iff *types.Interface, impl func(typ
 	return nil
 }
 
-func (gen *astHelperGen) visitStruct(t types.Type, stroct *types.Struct) error {
-	for _, g := range gen.gens {
-		err := g.visitStruct(t, stroct)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func (gen *astHelperGen) visitSlice(t types.Type, slice *types.Slice) error {
-	for _, g := range gen.gens {
-		err := g.visitSlice(t, slice)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
-func (gen *astHelperGen) visitInterface(t types.Type, iface *types.Interface) error {
-	for _, g := range gen.gens {
-		err := g.visitInterface(t, iface)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 // GenerateCode is the main loop where we build up the code per file.
 func (gen *astHelperGen) GenerateCode() (map[string]*jen.File, error) {
 	pkg := gen.namedIface.Obj().Pkg()
-	iface, ok := gen._iface.Underlying().(*types.Interface)
-	if !ok {
-		return nil, fmt.Errorf("expected interface, but got %T", gen.iface)
-	}
-
-	err := findImplementations(pkg.Scope(), iface, func(t types.Type) error {
-		switch n := t.Underlying().(type) {
-		case *types.Struct:
-			return gen.visitStruct(t, n)
-		case *types.Slice:
-			return gen.visitSlice(t, n)
-		case *types.Pointer:
-			strct, isStrct := n.Elem().Underlying().(*types.Struct)
-			if isStrct {
-				return gen.visitStruct(t, strct)
-			}
-		case *types.Interface:
-			return gen.visitInterface(t, n)
-		default:
-			// do nothing
-		}
-		return nil
-	})
-
-	if err != nil {
-		return nil, err
-	}
 
 	result := map[string]*jen.File{}
-	for _, g := range gen.gens {
-		file, code := g.createFile(pkg.Name())
-		fullPath := path.Join(gen.mod.Dir, strings.TrimPrefix(pkg.Path(), gen.mod.Path), file)
-		result[fullPath] = code
-	}
 
 	gen._scope = pkg.Scope()
 	gen.todo = append(gen.todo, gen.namedIface)
@@ -300,16 +252,12 @@ func GenerateASTHelpers(packagePatterns []string, rootIface, exceptCloneType str
 
 	nt := tt.Type().(*types.Named)
 
-	iface := nt.Underlying().(*types.Interface)
-
-	interestingType := func(t types.Type) bool {
-		return types.Implements(t, iface)
-	}
-	rewriter := newRewriterGen(interestingType, nt.Obj().Name())
-	generator := newGenerator(loaded[0].Module, loaded[0].TypesSizes, nt, rewriter)
-	generator.gens2 = append(generator.gens2, &equalsGen{})
-	generator.gens2 = append(generator.gens2, newCloneGen(exceptCloneType))
-	generator.gens2 = append(generator.gens2, &visitGen{})
+	generator := newGenerator(loaded[0].Module, loaded[0].TypesSizes, nt,
+		&equalsGen{},
+		newCloneGen(exceptCloneType),
+		&visitGen{},
+		&rewriteGen{types.TypeString(nt, noQualifier)},
+	)
 
 	it, err := generator.GenerateCode()
 	if err != nil {
@@ -335,19 +283,11 @@ const (
 	clone methodType = iota
 	equals
 	visit
+	rewrite
 )
 
 func (gen *astHelperGen) addFunc(name string, typ methodType, code jen.Code) {
-	var comment string
-	switch typ {
-	case clone:
-		comment = " creates a deep clone of the input."
-	case equals:
-		comment = " does deep equals between the two objects."
-	case visit:
-		comment = " will visit all parts of the AST"
-	}
-	gen.methods = append(gen.methods, jen.Comment(name+comment), code)
+	gen.functions = append(gen.functions, method{name: name, code: code, typ: typ})
 }
 
 func (gen *astHelperGen) createFile(pkgName string) (string, *jen.File) {
@@ -406,15 +346,23 @@ func (gen *astHelperGen) createFile(pkgName string) (string, *jen.File) {
 		alreadyDone[typeName] = true
 	}
 
-	for _, method := range gen.methods {
-		out.Add(method)
+	sort.Sort(gen.functions)
+
+	for _, m := range gen.functions {
+		switch m.typ {
+		case clone:
+			out.Add(jen.Comment(fmt.Sprintf("%s creates a deep clone of the input.", m.name)))
+		case equals:
+			out.Add(jen.Comment(fmt.Sprintf("%s does deep equals between the two objects.", m.name)))
+		}
+		out.Add(m.code)
 	}
 
 	return "ast_helper.go", out
 }
 
 func (gen *astHelperGen) allGenerators(f func(g generator2) error) {
-	for _, g := range gen.gens2 {
+	for _, g := range gen.gens {
 		err := f(g)
 
 		if err != nil {

go/tools/asthelpergen/equals_gen.go

@@ -235,7 +235,3 @@ func (e equalsGen) sliceMethod(t types.Type, slice *types.Slice, spi generatorSP
 func (e equalsGen) basicMethod(t types.Type, basic *types.Basic, spi generatorSPI) error {
 	return nil
 }
-
-func (e equalsGen) ptrToOtherMethod(types.Type, *types.Pointer, generatorSPI) error {
-	return nil
-}