mkwinsyscall_struct.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/token"
	"go/types"
	"golang.org/x/tools/go/packages"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"unicode"
)

var (
	output     = flag.String("output", "", "输出名称; 默认 srcdir/<struct_name>_struct.go")
	structName = flag.String("struct", "", "转成结构体名称")
)

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of mkwinsyscall_struct.go:\n")
	fmt.Fprintf(os.Stderr, "\tmkwinsyscall_struct [flags] -struct T [directory]\n")
	flag.PrintDefaults()
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("mkwinsyscall_struct: ")
	flag.Usage = Usage
	flag.Parse()
	if len(*structName) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	// We accept either one directory or a list of files. Which do we have?
	args := flag.Args()
	if len(args) == 0 {
		// Default: process whole package in current directory.
		args = []string{"github.com/gorpher/winspool-go"}
	}

	dir, _ := os.Getwd()
	// Parse the package once.
	//var dir = filepath.Dir(args[0])
	g := Generator{}

	g.parsePackage(args)

	// Print the header and package clause.
	g.Printf("// Code generated by \"mkwinsyscall_struct %s\"; DO NOT EDIT.\n", strings.Join(os.Args[1:], " "))
	g.Printf("\n")
	g.Printf("package %s", g.pkg.name)
	g.Printf("\n")

	g.generate(*structName)

	// Format the output.
	src := g.format()

	// Write to file.
	outputName := *output
	if outputName == "" {
		baseName := fmt.Sprintf("%s_struct.go", *structName)
		outputName = filepath.Join(dir, strings.ToLower(baseName))
	}
	err := ioutil.WriteFile(outputName, src, 0644)
	if err != nil {
		log.Fatalf("writing output: %s", err)
	}
}

// isDirectory reports whether the named file is a directory.
func isDirectory(name string) bool {
	info, err := os.Stat(name)
	if err != nil {
		log.Fatal(err)
	}
	return info.IsDir()
}

// Generator holds the state of the analysis. Primarily used to buffer
// the output for format.Source.
type Generator struct {
	buf bytes.Buffer // Accumulated output.
	pkg *Package     // Package we are scanning.

}

func (g *Generator) Printf(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, format, args...)
}

// File holds a single parsed file and associated data.
type File struct {
	pkg        *Package  // Package to which this file belongs.
	file       *ast.File // Parsed AST.
	structName string    // Name of the constant type.
	values     []Value   // Accumulator for constant values of that type.

}

type Package struct {
	name  string
	defs  map[*ast.Ident]types.Object
	files []*File
}

// parsePackage analyzes the single package constructed from the patterns and tags.
// parsePackage exits if there is an error.
func (g *Generator) parsePackage(patterns []string) {
	cfg := &packages.Config{
		Mode:       packages.LoadSyntax,
		Tests:      false,
		BuildFlags: []string{},
	}
	pkgs, err := packages.Load(cfg, patterns...)
	if err != nil {
		log.Fatal(err)
	}
	log.Println(pkgs)
	if len(pkgs) != 1 {
		log.Fatalf("error: %d packages found", len(pkgs))
	}
	g.addPackage(pkgs[0])
}

// addPackage adds a type checked Package and its syntax files to the generator.
func (g *Generator) addPackage(pkg *packages.Package) {
	g.pkg = &Package{
		name:  pkg.Name,
		defs:  pkg.TypesInfo.Defs,
		files: make([]*File, len(pkg.Syntax)),
	}

	for i, file := range pkg.Syntax {
		g.pkg.files[i] = &File{
			file: file,
			pkg:  g.pkg,
		}
	}
}

// generate produces the String method for the named type.
func (g *Generator) generate(structName string) {
	values := make([]Value, 0, 100)
	for _, file := range g.pkg.files {
		// Set the state for this run of the walker.
		file.structName = structName
		file.values = nil
		if file.file != nil {
			ast.Inspect(file.file, file.genDecl)
			values = append(values, file.values...)
		}
	}
	if len(values) == 0 {
		log.Fatalf("no values defined for type %s", structName)
	}

	g.genImportPackage(values)
	g.genGetterMethod(structName, values)
	g.genStruct(structName, values)
}
func (g *Generator) write(p []byte) {
	g.buf.Write(p)
}
func (g *Generator) writeString(s string) {
	g.buf.WriteString(s)
}

// format returns the gofmt-ed contents of the Generator's buffer.
func (g *Generator) format() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		// Should never happen, but can arise when developing this code.
		// The user can compile the output to see the error.
		log.Printf("warning: internal error: invalid Go generated: %s", err)
		log.Printf("warning: compile the package to analyze the error")
		return g.buf.Bytes()
	}
	return src
}

// Value represents a declared constant.
type Value struct {
	name           string
	Fields         map[string]string
	ImportPackages []string
}

func (v *Value) GetStringFields() string {
	b := strings.Builder{}
	for field, retype := range v.Fields {
		if retype == "string" {
			b.WriteString(fmt.Sprintf("func (p *%s) Get%s() %s {", v.name, processFieldName(field), retype))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    return utf16PtrToString(p.%s)", field))
			b.WriteString("\n")
			b.WriteString("}")
			b.WriteString("\n")
		}
		if retype == "time.Time" {
			b.WriteString(fmt.Sprintf("func (p *%s) Get%s() %s {", v.name, processFieldName(field), retype))
			b.WriteString("\n")
			b.WriteString("    return time.Date(")
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    int(p.%s.Year),", field))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    time.Month(int(p.%s.Month)),", field))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    int(p.%s.Day),", field))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    int(p.%s.Hour),", field))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    int(p.%s.Minute),", field))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    int(p.%s.Second),", field))
			b.WriteString("\n")
			b.WriteString(fmt.Sprintf("    int(1000*p.%s.Milliseconds),", field))
			b.WriteString("\n")
			b.WriteString("    time.Local,")
			b.WriteString("\n")
			b.WriteString("    ).UTC()")
			b.WriteString("\n")
			b.WriteString("}")
			b.WriteString("\n")
		}
	}
	return b.String()
}

func (v *Value) GetStringStruct(structName string) string {
	b := strings.Builder{}
	structName = strings.Title(structName)
	b.WriteString(fmt.Sprintf("func (p *%s) Get%s() %s {", v.name, structName, structName))
	b.WriteString("\n")
	b.WriteString(fmt.Sprintf("    return %s{", structName))
	b.WriteString("\n")
	for field, typ := range v.Fields {
		fld := processFieldName(field)
		if typ == "string" {
			b.WriteString(fmt.Sprintf("			%s:       p.Get%s(),", fld, fld))
			b.WriteString("\n")
		}
		if typ == "uint32" || typ == "int16" {
			b.WriteString(fmt.Sprintf("			%s:       p.%s,", fld, field))
			b.WriteString("\n")
		}
		if typ == "time.Time" {
			b.WriteString(fmt.Sprintf("			%s:       p.Get%s(),", fld, fld))
			b.WriteString("\n")
		}

	}
	b.WriteString(fmt.Sprintf("    }"))
	b.WriteString("\n")
	b.WriteString(fmt.Sprintf("}"))
	b.WriteString("\n")
	return b.String()
}

func processFieldName(field string) string {
	if strings.HasPrefix(field, "p") && len(field) > 1 {
		r := field[1]
		if unicode.IsUpper(rune(r)) {
			field = strings.TrimLeft(field, "p")
		}
	}
	return strings.Title(field)
}

func (g *Generator) genImportPackage(values []Value) {
	imps := map[string]struct{}{}
	for _, value := range values {
		for _, v := range value.ImportPackages {
			imps[v] = struct{}{}
		}
	}
	for k := range imps {
		g.Printf("import \"%s\"\n", k)
	}
}
func (g *Generator) genGetterMethod(structName string, values []Value) {
	for _, value := range values {
		g.writeString(value.GetStringFields())
		g.writeString(value.GetStringStruct(structName))
	}
}

func (g *Generator) genStruct(structName string, values []Value) {
	b := strings.Builder{}
	fields := map[string]string{}
	for _, value := range values {
		for k, v := range value.Fields {
			fields[processFieldName(k)] = v
		}
	}
	structName = strings.Title(structName)
	b.WriteString(fmt.Sprintf("type %s struct {", structName))
	b.WriteString("\n")
	for field, typ := range fields {
		if typ == "string" {
			b.WriteString(fmt.Sprintf("%s       %s", field, typ))
			b.WriteString("\n")
		}
		if typ == "uint32" || typ == "int16" {
			b.WriteString(fmt.Sprintf("%s       %s", field, typ))
			b.WriteString("\n")
		}
		if typ == "time.Time" {
			b.WriteString(fmt.Sprintf("%s       %s", field, typ))
			b.WriteString("\n")
		}
	}
	b.WriteString("}")
	b.WriteString("\n")
	g.writeString(b.String())
}

// genDecl processes one declaration clause.
func (f *File) genDecl(node ast.Node) bool {
	decl, ok := node.(*ast.GenDecl)
	if !ok || decl.Tok != token.TYPE {
		return true
	}
	for _, spec := range decl.Specs {
		vspec := spec.(*ast.TypeSpec)
		if vspec.Type != nil && len(vspec.Name.Name) > 0 {
			if !IsAllUpperCase(vspec.Name.Name) {
				continue
			}
			if !strings.HasPrefix(strings.ToLower(vspec.Name.Name), strings.ToLower(f.structName)) {
				continue
			}
			sType, ok := vspec.Type.(*ast.StructType)
			if ok && sType.Fields.NumFields() > 0 {
				value := Value{
					name:           vspec.Name.Name,
					Fields:         map[string]string{},
					ImportPackages: []string{},
				}
				for i := range sType.Fields.List {
					//fmt.Println(sType.Fields.List[i].Names, sType.Fields.List[i].Type)
					star, ok := sType.Fields.List[i].Type.(*ast.StarExpr)
					if ok {
						indent, ok := star.X.(*ast.Ident)
						if ok {
							if indent.Name == "uint16" {
								value.Fields[getFiledName(sType.Fields.List[i].Names)] = "string"
								continue
							}
						}
					}
					indent, ok := sType.Fields.List[i].Type.(*ast.Ident)
					if ok {
						if indent.Name == "uintptr" {
							// TODO
							//value.Fields[getFiledName(sType.Fields.List[i].Names)] = "string"
						}
						if indent.Name == "uint32" {
							value.Fields[getFiledName(sType.Fields.List[i].Names)] = "uint32"
							continue
						}
						if indent.Name == "int16" {
							value.Fields[getFiledName(sType.Fields.List[i].Names)] = "int16"
							continue
						}
						//fmt.Println(sType.Fields.List[i].Names, indent.Name)
					}
					selector, ok := sType.Fields.List[i].Type.(*ast.SelectorExpr)
					if ok {
						if selector.Sel.Name == "Systemtime" {
							value.Fields[getFiledName(sType.Fields.List[i].Names)] = "time.Time"
							value.ImportPackages = append(value.ImportPackages, "time")
							continue
						}
					}
				}
				f.values = append(f.values, value)
			}
		}
	}
	return false
}

func getFiledName(idents []*ast.Ident) string {
	return idents[0].Name
}

func IsAllUpperCase(s string) bool {
	for _, r := range s {
		if !(unicode.IsUpper(r) || '_' == r || unicode.IsNumber(r)) {
			return false
		}
	}
	return true
}