diff --git a/go.mod b/go.mod index 70a0b1b96..f0191bd49 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/bytedance/sonic/loader v0.2.4 github.com/cloudwego/base64x v0.1.5 github.com/davecgh/go-spew v1.1.1 + github.com/fatih/structtag v1.2.0 github.com/klauspost/cpuid/v2 v2.0.9 github.com/stretchr/testify v1.8.1 github.com/twitchyliquid64/golang-asm v0.15.1 diff --git a/go.sum b/go.sum index 7839159a9..ed7abace2 100644 --- a/go.sum +++ b/go.sum @@ -7,6 +7,8 @@ github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/fatih/structtag v1.2.0 h1:/OdNE99OxoI/PqaW/SuSK9uxxT3f/tcSZgon/ssNSx4= +github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94= github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= diff --git a/utils/tags.go b/utils/tags.go new file mode 100644 index 000000000..1de951e56 --- /dev/null +++ b/utils/tags.go @@ -0,0 +1,182 @@ +/** + * Copyright 2025 ByteDance Inc. + * + * 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 + * + * https://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 utils + +import ( + "fmt" + "reflect" + "unsafe" + + "github.com/bytedance/sonic" + "github.com/fatih/structtag" +) + +// TagChanger is used to change the tag of a struct. +// +// WARNING: it is not thread-safe, must be used in a single thread or in initialization phase. +type TagChanger struct { + types map[reflect.Type]reflect.Type // old -> new + replacer func(field reflect.StructField, tag reflect.StructTag) reflect.StructTag +} + +func NewTagChanger(replacer func(field reflect.StructField, tag reflect.StructTag) reflect.StructTag) *TagChanger { + return &TagChanger{ + types: make(map[reflect.Type]reflect.Type), + replacer: replacer, + } +} + +// Replace returns the type whose tag is replaced by replacer. +// +// WARNING: it is not thread-safe, must be used in a single thread or in initialization phase. +func (t *TagChanger) Replace(old reflect.Type) reflect.Type { + if t.types == nil { + t.types = make(map[reflect.Type]reflect.Type) + } + + if n, ok := t.types[old]; ok { + return n + } + + if old.Kind() == reflect.Struct { + newType := reflect.StructOf(t.createFieldsWithReplacedTags(old)) + t.types[old] = newType + return newType + } + + if old.Kind() == reflect.Map { + nn := t.Replace(old.Elem()) + n := reflect.MapOf(old.Key(), nn) + t.types[old] = n + return n + } + + if old.Kind() == reflect.Slice || old.Kind() == reflect.Array { + nn := t.Replace(old.Elem()) + n := reflect.SliceOf(nn) + t.types[old] = n + return n + } + + if old.Kind() == reflect.Ptr { + nn := t.Replace(old.Elem()) + n := reflect.PtrTo(nn) + t.types[old] = n + return n + } + + return old +} + +func (t *TagChanger) createFieldsWithReplacedTags(old reflect.Type) []reflect.StructField { + numField := old.NumField() + fields := make([]reflect.StructField, numField) + + for i := 0; i < numField; i++ { + oldField := old.Field(i) + tag := oldField.Tag + + fieldType := t.Replace(oldField.Type) + + newTag := t.replacer(oldField, tag) + + fields[i] = reflect.StructField{ + Name: oldField.Name, + Type: fieldType, + Tag: newTag, + } + } + + return fields +} + +// ReplacerAddOmitempty is a replacer that adds "omitempty" to the json tag. +func ReplacerAddOmitempty(field reflect.StructField, tag reflect.StructTag) reflect.StructTag { + tags, err := structtag.Parse(string(tag)) + if err != nil { + return tag + } + + jsonTag, err := tags.Get("json") + if jsonTag == nil { + if err := tags.Set(&structtag.Tag{Key: "json", Name: field.Name, Options: []string{"omitempty"}}); err != nil { + panic(err) + } + return reflect.StructTag(tags.String()) + } + + if !jsonTag.HasOption("omitempty") { + jsonTag.Options = append(jsonTag.Options, "omitempty") + } + + tags.Set(jsonTag) + newTag := tags.String() + return reflect.StructTag(newTag) +} + +func ReplaceerAddStringInt64(field reflect.StructField, tag reflect.StructTag) reflect.StructTag { + typ := field.Type + if typ.Kind() == reflect.Ptr { + typ = typ.Elem() + } + if typ.Kind() != reflect.Int64 && typ.Kind() != reflect.Uint64 { + return tag + } + + tags, err := structtag.Parse(string(tag)) + if err != nil { + return tag + } + + jsonTag, err := tags.Get("json") + if jsonTag == nil { + if err := tags.Set(&structtag.Tag{Key: "json", Name: field.Name, Options: []string{"string"}}); err != nil { + panic(err) + } + return reflect.StructTag(tags.String()) + } + + if !jsonTag.HasOption("string") { + jsonTag.Options = append(jsonTag.Options, "string") + } + + tags.Set(jsonTag) + newTag := tags.String() + return reflect.StructTag(newTag) +} + +// Marshal is used to marshal a struct whose tag is replaced by replacer into json. +func (t *TagChanger) Marshal(v interface{}) ([]byte, error) { + oldType := reflect.TypeOf(v) + if oldType.Kind() != reflect.Ptr { + return nil, fmt.Errorf("v must be a pointer") + } + newType := t.Replace(oldType.Elem()) + newObj := reflect.NewAt(newType, unsafe.Pointer(reflect.ValueOf(v).Pointer())).Interface() + return sonic.Marshal(newObj) +} + +// Unmarshal is used to unmarshal json into a struct whose tag is replaced by replacer. +func (t *TagChanger) Unmarshal(data []byte, v interface{}) error { + oldType := reflect.TypeOf(v) + if oldType.Kind()!= reflect.Ptr { + return fmt.Errorf("v must be a pointer") + } + newType := t.Replace(oldType.Elem()) + newObj := reflect.NewAt(newType, unsafe.Pointer(reflect.ValueOf(v).Pointer())).Interface() + return sonic.Unmarshal(data, newObj) +} diff --git a/utils/tags_test.go b/utils/tags_test.go new file mode 100644 index 000000000..bef5f77af --- /dev/null +++ b/utils/tags_test.go @@ -0,0 +1,170 @@ +/** + * Copyright 2025 ByteDance Inc. + * + * 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 + * + * https://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 utils + +import ( + "encoding/json" + "fmt" + "reflect" + "testing" + "unsafe" +) + +func TestTagChanger_ReplaceTag(t *testing.T) { + + type caseStruct struct { + A *int `xx:"yy"` + B map[int]struct { + C *int `json:"c" xx:"yy"` + } + D []struct { + E *int `json:"e"` + } + F *struct { + G *int `json:"g"` + } + H struct { + I *int `json:"i"` + } + } + + type fields struct { + types map[reflect.Type]reflect.Type + replacer func(field reflect.StructField, tag reflect.StructTag) reflect.StructTag + } + type args struct { + old interface{} + } + tests := []struct { + name string + args args + want string + }{ + { + name: "test1", + args: args{ + old: caseStruct{}, + }, + want: `{"H":{}}`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tr := NewTagChanger(ReplacerAddOmitempty) + oldObj := caseStruct{} + oldJs, err := json.Marshal(oldObj) + if err != nil { + t.Errorf("TagChanger.ReplaceTag() error = %v", err) + return + } + fmt.Println(string(oldJs)) + got := tr.Replace(reflect.TypeOf(tt.args.old)) + gotObj := reflect.NewAt(got, unsafe.Pointer(&oldObj)).Interface() + js, err := json.Marshal(gotObj) + println(string(js)) + if err != nil { + t.Errorf("TagChanger.ReplaceTag() error = %v", err) + return + } else if string(js) != tt.want { + t.Errorf("TagChanger.ReplaceTag() = %v, want %v", string(js), tt.want) + } + }) + } +} + +func TestTagChanger_Marshal(t *testing.T) { + type args struct { + v interface{} + } + tests := []struct { + name string + args args + want []byte + wantErr bool + }{ + { + name: "1", + args: args{ + v: &struct { + A *int `json:"a"` + B *string + }{}, + }, + want: []byte(`{}`), + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tr := NewTagChanger(ReplacerAddOmitempty) + got, err := tr.Marshal(tt.args.v) + if (err != nil) != tt.wantErr { + t.Errorf("TagChanger.Marshal() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !reflect.DeepEqual(got, tt.want) { + t.Errorf("TagChanger.Marshal() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestTagChanger_Unmarshal(t *testing.T) { + type args struct { + v interface{} + } + type testStruct struct { + A int64 `json:"a"` + B uint64 + C *int64 + } + vc := int64(3) + tests := []struct { + name string + args args + js []byte + want interface{} + wantErr bool + }{ + { + name: "1", + args: args{ + v: &testStruct{}, + }, + js: []byte(`{"a":"1","b":"2","c":"3"}`), + want: &testStruct{ + A: 1, + B: 2, + C: &vc, + }, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tr := NewTagChanger(ReplaceerAddStringInt64) + err := tr.Unmarshal(tt.js, tt.args.v) + if (err != nil) != tt.wantErr { + t.Errorf("TagChanger.Marshal() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !reflect.DeepEqual(tt.args.v, tt.want) { + t.Errorf("TagChanger.Marshal() = %v, want %v", tt.args.v, tt.want) + } + }) + } +}