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

/*
Copyright 2022 The Karmada 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 (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/olekukonko/tablewriter"
	"github.com/spf13/pflag"
)

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 (
	liftedDir    = ""
	output       = ""
	excludeFiles = []string{"doc.go"}
)

func main() {
	pflag.StringVar(&liftedDir, "lifted-dir", "pkg/util/lifted", "The path of lifted dir")
	pflag.StringVarP(&output, "output", "o", "pkg/util/lifted/doc.go", "The path of lifted-gen output")
	pflag.Parse()

	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
	}
}