github.com/cloudwego/kitex@v0.9.0/tool/internal_pkg/generator/completor.go

// Copyright 2021 CloudWeGo 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 generator

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"golang.org/x/tools/go/ast/astutil"

	"github.com/cloudwego/kitex/tool/internal_pkg/log"
	"github.com/cloudwego/kitex/tool/internal_pkg/tpl"
)

var errNoNewMethod = fmt.Errorf("no new method")

type completer struct {
	allMethods  []*MethodInfo
	handlerPath string
	serviceName string
}

func newCompleter(allMethods []*MethodInfo, handlerPath, serviceName string) *completer {
	return &completer{
		allMethods:  allMethods,
		handlerPath: handlerPath,
		serviceName: serviceName,
	}
}

func parseFuncDecl(fd *ast.FuncDecl) (recvName, funcName string) {
	funcName = fd.Name.String()
	if fd.Recv != nil && len(fd.Recv.List) > 0 {
		v := fd.Recv.List[0]
		switch xv := v.Type.(type) {
		case *ast.StarExpr:
			if si, ok := xv.X.(*ast.Ident); ok {
				recvName = si.Name
			}
		case *ast.Ident:
			recvName = xv.Name
		}
	}
	return
}

func (c *completer) compare(pkg *ast.Package) []*MethodInfo {
	var newMethods []*MethodInfo
	for _, m := range c.allMethods {
		var have bool
	PKGFILES:
		for _, file := range pkg.Files {
			for _, d := range file.Decls {
				if fd, ok := d.(*ast.FuncDecl); ok {
					rn, fn := parseFuncDecl(fd)
					if rn == c.serviceName+"Impl" && fn == m.Name {
						have = true
						break PKGFILES
					}
				}
			}
		}
		if !have {
			log.Infof("[complete handler] add '%s' to handler.go\n", m.Name)
			newMethods = append(newMethods, m)
		}
	}

	return newMethods
}

func (c *completer) addImplementations(w io.Writer, newMethods []*MethodInfo) error {
	// generate implements of new methods
	mt := template.New(HandlerFileName).Funcs(funcs)
	mt = template.Must(mt.Parse(`{{template "HandlerMethod" .}}`))
	mt = template.Must(mt.Parse(tpl.HandlerMethodsTpl))
	data := struct {
		AllMethods  []*MethodInfo
		ServiceName string
	}{
		AllMethods:  newMethods,
		ServiceName: c.serviceName,
	}
	var buf bytes.Buffer
	if err := mt.ExecuteTemplate(&buf, HandlerFileName, data); err != nil {
		return err
	}
	_, err := w.Write(buf.Bytes())
	return err
}

// add imports for new methods
func (c *completer) addImport(w io.Writer, newMethods []*MethodInfo, fset *token.FileSet, handlerAST *ast.File) error {
	newImports := make(map[string]bool)
	for _, m := range newMethods {
		for _, arg := range m.Args {
			for _, dep := range arg.Deps {
				newImports[dep.PkgRefName+" "+dep.ImportPath] = true
			}
		}
		if m.Resp != nil {
			for _, dep := range m.Resp.Deps {
				newImports[dep.PkgRefName+" "+dep.ImportPath] = true
			}
		}
	}
	imports := handlerAST.Imports
	for _, i := range imports {
		path := strings.Trim(i.Path.Value, "\"")
		var aliasPath string
		// remove imports that already in handler.go
		if i.Name != nil {
			aliasPath = i.Name.String() + " " + path
		} else {
			aliasPath = filepath.Base(path) + " " + path
			delete(newImports, path)
		}
		delete(newImports, aliasPath)
	}
	for path := range newImports {
		s := strings.Split(path, " ")
		switch len(s) {
		case 1:
			astutil.AddImport(fset, handlerAST, strings.Trim(s[0], "\""))
		case 2:
			astutil.AddNamedImport(fset, handlerAST, s[0], strings.Trim(s[1], "\""))
		default:
			log.Warn("cannot recognize import path", path)
		}
	}
	printer.Fprint(w, fset, handlerAST)
	return nil
}

func (c *completer) process(w io.Writer) error {
	// get AST of main package
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, filepath.Dir(c.handlerPath), nil, parser.ParseComments)
	if err != nil {
		err = fmt.Errorf("go/parser failed to parse the main package: %w", err)
		log.Warn("NOTICE: This is not a bug. We cannot add new methods to handler.go because your codes failed to compile. Fix the compile errors and try again.\n%s", err.Error())
		return err
	}
	main, ok := pkgs["main"]
	if !ok {
		return fmt.Errorf("main package not found")
	}

	newMethods := c.compare(main)
	if len(newMethods) == 0 {
		return errNoNewMethod
	}
	err = c.addImport(w, newMethods, fset, main.Files[c.handlerPath])
	if err != nil {
		return fmt.Errorf("add imports failed error: %v", err)
	}
	err = c.addImplementations(w, newMethods)
	if err != nil {
		return fmt.Errorf("add implements failed error: %v", err)
	}
	return nil
}

func (c *completer) CompleteMethods() (*File, error) {
	var buf bytes.Buffer
	err := c.process(&buf)
	if err != nil {
		return nil, err
	}
	return &File{Name: c.handlerPath, Content: buf.String()}, nil
}

type commonCompleter struct {
	path   string
	pkg    *PackageInfo
	update *Update
}

func (c *commonCompleter) Complete() (*File, error) {
	var w bytes.Buffer
	// get AST of main package
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, c.path, nil, parser.ParseComments)
	if err != nil {
		err = fmt.Errorf("go/parser failed to parse the file: %s, err: %v", c.path, err)
		log.Warnf("NOTICE: This is not a bug. We cannot update the file %s because your codes failed to compile. Fix the compile errors and try again.\n%s", c.path, err.Error())
		return nil, err
	}

	newMethods, err := c.compare()
	if err != nil {
		return nil, err
	}
	if len(newMethods) == 0 {
		return nil, errNoNewMethod
	}
	err = c.addImport(&w, newMethods, fset, f)
	if err != nil {
		return nil, fmt.Errorf("add imports failed error: %v", err)
	}
	err = c.addImplementations(&w, newMethods)
	if err != nil {
		return nil, fmt.Errorf("add implements failed error: %v", err)
	}
	return &File{Name: c.path, Content: w.String()}, nil
}

func (c *commonCompleter) compare() ([]*MethodInfo, error) {
	var newMethods []*MethodInfo
	for _, m := range c.pkg.Methods {
		c.pkg.Methods = []*MethodInfo{m}
		keyTask := &Task{
			Text: c.update.Key,
		}
		key, err := keyTask.RenderString(c.pkg)
		if err != nil {
			return newMethods, err
		}
		have := false

		dir := c.path
		if strings.HasSuffix(c.path, ".go") {
			dir = path.Dir(c.path)
		}
		filepath.Walk(dir, func(fullPath string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}

			if path.Base(dir) == info.Name() && info.IsDir() {
				return nil
			}
			if info.IsDir() {
				return filepath.SkipDir
			}
			if !strings.HasSuffix(fullPath, ".go") {
				return nil
			}
			// get AST of main package
			fset := token.NewFileSet()
			f, err := parser.ParseFile(fset, fullPath, nil, parser.ParseComments)
			if err != nil {
				err = fmt.Errorf("go/parser failed to parse the file: %s, err: %v", c.path, err)
				log.Warnf("NOTICE: This is not a bug. We cannot update the file %s because your codes failed to compile. Fix the compile errors and try again.\n%s", c.path, err.Error())
				return err
			}

			for _, d := range f.Decls {
				if fd, ok := d.(*ast.FuncDecl); ok {
					_, fn := parseFuncDecl(fd)
					if fn == key {
						have = true
						break
					}
				}
			}
			return nil
		})

		if !have {
			newMethods = append(newMethods, m)
		}
	}

	return newMethods, nil
}

// add imports for new methods
func (c *commonCompleter) addImport(w io.Writer, newMethods []*MethodInfo, fset *token.FileSet, handlerAST *ast.File) error {
	existImports := make(map[string]bool)
	for _, i := range handlerAST.Imports {
		existImports[strings.Trim(i.Path.Value, "\"")] = true
	}
	tmp := c.pkg.Methods
	defer func() {
		c.pkg.Methods = tmp
	}()
	c.pkg.Methods = newMethods
	for _, i := range c.update.ImportTpl {
		importTask := &Task{
			Text: i,
		}
		content, err := importTask.RenderString(c.pkg)
		if err != nil {
			return err
		}
		imports := c.parseImports(content)
		for idx := range imports {
			if _, ok := existImports[strings.Trim(imports[idx][1], "\"")]; !ok {
				astutil.AddImport(fset, handlerAST, strings.Trim(imports[idx][1], "\""))
			}
		}
	}
	printer.Fprint(w, fset, handlerAST)
	return nil
}

func (c *commonCompleter) addImplementations(w io.Writer, newMethods []*MethodInfo) error {
	tmp := c.pkg.Methods
	defer func() {
		c.pkg.Methods = tmp
	}()
	c.pkg.Methods = newMethods
	// generate implements of new methods
	appendTask := &Task{
		Text: c.update.AppendTpl,
	}
	content, err := appendTask.RenderString(c.pkg)
	if err != nil {
		return err
	}
	_, err = w.Write([]byte(content))
	c.pkg.Methods = tmp
	return err
}

// imports[2] is alias, import
func (c *commonCompleter) parseImports(content string) (imports [][2]string) {
	if !strings.Contains(content, "\"") {
		imports = append(imports, [2]string{"", content})
		return imports
	}
	for i := 0; i < len(content); i++ {
		if content[i] == ' ' {
			continue
		}
		isAlias := content[i] != '"'

		start := i
		for ; i < len(content); i++ {
			if content[i] == ' ' {
				break
			}
		}
		sub := content[start:i]
		switch {
		case isAlias:
			imports = append(imports, [2]string{sub, ""})
		case len(imports) > 0 && imports[len(imports)-1][1] == "":
			imports[len(imports)-1][1] = sub
		default:
			imports = append(imports, [2]string{"", sub})
		}
	}
	return imports
}