karmada/hack/tools/lifted-gen/lifted-gen.go

package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/olekukonko/tablewriter"
)

const (
	// a lifted comment line shall be like `+lifted:key=value`
	liftedPrefix            = "+lifted:"
	liftedKeyValueSeparator = "="

	liftedKeySource  = "source"
	liftedKeyChanged = "changed"

	docPrefix = "// Code generated by hack/update-lifted. DO NOT EDIT.\n\n" +
		"// Package lifted contains the files lifted from other projects.\n" +
		"package lifted\n\n/*\n"
	docSuffix = "*/\n"
)

var (
	// TODO: add to flag
	liftedDir    = "pkg/util/lifted"
	output       = "pkg/util/lifted/doc.go"
	excludeFiles = []string{"doc.go"}
)

func main() {
	a := newAnalyzer()
	a.collect(liftedDir)
	a.dump(output)

	fmt.Println()
	if a.failed {
		os.Exit(1)
	}
}

type analyzer struct {
	fset   *token.FileSet
	failed bool
	items  []item
}

func newAnalyzer() *analyzer {
	a := &analyzer{
		fset: token.NewFileSet(),
	}
	return a
}

func (a *analyzer) collect(dir string) {
	pkgs, err := parser.ParseDir(a.fset, dir, a.filter, parser.AllErrors|parser.ParseComments)
	if err != nil {
		a.errorf("syntax error: %v", err)
		return
	}

	for _, pkg := range pkgs {
		sortedFilenames := a.sortFilenames(pkg.Files)
		for _, filename := range sortedFilenames {
			file := pkg.Files[filename]
			basename := filepath.Base(filename)
			a.collectFile(basename, file)
		}
	}
}

func (a *analyzer) filter(info fs.FileInfo) bool {
	for _, pattern := range excludeFiles {
		match, err := filepath.Match(pattern, info.Name())
		if err != nil {
			a.errorf("match %v %v error: %v", pattern, info.Name(), err)
			return false
		}
		if match {
			return false
		}
	}
	return true
}

func (a *analyzer) sortFilenames(m map[string]*ast.File) []string {
	files := make([]string, 0, len(m))
	for n := range m {
		files = append(files, n)
	}
	sort.Strings(files)
	return files
}

func (a *analyzer) collectFile(filename string, file *ast.File) {
	cmap := ast.NewCommentMap(a.fset, file, file.Comments)
	for _, decl := range file.Decls {
		switch decl := decl.(type) {
		case *ast.FuncDecl:
			a.collectFuncDecl(cmap, filename, decl)
		case *ast.GenDecl:
			a.collectGenDecl(cmap, filename, decl)
		}
	}
}

func (a *analyzer) collectGenDecl(cmap ast.CommentMap, file string, decl *ast.GenDecl) {
	for _, spec := range decl.Specs {
		switch spec := spec.(type) {
		case *ast.TypeSpec:
			a.collectTypeSpec(cmap, file, decl, spec)
		case *ast.ValueSpec:
			a.collectValueSpec(cmap, file, decl, spec)
		}
	}
}

func (a *analyzer) collectTypeSpec(cmap ast.CommentMap, file string, decl *ast.GenDecl, spec *ast.TypeSpec) {
	// for GenDecl, comments are in decl, not spec
	item, ok := a.parseComments(cmap[decl])
	if !ok {
		return
	}
	item.file = file
	item.kindName += decl.Tok.String() /* type */ + " " + spec.Name.Name
	a.items = append(a.items, item)
}

func (a *analyzer) collectValueSpec(cmap ast.CommentMap, file string, decl *ast.GenDecl, spec *ast.ValueSpec) {
	// for GenDecl, comments are in decl, not spec
	item, ok := a.parseComments(cmap[decl])
	if !ok {
		return
	}
	item.file = file
	// One spec may contains one or more name. we combine then to one item.
	// e.g.
	// ```
	// var a, b = 1, 2
	// ```
	// spec.Names will be [a, b]. And generate `var a b` as item's kindName.
	item.kindName = decl.Tok.String() /* var/const */
	for _, name := range spec.Names {
		item.kindName += " " + name.Name
	}
	a.items = append(a.items, item)
}

func (a *analyzer) collectFuncDecl(cmap ast.CommentMap, file string, decl *ast.FuncDecl) {
	item, ok := a.parseComments(cmap[decl])
	if !ok {
		return
	}
	item.file = file
	item.kindName = "func " + decl.Name.Name
	a.items = append(a.items, item)
}

func (a *analyzer) parseComments(comments []*ast.CommentGroup) (item item, ok bool) {
	for _, comment := range comments {
		for _, c := range comment.List {
			key, value, isLifted := parseLiftedLine(c.Text)
			if !isLifted {
				continue
			}

			switch key {
			case liftedKeySource:
				item.source = value
				ok = true
			case liftedKeyChanged:
				item.changed = parseBool(value, true)
			}
		}
	}
	return
}

func (a *analyzer) dump(path string) {
	file, err := os.Create(path)
	if err != nil {
		a.errorf("open %s error: %v", path, err)
		return
	}
	defer file.Close()

	fmt.Fprint(file, docPrefix)
	defer fmt.Fprint(file, docSuffix)

	table := newTableWriter(file)
	for _, lifted := range a.items {
		row := []string{lifted.file, lifted.source, lifted.kindName, "N"}
		if lifted.changed {
			row[3] = "Y"
		}
		table.Append(row)
	}
	table.Render()
}

func newTableWriter(w io.Writer) *tablewriter.Table {
	table := tablewriter.NewWriter(w)
	table.SetHeader([]string{"lifted file ", "source file", "const/var/type/func", "changed"})
	table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
	table.SetAutoFormatHeaders(false)
	table.SetAutoWrapText(false)
	table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false})
	table.SetCenterSeparator("|")
	return table
}

func (a *analyzer) errorf(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "ERROR "+format, args...)
	fmt.Fprintln(os.Stderr)
	a.failed = true
}

type item struct {
	file     string
	source   string
	kindName string
	changed  bool
}

func parseLiftedLine(line string) (key, value string, isLifted bool) {
	// a lifted lie is like: `// +lifted:key=value`
	line = strings.TrimSpace(line)
	line = strings.TrimPrefix(line, "//")
	line = strings.TrimSpace(line)

	isLifted = strings.HasPrefix(line, liftedPrefix)
	if !isLifted {
		return
	}
	line = strings.TrimPrefix(line, liftedPrefix)
	parts := strings.SplitN(line, liftedKeyValueSeparator, 2)
	if len(parts) > 0 {
		key = strings.TrimSpace(parts[0])
	}
	if len(parts) > 1 {
		value = strings.TrimSpace(parts[1])
	}
	return
}

func parseBool(s string, defaultValue bool) bool {
	s = strings.ToLower(s)
	switch s {
	case "true":
		return true
	case "false":
		return false
	default:
		return defaultValue
	}
}