diff --git a/generator/golang/option.go b/generator/golang/option.go index e894f57..f8dcb4d 100644 --- a/generator/golang/option.go +++ b/generator/golang/option.go @@ -52,8 +52,8 @@ type Features struct { EnumAsINT32 bool `enum_as_int_32:"Generate enum type as int32"` CodeRefSlim bool `code_ref_slim:"Generate code ref by given idl-ref.yaml with less refs to avoid conflict"` CodeRef bool `code_ref:"Generate code ref by given idl-ref.yaml"` - ExpCodeRef bool `exp_code_ref:"Generate code ref by given idl-ref.yaml with less refs to avoid conflict, but remind some struct as local.( this is a exp feature )"` - KeepCodeRefName bool `keep_code_ref_name:"Generate code ref but still keep file name."` + ExpCodeRef bool `exp_code_ref:"Generate code ref by given idl-ref.yaml with less refs to avoid conflict, but remind some struct as local.( this is a exp feature )"` + KeepCodeRefName bool `keep_code_ref_name:"Generate code ref but still keep file name."` TrimIDL bool `trim_idl:"Simplify IDL to the most concise form before generating code."` EnableNestedStruct bool `enable_nested_struct:"Generate nested field when 'thrift.nested=\"true\"' annotation is set to field, valid only in 'slim and raw_struct template'"` JSONStringer bool `json_stringer:"Generate the JSON marshal method in String() method."` @@ -70,6 +70,7 @@ type Features struct { SkipEmpty bool `skip_empty:"If there's not content in file, just skip it. Later this feature will be a default feature."` NoProcessor bool `no_processor:" Do not generate default thrift processor and client. Later this feature will be a default feature."` GetEnumAnnotation bool `get_enum_annotation:"Generate GetAnnotation method for enum types."` + GenMapForIDLSet bool `gen_map_for_idl_set:"Generate map[type]struct{} instead of []type for IDL set."` } var defaultFeatures = Features{ @@ -108,6 +109,7 @@ var defaultFeatures = Features{ NoAliasTypeReflectionMethod: false, EnableRefInterface: false, GetEnumAnnotation: false, + GenMapForIDLSet: false, } type param struct { diff --git a/generator/golang/read_write_context.go b/generator/golang/read_write_context.go index 045d447..a2929bd 100644 --- a/generator/golang/read_write_context.go +++ b/generator/golang/read_write_context.go @@ -39,6 +39,8 @@ type ReadWriteContext struct { ids map[string]int // Prefix => local variable index FieldMask string + + IsGenMap bool // whether set field would be generated as map } // GenID returns a local variable with the given name as prefix. @@ -132,5 +134,13 @@ func mkRWCtx(r *Resolver, s *Scope, t *parser.Type, top *ReadWriteContext) (*Rea } } + if r.util.Features().GenMapForIDLSet && t.Category == parser.Category_Set { + ctx.IsGenMap = true + // since we can not use []byte as the key of map, convert it to string + if ctx.ValCtx.TypeID == typeids.Binary { + ctx.ValCtx.TypeName = "string" + } + } + return ctx, nil } diff --git a/generator/golang/resolver.go b/generator/golang/resolver.go index 0b09415..d4ea133 100644 --- a/generator/golang/resolver.go +++ b/generator/golang/resolver.go @@ -123,16 +123,23 @@ func (r *Resolver) getTypeName(g *Scope, t *parser.Type) (name string, err error } func (r *Resolver) getContainerTypeName(g *Scope, t *parser.Type) (name string, err error) { - if t.Name == "map" { + // when isGenMapForIDLSet is true, the generated code for set would be map[type]struct{} + isGenMapForIDLSet := r.util.Features().GenMapForIDLSet && t.Name == "set" + if t.Name == "map" || isGenMapForIDLSet { var k string - if t.KeyType.Category == parser.Category_Binary { + keyType := t.KeyType + if t.Name == "set" { + keyType = t.ValueType + } + + if keyType.Category == parser.Category_Binary { k = "string" // 'binary => string' for key type in map } else { - k, err = r.getTypeName(g, t.KeyType) + k, err = r.getTypeName(g, keyType) if err != nil { return "", fmt.Errorf("resolve key type of '%s' failed: %w", t, err) } - if t.KeyType.Category.IsStructLike() && !checkRefInterfaceType(r.util, g, t.KeyType) { + if keyType.Category.IsStructLike() && !checkRefInterfaceType(r.util, g, keyType) { // when a struct-like is used as key of a map, it must // generte a pointer type instead of the struct itself k = "*" + k @@ -143,14 +150,19 @@ func (r *Resolver) getContainerTypeName(g *Scope, t *parser.Type) (name string, name = "[]" // sets and lists compile into slices } - v, err := r.getTypeName(g, t.ValueType) - if err != nil { - return "", fmt.Errorf("resolve value type of '%s' failed: %w", t, err) + var v string + if isGenMapForIDLSet { + v = "struct{}" + } else { + v, err = r.getTypeName(g, t.ValueType) + if err != nil { + return "", fmt.Errorf("resolve value type of '%s' failed: %w", t, err) + } + if t.ValueType.Category.IsStructLike() && !r.util.Features().ValueTypeForSIC && !checkRefInterfaceType(r.util, g, t.ValueType) { + v = "*" + v // generate pointer type for struct-like by default + } } - if t.ValueType.Category.IsStructLike() && !r.util.Features().ValueTypeForSIC && !checkRefInterfaceType(r.util, g, t.ValueType) { - v = "*" + v // generate pointer type for struct-like by default - } return name + v, nil // map[k]v or []v } @@ -348,7 +360,11 @@ func (r *Resolver) onSetOrList(g *Scope, name string, t *parser.Type, v *parser. if err != nil { return "", err } - ss = append(ss, str+",") + if !r.util.Features().GenMapForIDLSet { + ss = append(ss, str+",") + } else { + ss = append(ss, fmt.Sprintf("%s: struct{}{},", str)) + } } if len(ss) == 0 { return goType + "{}", nil diff --git a/generator/golang/templates/struct.go b/generator/golang/templates/struct.go index bcb8187..bfe478e 100644 --- a/generator/golang/templates/struct.go +++ b/generator/golang/templates/struct.go @@ -675,13 +675,18 @@ var FieldReadMap = ` var FieldReadSet = ` {{define "FieldReadSet"}} {{- $isBaseVal := .ValCtx.Type | IsBaseType -}} +{{- $isGenMap := .IsGenMap -}} {{- $curFieldMask := .FieldMask -}} {{- $isStructVal := .ValCtx.Type.Category.IsStructLike -}} _, size, err := iprot.ReadSetBegin() if err != nil { return err } + {{- if $isGenMap}} + {{.Target}} {{if .NeedDecl}}:{{end}}= make({{.TypeName}}, size) + {{- else}} {{.Target}} {{if .NeedDecl}}:{{end}}= make({{.TypeName}}, 0, size) + {{- end}} {{- if $isStructVal}} values := make([]{{.ValCtx.TypeName.Deref}}, size) {{- end}} @@ -708,8 +713,12 @@ var FieldReadSet = ` {{if and .ValCtx.Type.Category.IsStructLike Features.ValueTypeForSIC}} {{$val = printf "*%s" $val}} {{end}} - + + {{- if $isGenMap}} + {{.Target}}[{{$val}}] = struct{}{} + {{- else}} {{.Target}} = append({{.Target}}, {{$val}}) + {{- end}} {{- if Features.WithFieldMask}} } {{- end}} @@ -904,6 +913,7 @@ var FieldWriteSet = ` {{define "FieldWriteSet"}} {{- $isBaseVal := .ValCtx.Type | IsBaseType -}} {{- $curFieldMask := .FieldMask -}} +{{- $isGenMap := .IsGenMap -}} {{- if Features.WithFieldMask}} if !{{.FieldMask}}.All() { l := len({{.Target}}) @@ -931,7 +941,7 @@ var FieldWriteSet = ` return err } {{- end}} - {{- if Features.ValidateSet}} + {{- if and Features.ValidateSet (not $isGenMap)}} {{- $ctx := (.ValCtx.WithTarget "tgt").WithSource "src"}} for i := 0; i < len({{.Target}}); i++ { for j := i + 1; j < len({{.Target}}); j++ { @@ -950,7 +960,7 @@ var FieldWriteSet = ` } } {{- end}} - for {{if Features.WithFieldMask}}i{{else}}_{{end}}, v := range {{.Target}} { + for {{if not $isGenMap}}{{if Features.WithFieldMask }}i{{else}}_{{end}}, {{end}}v := range {{.Target}} { {{- if Features.WithFieldMask}} {{- $curFieldMask = "nfm"}} if {{if $isBaseVal}}_{{else}}{{$curFieldMask}}{{end}}, ex := {{.FieldMask}}.Int(i); !ex { diff --git a/generator/golang/util.go b/generator/golang/util.go index c849ce8..d17d3e1 100644 --- a/generator/golang/util.go +++ b/generator/golang/util.go @@ -97,7 +97,7 @@ func (cu *CodeUtils) Features() Features { return cu.features } -// GetPackagePrefix sets the package prefix in generated codes. +// GetPackagePrefix gets the package prefix in generated codes. func (cu *CodeUtils) GetPackagePrefix() (pp string) { return cu.packagePrefix }