一般我们会起一个main函数的go主文件来写逻辑用于生成代码，但有的时候不太方便在一个包里面有main函数，而生成代码的逻辑又是必须的且最好是就在同一个包内或临近包中，此时，可以利用Test的编译机制，将我们需要的代码生成逻辑写在TestAbc函数中。
// file: github.com/cloudwego/frugal/internal/reflect/append_gen.sh
#!/bin/bash
FRUGAL_GEN_APPEND_MAP_FILE="append_map_gen.go"
FRUGAL_GEN_APPEND_LIST_FILE="append_list_gen.go"
rm -f $FRUGAL_GEN_APPEND_MAP_FILE
rm -f $FRUGAL_GEN_APPEND_LIST_FILE
exec go test -v -run=TestGenAppend -gencode=true
// file: github.com/cloudwego/frugal/internal/reflect/append_gen_test.go
package reflect
var (
gencode = flag.Bool("gencode", false, "generate list/map code for better performance")
)
// file: github.com/cloudwego/frugal/internal/reflect/append_list_gen_test.go
package reflect
const appendListFileName = "append_list_gen.go"
func TestGenAppendListCode(t *testing.T) {
if *gencode {
genAppendListCode(t, appendListFileName)
return
}
// ...other logic...
}
func genAppendListCode(t *testing.T, filename string) {
defineErr := map[ttype]bool{tOTHER: true}
defineStr := map[ttype]bool{tSTRING: true}
f := &bytes.Buffer{}
f.WriteString(appendListGenFileHeader)
// func init
fmt.Fprintln(f, "func init() {")
supportTypes := []ttype{
tBYTE, tI16, tI32, tI64, tDOUBLE,
tENUM, tSTRING, tSTRUCT, tMAP, tSET, tLIST,
}
t2var := map[ttype]string{
tBYTE: "tBYTE", tI16: "tI16", tI32: "tI32", tI64: "tI64", tDOUBLE: "tDOUBLE",
tENUM: "tENUM", tSTRING: "tSTRING",
tSTRUCT: "tSTRUCT", tMAP: "tMAP", tSET: "tSET", tLIST: "tLIST",
}
for _, v := range supportTypes {
fmt.Fprintf(f, "registerListAppendFunc(%s, %s)\n",
t2var[v], appendListFuncName(v))
}
fmt.Fprintln(f, "}")
fmt.Fprintln(f, "")
// func appendList_XXX
for _, v := range []ttype{tBYTE, tI16, tI32, tI64, tENUM, tSTRING, tOTHER} {
fmt.Fprintf(f, "func %s(t *tType, b []byte, p unsafe.Pointer) ([]byte, error) {\n",
appendListFuncName(v))
fmt.Fprintln(f, "t = t.V")
fmt.Fprintln(f, "b, n, vp := appendListHeader(t, b, p)")
fmt.Fprintln(f, "if n == 0 { return b, nil }")
if defineErr[v] {
fmt.Fprintln(f, "var err error")
} else if defineStr[v] {
fmt.Fprintln(f, "var s string")
}
fmt.Fprintln(f, "for i := uint32(0); i < n; i++ {")
fmt.Fprintln(f, "if i != 0 { vp = unsafe.Add(vp, t.Size) }")
fmt.Fprintln(f, getAppendCode(v, "t", "vp"))
fmt.Fprintln(f, "}")
fmt.Fprintln(f, "return b, nil")
fmt.Fprintln(f, "}")
fmt.Fprintln(f, "")
}
fileb, err := format.Source(f.Bytes())
if err != nil {
t.Log(codeWithLine(f.Bytes()))
t.Fatal(err)
}
err = os.WriteFile(filename, fileb, 0o644)
if err != nil {
t.Fatal(err)
}
t.Logf("generated: %s", filename)
}
// file: github.com/cloudwego/frugal/internal/reflect/append_map_gen_test.go
package reflect
const appendMapFileName = "append_map_gen.go"
func TestGenAppendMapCode(t *testing.T) {
if *gencode {
genAppendMapCode(t, appendMapFileName)
return
}
// ...other logic...
}
func genAppendMapCode(t *testing.T, filename string) {
f := &bytes.Buffer{}
f.WriteString(appendMapGenFileHeader)
// func init
fmt.Fprintln(f, "func init() {")
supportTypes := []ttype{
tBYTE, tI16, tI32, tI64, tDOUBLE,
tENUM, tSTRING, tSTRUCT, tMAP, tSET, tLIST,
}
t2var := map[ttype]string{
tBYTE: "tBYTE", tI16: "tI16", tI32: "tI32", tI64: "tI64", tDOUBLE: "tDOUBLE",
tENUM: "tENUM", tSTRING: "tSTRING",
tSTRUCT: "tSTRUCT", tMAP: "tMAP", tSET: "tSET", tLIST: "tLIST",
}
for _, k := range supportTypes {
for _, v := range supportTypes {
fmt.Fprintf(f, "registerMapAppendFunc(%s, %s, %s)\n",
t2var[k], t2var[v], appendMapFuncName(k, v))
}
}
fmt.Fprintln(f, "}")
fmt.Fprintln(f, "")
// func appendMapXXX
for _, k := range []ttype{tBYTE, tI16, tI32, tI64, tENUM, tSTRING, tOTHER} {
for _, v := range []ttype{tBYTE, tI16, tI32, tI64, tENUM, tSTRING, tOTHER} {
fmt.Fprintf(f, "func %s(t *tType, b []byte, p unsafe.Pointer) ([]byte, error) {\n",
appendMapFuncName(k, v))
fmt.Fprintln(f, "b, n := appendMapHeader(t, b, p)")
fmt.Fprintln(f, "if n == 0 { return b, nil }")
if defineErr[k] || defineErr[v] {
fmt.Fprintln(f, "var err error")
}
if defineStr[k] || defineStr[v] {
fmt.Fprintln(f, "var s string")
}
fmt.Fprintln(f, "it := newMapIter(rvWithPtr(t.RV, p))")
fmt.Fprintln(f, "for kp, vp := it.Next(); kp != nil;kp, vp = it.Next() {")
fmt.Fprintln(f, "n--")
fmt.Fprintln(f, getAppendCode(k, "t.K", "kp"))
fmt.Fprintln(f, getAppendCode(v, "t.V", "vp"))
fmt.Fprintln(f, "}")
fmt.Fprintln(f, "return b, checkMapN(n)")
fmt.Fprintln(f, "}")
fmt.Fprintln(f, "")
}
}
fileb, err := format.Source(f.Bytes())
if err != nil {
t.Log(codeWithLine(f.Bytes()))
t.Fatal(err)
}
err = os.WriteFile(filename, fileb, 0o644)
if err != nil {
t.Fatal(err)
}
t.Logf("generated: %s", filename)
}
这种法子我也是首见，确实是妙，如果是使用一个单独的辅组包专门用于生成所需的代码，也不是不行，就是看着别扭，毕竟frugal是一个库，终端用户是导入这个库然后使用，并不是一个cli可执行程序，潜意识里如果这个是一个go功能库，就不应该或者最好不要存在有main函数。