VisitorGen rewrite

go/tools/asthelpergen/asthelpergen.go

@@ -0,0 +1,236 @@
+/*
+Copyright 2021 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 main
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"go/types"
+	"io/ioutil"
+	"log"
+	"path"
+	"strings"
+
+	"github.com/dave/jennifer/jen"
+	"golang.org/x/tools/go/packages"
+)
+
+const licenseFileHeader = `Copyright 2021 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.`
+
+type astHelperGen struct {
+	DebugTypes bool
+	mod        *packages.Module
+	sizes      types.Sizes
+	namedIface *types.Named
+	iface      *types.Interface
+}
+
+type rewriterFile struct {
+	cases          []jen.Code
+	replaceMethods []jen.Code
+}
+
+func newGenerator(mod *packages.Module, sizes types.Sizes, named *types.Named) *astHelperGen {
+	return &astHelperGen{
+		DebugTypes: true,
+		mod:        mod,
+		sizes:      sizes,
+		namedIface: named,
+		iface:      named.Underlying().(*types.Interface),
+	}
+}
+
+func findImplementations(scope *types.Scope, iff *types.Interface, impl func(types.Type)) {
+	for _, name := range scope.Names() {
+		obj := scope.Lookup(name)
+		baseType := obj.Type()
+		if types.Implements(baseType, iff) || types.Implements(types.NewPointer(baseType), iff) {
+			impl(baseType)
+		}
+	}
+}
+
+func (gen *astHelperGen) doIt() (map[string]*jen.File, error) {
+	pkg := gen.namedIface.Obj().Pkg()
+	rewriter := &rewriterFile{}
+
+	iface, ok := gen.iface.Underlying().(*types.Interface)
+	if !ok {
+		return nil, fmt.Errorf("expected interface, but got %T", gen.iface)
+	}
+
+	var outerErr error
+
+	findImplementations(pkg.Scope(), iface, func(t types.Type) {
+		nt := t.(*types.Named)
+
+		switch n := t.Underlying().(type) {
+		case *types.Struct:
+			switchCase, err := gen.rewriterStructCase(nt.Obj().Name(), n)
+			if err != nil {
+				outerErr = err
+				return
+			}
+			rewriter.cases = append(rewriter.cases, switchCase)
+
+			replaceMethods, err := gen.rewriterReplaceMethods(t.String(), n)
+			if err != nil {
+				outerErr = err
+				return
+			}
+			rewriter.replaceMethods = append(rewriter.cases, replaceMethods...)
+		case *types.Interface:
+
+		default:
+			fmt.Printf("unknown %T\n", t)
+		}
+	})
+
+	if outerErr != nil {
+		return nil, outerErr
+	}
+
+	result := map[string]*jen.File{}
+	fullPath := path.Join(gen.mod.Dir, strings.TrimPrefix(pkg.Path(), gen.mod.Path), "rewriter.go")
+	result[fullPath] = gen.rewriterFile(pkg.Name(), rewriter)
+
+	return result, nil
+}
+
+type typePaths []string
+
+func (t *typePaths) String() string {
+	return fmt.Sprintf("%v", *t)
+}
+
+func (t *typePaths) Set(path string) error {
+	*t = append(*t, path)
+	return nil
+}
+
+func main() {
+	var patterns typePaths
+	var generate string
+	var verify bool
+
+	flag.Var(&patterns, "in", "Go packages to load the generator")
+	flag.StringVar(&generate, "iface", "", "Root interface generate rewriter for")
+	flag.BoolVar(&verify, "verify", false, "ensure that the generated files are correct")
+	flag.Parse()
+
+	result, err := GenerateASTHelpers(patterns, generate)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	if verify {
+		for _, err := range VerifyFilesOnDisk(result) {
+			log.Fatal(err)
+		}
+		log.Printf("%d files OK", len(result))
+	} else {
+		for fullPath, file := range result {
+			if err := file.Save(fullPath); err != nil {
+				log.Fatalf("failed to save file to '%s': %v", fullPath, err)
+			}
+			log.Printf("saved '%s'", fullPath)
+		}
+	}
+}
+
+// VerifyFilesOnDisk compares the generated results from the codegen against the files that
+// currently exist on disk and returns any mismatches
+func VerifyFilesOnDisk(result map[string]*jen.File) (errors []error) {
+	for fullPath, file := range result {
+		existing, err := ioutil.ReadFile(fullPath)
+		if err != nil {
+			errors = append(errors, fmt.Errorf("missing file on disk: %s (%w)", fullPath, err))
+			continue
+		}
+
+		var buf bytes.Buffer
+		if err := file.Render(&buf); err != nil {
+			errors = append(errors, fmt.Errorf("render error for '%s': %w", fullPath, err))
+			continue
+		}
+
+		if !bytes.Equal(existing, buf.Bytes()) {
+			errors = append(errors, fmt.Errorf("'%s' has changed", fullPath))
+			continue
+		}
+	}
+	return errors
+}
+
+// GenerateASTHelpers generates the auxiliary code that implements CachedSize helper methods
+// for all the types listed in typePatterns
+func GenerateASTHelpers(packagePatterns []string, rootIface string) (map[string]*jen.File, error) {
+	loaded, err := packages.Load(&packages.Config{
+		Mode: packages.NeedName | packages.NeedTypes | packages.NeedTypesSizes | packages.NeedTypesInfo | packages.NeedDeps | packages.NeedImports | packages.NeedModule,
+		Logf: log.Printf,
+	}, packagePatterns...)
+
+	if err != nil {
+		return nil, err
+	}
+
+	scopes := make(map[string]*types.Scope)
+	for _, pkg := range loaded {
+		scopes[pkg.PkgPath] = pkg.Types.Scope()
+	}
+
+	pos := strings.LastIndexByte(rootIface, '.')
+	if pos < 0 {
+		return nil, fmt.Errorf("unexpected input type: %s", rootIface)
+	}
+
+	pkgname := rootIface[:pos]
+	typename := rootIface[pos+1:]
+
+	scope := scopes[pkgname]
+	if scope == nil {
+		return nil, fmt.Errorf("no scope found for type '%s'", rootIface)
+	}
+
+	tt := scope.Lookup(typename)
+	if tt == nil {
+		return nil, fmt.Errorf("no type called '%s' found in '%s'", typename, pkgname)
+	}
+
+	nt := tt.Type().(*types.Named)
+
+	generator := newGenerator(loaded[0].Module, loaded[0].TypesSizes, nt)
+	it, err := generator.doIt()
+	if err != nil {
+		return nil, err
+	}
+
+	return it, nil
+}

go/tools/asthelpergen/asthelpergen_test.go

@@ -0,0 +1,36 @@
+/*
+Copyright 2021 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 main
+
+import (
+	"testing"
+)
+
+func TestFullGeneration(t *testing.T) {
+	//result, err := GenerateASTHelpers([]string{"./integration/..."}, []string{"vitess.io/vitess/go/tools/sizegen/integration.*"})
+	//require.NoError(t, err)
+	//
+	//verifyErrors := VerifyFilesOnDisk(result)
+	//require.Empty(t, verifyErrors)
+	//
+	//for _, file := range result {
+	//	contents := fmt.Sprintf("%#v", file)
+	//	require.Contains(t, contents, "http://www.apache.org/licenses/LICENSE-2.0")
+	//	require.Contains(t, contents, "type cachedObject interface")
+	//	require.Contains(t, contents, "//go:nocheckptr")
+	//}
+}

go/tools/asthelpergen/integration/integration_rewriter_test.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2021 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 (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestVisit(t *testing.T) {
+	one := &LiteralInt{1}
+	two := &LiteralInt{1}
+	plus := &Plus{Left: one, Right: two}
+
+	var preOrder, postOrder []AST
+
+	a := &application{
+		pre: func(cursor *Cursor) bool {
+			preOrder = append(preOrder, cursor.node)
+			return true
+		},
+		post: func(cursor *Cursor) bool {
+			postOrder = append(postOrder, cursor.node)
+			return true
+		},
+		cursor: Cursor{},
+	}
+
+	// visit
+	a.apply(nil, plus, nil)
+
+	assert.Equal(t, []AST{plus, one, two}, preOrder, "pre-order wrong")
+	assert.Equal(t, []AST{one, two, plus}, postOrder, "post-order wrong")
+
+}