Skip to content

Commit 8d4888b

Browse files
AsterDYAsterDY
committed
feat: Descriptor support marshal json
1 parent d4d60c0 commit 8d4888b

File tree

3 files changed

+368
-3
lines changed

3 files changed

+368
-3
lines changed

‎trim/desc.go‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
package trim
1818

1919
import (
20+
"encoding/json"
21+
"fmt"
2022
"strings"
2123
"sync/atomic"
2224
)
@@ -154,3 +156,141 @@ func (d *Descriptor) String() string {
154156
printer(d, "")
155157
return sb.String()
156158
}
159+
160+
// descriptorJSON is the JSON representation of Descriptor
161+
type descriptorJSON struct {
162+
Kind TypeKind `json:"kind"`
163+
Name string `json:"name"`
164+
Children []fieldJSON `json:"children,omitempty"`
165+
}
166+
167+
// fieldJSON is the JSON representation of Field
168+
type fieldJSON struct {
169+
Name string `json:"name"`
170+
ID int `json:"id"`
171+
Desc *descriptorJSON `json:"desc,omitempty"`
172+
Ref string `json:"$ref,omitempty"` // reference to another descriptor by path
173+
}
174+
175+
// MarshalJSON implements json.Marshaler interface for Descriptor
176+
// It handles circular references by using $ref to reference already visited descriptors
177+
func (d *Descriptor) MarshalJSON() ([]byte, error) {
178+
visited := make(map[*Descriptor]string) // maps pointer to path
179+
result := d.marshalWithPath("$", visited)
180+
return json.Marshal(result)
181+
}
182+
183+
// marshalWithPath recursively marshals the descriptor, tracking visited nodes
184+
func (d *Descriptor) marshalWithPath(path string, visited map[*Descriptor]string) *descriptorJSON {
185+
if d == nil {
186+
return nil
187+
}
188+
189+
// Check if we've already visited this descriptor (circular reference)
190+
if existingPath, ok := visited[d]; ok {
191+
// Return a reference placeholder - this will be handled specially
192+
return &descriptorJSON{
193+
Kind: d.Kind,
194+
Name: fmt.Sprintf("$ref:%s", existingPath),
195+
}
196+
}
197+
198+
// Mark as visited
199+
visited[d] = path
200+
201+
result := &descriptorJSON{
202+
Kind: d.Kind,
203+
Name: d.Name,
204+
Children: make([]fieldJSON, 0, len(d.Children)),
205+
}
206+
207+
for i, f := range d.Children {
208+
childPath := fmt.Sprintf("%s.children[%d].desc", path, i)
209+
fj := fieldJSON{
210+
Name: f.Name,
211+
ID: f.ID,
212+
}
213+
214+
if f.Desc != nil {
215+
// Check if child descriptor was already visited
216+
if existingPath, ok := visited[f.Desc]; ok {
217+
fj.Ref = existingPath
218+
} else {
219+
fj.Desc = f.Desc.marshalWithPath(childPath, visited)
220+
}
221+
}
222+
223+
result.Children = append(result.Children, fj)
224+
}
225+
226+
return result
227+
}
228+
229+
// UnmarshalJSON implements json.Unmarshaler interface for Descriptor
230+
// It handles circular references by resolving $ref references after initial parsing
231+
func (d *Descriptor) UnmarshalJSON(data []byte) error {
232+
var raw descriptorJSON
233+
if err := json.Unmarshal(data, &raw); err != nil {
234+
return err
235+
}
236+
237+
// First pass: build all descriptors and collect references
238+
refs := make(map[string]*Descriptor) // path -> descriptor
239+
d.unmarshalFromJSON(&raw, "$", refs)
240+
241+
// Second pass: resolve references
242+
d.resolveRefs("$", refs)
243+
244+
return nil
245+
}
246+
247+
// unmarshalFromJSON populates the descriptor from JSON representation
248+
func (d *Descriptor) unmarshalFromJSON(raw *descriptorJSON, path string, refs map[string]*Descriptor) {
249+
d.Kind = raw.Kind
250+
d.Name = raw.Name
251+
d.Children = make([]Field, 0, len(raw.Children))
252+
d.ids = nil
253+
d.names = nil
254+
255+
// Register this descriptor
256+
refs[path] = d
257+
258+
for i, fj := range raw.Children {
259+
childPath := fmt.Sprintf("%s.children[%d].desc", path, i)
260+
f := Field{
261+
Name: fj.Name,
262+
ID: fj.ID,
263+
}
264+
265+
if fj.Ref != "" {
266+
// This is a reference, will be resolved later
267+
// Create a placeholder descriptor with special name
268+
f.Desc = &Descriptor{Name: "$ref:" + fj.Ref}
269+
} else if fj.Desc != nil {
270+
f.Desc = &Descriptor{}
271+
f.Desc.unmarshalFromJSON(fj.Desc, childPath, refs)
272+
}
273+
274+
d.Children = append(d.Children, f)
275+
}
276+
}
277+
278+
// resolveRefs resolves all $ref references in the descriptor tree
279+
func (d *Descriptor) resolveRefs(path string, refs map[string]*Descriptor) {
280+
for i := range d.Children {
281+
if d.Children[i].Desc != nil {
282+
childPath := fmt.Sprintf("%s.children[%d].desc", path, i)
283+
284+
// Check if this is a reference
285+
if strings.HasPrefix(d.Children[i].Desc.Name, "$ref:") {
286+
refPath := strings.TrimPrefix(d.Children[i].Desc.Name, "$ref:")
287+
if target, ok := refs[refPath]; ok {
288+
d.Children[i].Desc = target
289+
}
290+
} else {
291+
// Recursively resolve references
292+
d.Children[i].Desc.resolveRefs(childPath, refs)
293+
}
294+
}
295+
}
296+
}

‎trim/desc_test.go‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
/**
2+
* Copyright 2025 ByteDance Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package trim
18+
19+
import (
20+
"encoding/json"
21+
"testing"
22+
23+
"github.com/stretchr/testify/require"
24+
)
25+
26+
func TestDescriptorMarshalJSON(t *testing.T) {
27+
// Create a simple descriptor without circular reference
28+
desc := &Descriptor{
29+
Kind: TypeKind_Struct,
30+
Name: "Root",
31+
Children: []Field{
32+
{
33+
Name: "field1",
34+
ID: 1,
35+
Desc: &Descriptor{
36+
Kind: TypeKind_Scalar,
37+
Name: "Leaf1",
38+
},
39+
},
40+
{
41+
Name: "field2",
42+
ID: 2,
43+
Desc: &Descriptor{
44+
Kind: TypeKind_StrMap,
45+
Name: "Map1",
46+
Children: []Field{
47+
{
48+
Name: "key1",
49+
ID: 0,
50+
Desc: &Descriptor{
51+
Kind: TypeKind_Scalar,
52+
Name: "Leaf2",
53+
},
54+
},
55+
},
56+
},
57+
},
58+
},
59+
}
60+
61+
// Marshal to JSON
62+
data, err := json.Marshal(desc)
63+
require.NoError(t, err)
64+
t.Logf("JSON: %s", string(data))
65+
66+
// Unmarshal back
67+
var desc2 Descriptor
68+
err = json.Unmarshal(data, &desc2)
69+
require.NoError(t, err)
70+
71+
// Verify structure
72+
require.Equal(t, desc.Kind, desc2.Kind)
73+
require.Equal(t, desc.Name, desc2.Name)
74+
require.Len(t, desc2.Children, 2)
75+
require.Equal(t, desc.Children[0].Name, desc2.Children[0].Name)
76+
require.Equal(t, desc.Children[0].ID, desc2.Children[0].ID)
77+
require.NotNil(t, desc2.Children[0].Desc)
78+
require.Equal(t, desc.Children[0].Desc.Name, desc2.Children[0].Desc.Name)
79+
}
80+
81+
func TestDescriptorMarshalJSONWithCircularReference(t *testing.T) {
82+
// Create a descriptor with circular reference
83+
root := &Descriptor{
84+
Kind: TypeKind_Struct,
85+
Name: "Root",
86+
}
87+
88+
child := &Descriptor{
89+
Kind: TypeKind_Struct,
90+
Name: "Child",
91+
}
92+
93+
// Root -> Child -> Root (circular reference)
94+
root.Children = []Field{
95+
{
96+
Name: "child",
97+
ID: 1,
98+
Desc: child,
99+
},
100+
}
101+
102+
child.Children = []Field{
103+
{
104+
Name: "parent",
105+
ID: 1,
106+
Desc: root, // circular reference back to root
107+
},
108+
}
109+
110+
// Marshal to JSON - should not panic or infinite loop
111+
data, err := json.Marshal(root)
112+
require.NoError(t, err)
113+
t.Logf("JSON with circular ref: %s", string(data))
114+
115+
// Unmarshal back
116+
var root2 Descriptor
117+
err = json.Unmarshal(data, &root2)
118+
require.NoError(t, err)
119+
120+
// Verify structure
121+
require.Equal(t, root.Kind, root2.Kind)
122+
require.Equal(t, root.Name, root2.Name)
123+
require.Len(t, root2.Children, 1)
124+
require.NotNil(t, root2.Children[0].Desc)
125+
require.Equal(t, "Child", root2.Children[0].Desc.Name)
126+
127+
// Verify circular reference is resolved
128+
require.NotNil(t, root2.Children[0].Desc.Children)
129+
require.Len(t, root2.Children[0].Desc.Children, 1)
130+
require.NotNil(t, root2.Children[0].Desc.Children[0].Desc)
131+
// The circular reference should point back to root2
132+
require.Equal(t, &root2, root2.Children[0].Desc.Children[0].Desc)
133+
}
134+
135+
func TestDescriptorMarshalJSONSelfReference(t *testing.T) {
136+
// Create a descriptor that references itself
137+
self := &Descriptor{
138+
Kind: TypeKind_Struct,
139+
Name: "Self",
140+
}
141+
142+
self.Children = []Field{
143+
{
144+
Name: "self",
145+
ID: 1,
146+
Desc: self, // self reference
147+
},
148+
}
149+
150+
// Marshal to JSON
151+
data, err := json.Marshal(self)
152+
require.NoError(t, err)
153+
t.Logf("JSON with self ref: %s", string(data))
154+
155+
// Unmarshal back
156+
var self2 Descriptor
157+
err = json.Unmarshal(data, &self2)
158+
require.NoError(t, err)
159+
160+
// Verify self reference is resolved
161+
require.Equal(t, &self2, self2.Children[0].Desc)
162+
}
163+
164+
func TestDescriptorMarshalJSONNil(t *testing.T) {
165+
// Descriptor with nil child
166+
desc := &Descriptor{
167+
Kind: TypeKind_Struct,
168+
Name: "Root",
169+
Children: []Field{
170+
{
171+
Name: "field1",
172+
ID: 1,
173+
Desc: nil, // nil descriptor
174+
},
175+
},
176+
}
177+
178+
data, err := json.Marshal(desc)
179+
require.NoError(t, err)
180+
t.Logf("JSON with nil: %s", string(data))
181+
182+
var desc2 Descriptor
183+
err = json.Unmarshal(data, &desc2)
184+
require.NoError(t, err)
185+
186+
require.Nil(t, desc2.Children[0].Desc)
187+
}
188+
189+
func TestDescriptorMarshalJSONMultipleReferences(t *testing.T) {
190+
// Create a shared descriptor referenced by multiple fields
191+
shared := &Descriptor{
192+
Kind: TypeKind_Scalar,
193+
Name: "Shared",
194+
}
195+
196+
root := &Descriptor{
197+
Kind: TypeKind_Struct,
198+
Name: "Root",
199+
Children: []Field{
200+
{
201+
Name: "ref1",
202+
ID: 1,
203+
Desc: shared,
204+
},
205+
{
206+
Name: "ref2",
207+
ID: 2,
208+
Desc: shared, // same descriptor referenced again
209+
},
210+
},
211+
}
212+
213+
data, err := json.Marshal(root)
214+
require.NoError(t, err)
215+
t.Logf("JSON with multiple refs: %s", string(data))
216+
217+
var root2 Descriptor
218+
err = json.Unmarshal(data, &root2)
219+
require.NoError(t, err)
220+
221+
// Both refs should point to the same descriptor after unmarshaling
222+
require.NotNil(t, root2.Children[0].Desc)
223+
require.NotNil(t, root2.Children[1].Desc)
224+
require.Equal(t, root2.Children[0].Desc, root2.Children[1].Desc)
225+
}

‎trim/fetch.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -97,12 +97,12 @@ func getStructFieldInfo(t reflect.Type) *structFieldInfo {
9797
}
9898

9999
// Parse thrift tag: "FieldName,ID" - use IndexByte for better performance
100-
idx := strings.IndexByte(tag, ',')
101-
if idx < 0 {
100+
idx := strings.Split(tag, ",")
101+
if len(idx) < 2 {
102102
continue
103103
}
104104

105-
fieldID, err := strconv.Atoi(tag[idx+1:])
105+
fieldID, err := strconv.Atoi(idx[1])
106106
if err != nil {
107107
continue
108108
}

0 commit comments

Comments
 (0)