codegen.go

package ens

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/things-go/ens/matcher"
	"github.com/things-go/ens/utils"
	"golang.org/x/tools/imports"
)

var mustEscapeNames = []string{"TableName"}

type CodeGen struct {
	buf               bytes.Buffer
	Entities          []*EntityDescriptor
	ByName            string
	Version           string
	PackageName       string
	DisableDocComment bool
	Option
}

// Bytes returns the CodeBuf's buffer.
func (g *CodeGen) Bytes() []byte {
	return g.buf.Bytes()
}

// FormatSource return formats and adjusts imports contents of the CodeGen's buffer.
func (g *CodeGen) FormatSource() ([]byte, error) {
	data := g.buf.Bytes()
	if len(data) == 0 {
		return data, nil
	}
	// return format.Source(data)
	// 格式化时, 如果需要插入或删除包是非常耗时的
	return imports.Process("", data, &imports.Options{
		Fragment:   false,
		AllErrors:  false,
		Comments:   true,
		TabIndent:  true,
		TabWidth:   4,
		FormatOnly: false,
	})
}

// Write appends the contents of p to the buffer,
func (g *CodeGen) Write(b []byte) (n int, err error) {
	return g.buf.Write(b)
}

// Print formats using the default formats for its operands and writes to the generated output.
// Spaces are added between operands when neither is a string.
// It returns the number of bytes written and any write error encountered.
func (g *CodeGen) Print(a ...any) (n int, err error) {
	return fmt.Fprint(&g.buf, a...)
}

// Printf formats according to a format specifier for its operands and writes to the generated output.
// It returns the number of bytes written and any write error encountered.
func (g *CodeGen) Printf(format string, a ...any) (n int, err error) {
	return fmt.Fprintf(&g.buf, format, a...)
}

// Fprintln formats using the default formats to the generated output.
// Spaces are always added between operands and a newline is appended.
// It returns the number of bytes written and any write error encountered.
func (g *CodeGen) Println(a ...any) (n int, err error) {
	return fmt.Fprintln(&g.buf, a...)
}

func (g *CodeGen) Gen() *CodeGen {
	if !g.DisableDocComment {
		g.Printf("// Code generated by %s. DO NOT EDIT.\n", g.ByName)
		g.Printf("// version: %s\n", g.Version)
		g.Println()
	}
	g.Printf("package %s\n", g.PackageName)
	g.Println()

	//* 先处理转义, 主要是一些需要导入的包, 各种选项. 避免格式化耗时.
	for _, et := range g.Entities {
		et.fixEntityField(&g.Option)
	}
	//* import
	imports := make(map[string]struct{})
	for _, st := range g.Entities {
		for _, field := range st.Fields {
			if field.Type.PkgPath != "" {
				imports[field.Type.PkgPath] = struct{}{}
			}
		}
	}
	if len(imports) > 0 {
		g.Println("import (")
		for k := range imports {
			g.Printf("\"%s\"\n", k)
		}
		g.Println(")")
	}
	//* struct
	for _, et := range g.Entities {
		structName := utils.PascalCase(et.Name)
		tableName := et.Name
		g.Printf("// %s %s\n", structName, strings.ReplaceAll(strings.TrimSpace(et.Comment), "\n", "\n// "))
		g.Printf("type %s struct {\n", structName)
		for _, field := range et.Fields {
			g.Println(g.genModelStructField(field))
		}
		g.Println("}")
		g.Println()
		g.Println("// TableName implement schema.Tabler interface")
		g.Printf("func (*%s) TableName() string {\n", structName)
		g.Printf("return \"%s\"\n", tableName)
		g.Println("}")
		g.Println()
	}
	return g
}

func (g *CodeGen) genModelStructField(field *FieldDescriptor) string {
	b := strings.Builder{}
	b.Grow(128)
	ident := field.Type.Ident
	if field.GoPointer && !field.Type.NonPointer {
		ident = "*" + field.Type.Ident
	}
	// field
	b.WriteString(field.GoName)
	b.WriteString(" ")
	b.WriteString(ident)
	if len(field.Tags) > 0 {
		b.WriteString(" `")
		b.WriteString(strings.Join(field.Tags, " "))
		b.WriteString("`")
	}
	if field.Comment != "" {
		b.WriteString(" // ")
		b.WriteString(field.Comment)
	}
	return b.String()
}

func (et *EntityDescriptor) fixEntityField(opt *Option) {
	if opt == nil {
		opt = defaultOption()
	}
	escapeNames := make(map[string]struct{}, len(mustEscapeNames)+len(opt.EscapeName)+len(et.Fields))
	for _, v := range mustEscapeNames {
		escapeNames[v] = struct{}{}
	}
	for _, k := range opt.EscapeName {
		escapeNames[k] = struct{}{}
	}
	allFieldName := make(map[string]struct{}, len(et.Fields))
	for _, field := range et.Fields {
		allFieldName[field.GoName] = struct{}{}
	}

	for _, field := range et.Fields {
		field.fixField(allFieldName, escapeNames, opt)
	}
}

// 根据规则转义一些数据
func (field *FieldDescriptor) fixField(allFieldName, escapeFieldNames map[string]struct{}, opt *Option) {
	if field.ColumnName == "deleted_at" {
		if field.Type.IsInteger() {
			field.GoPointer = false
			field.Type = SoftDeleteType().WithNewType(field.Type.Type)
		} else if field.Type.IsInteger() {
			field.GoPointer = false
			field.Type = GormDeletedAtType().WithNewType(field.Type.Type)
		}
	}
	if opt == nil {
		opt = defaultOption()
	}
	if opt.EnableInt {
		switch field.Type.Type {
		case TypeInt8, TypeInt16, TypeInt32:
			field.Type = IntType().WithNewType(field.Type.Type)
		case TypeUint8, TypeUint16, TypeUint32:
			field.Type = UintType().WithNewType(field.Type.Type)
		}
	}
	if opt.EnableBoolInt && field.Type.IsBool() {
		field.Type = IntType().WithNewType(field.Type.Type)
	}
	if field.Nullable && opt.DisableNullToPoint {
		gt, ok := getSQLNullValueGoType(field.Type.Type)
		if ok {
			field.Type = gt
			field.GoPointer = false
		}
	}
	for tag, kind := range opt.Tags {
		if tag == "json" {
			if vv := matcher.JsonTag(field.Comment); vv != "" {
				field.Tags = append(field.Tags, fmt.Sprintf(`%s:"%s"`, tag, vv))
				continue
			}
		}
		vv := utils.StyleName(kind, field.ColumnName)
		if vv == "" {
			continue
		}
		format := ""
		if tag == "json" && matcher.HasAffixJSONTag(field.Comment) {
			if opt.IgnoreOmitempty {
				format = `%s:"%s,string"`
			} else {
				format = `%s:"%s,omitempty,string"`
			}
		} else if opt.IgnoreOmitempty {
			format = `%s:"%s"`
		} else {
			format = `%s:"%s,omitempty"`
		}
		field.Tags = append(field.Tags, fmt.Sprintf(format, tag, vv))
	}
	goName := field.GoName
	for {
		_, ok := escapeFieldNames[goName]
		if !ok { // need to escape
			break
		}
		goName = "X" + goName
		// 和当前字段存在的重复, 再追加一个
		_, ok = allFieldName[goName]
		if ok {
			goName = "X" + goName
		}
	}
	if field.GoName != goName {
		field.GoName = goName
		escapeFieldNames[goName] = struct{}{} // 添加为必须转义
	}
}