From 596843a62c6e4fca4fe6e7bebf4afa34594d1ca2 Mon Sep 17 00:00:00 2001
From: "felix.fengmin" <felix.fengmin@bytedance.com>
Date: Mon, 6 Mar 2023 11:13:47 +0800
Subject: [PATCH 1/2] feat: support json marshal on general objects

---
 .gitignore                |  1 +
 general/general_object.go | 41 ++++++++++++++++++++++++++++++++++++++-
 go.mod                    |  3 +++
 3 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 .gitignore
 create mode 100644 go.mod

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..485dee6
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+.idea
diff --git a/general/general_object.go b/general/general_object.go
index a3eb314..f86c3dd 100644
--- a/general/general_object.go
+++ b/general/general_object.go
@@ -1,6 +1,16 @@
 package general
 
-import "github.com/thrift-iterator/go/protocol"
+import (
+	"bytes"
+	"encoding/json"
+	"github.com/thrift-iterator/go/protocol"
+)
+
+var (
+	_ json.Marshaler = (*List)(nil)
+	_ json.Marshaler = (*Map)(nil)
+	_ json.Marshaler = (*Struct)(nil)
+)
 
 type Object interface {
 	Get(path ...interface{}) interface{}
@@ -19,6 +29,10 @@ func (obj List) Get(path ...interface{}) interface{} {
 	return elem.(Object).Get(path[1:]...)
 }
 
+func (l *List) MarshalJSON() ([]byte, error) {
+	return json.Marshal((*[]interface{})(l))
+}
+
 type Map map[interface{}]interface{}
 
 func (obj Map) Get(path ...interface{}) interface{} {
@@ -32,6 +46,27 @@ func (obj Map) Get(path ...interface{}) interface{} {
 	return elem.(Object).Get(path[1:]...)
 }
 
+func (m Map) MarshalJSON() ([]byte, error) {
+	if len(m) == 0 {
+		return []byte("{}"), nil
+	}
+	buf := bytes.NewBuffer([]byte("{"))
+	for k, v := range m {
+		buf.WriteString(`"`)
+		buf.WriteString(k.(string))
+		buf.WriteString(`":`)
+		b, err := json.Marshal(v)
+		if err != nil {
+			return nil, err
+		}
+		buf.Write(b)
+		buf.WriteString(",")
+	}
+	buf.Truncate(buf.Len() - 1)
+	buf.WriteString("}")
+	return buf.Bytes(), nil
+}
+
 type Struct map[protocol.FieldId]interface{}
 
 func (obj Struct) Get(path ...interface{}) interface{} {
@@ -45,6 +80,10 @@ func (obj Struct) Get(path ...interface{}) interface{} {
 	return elem.(Object).Get(path[1:]...)
 }
 
+func (s *Struct) MarshalJSON() ([]byte, error) {
+	return json.Marshal((*map[protocol.FieldId]interface{})(s))
+}
+
 type Message struct {
 	protocol.MessageHeader
 	Arguments Struct
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..d68272a
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,3 @@
+module github.com/thrift-iterator/go
+
+go 1.19

From dbebdf85ec7487ee9f2fadfb3249b31765081f6b Mon Sep 17 00:00:00 2001
From: "felix.fengmin" <felix.fengmin@bytedance.com>
Date: Mon, 6 Mar 2023 17:00:58 +0800
Subject: [PATCH 2/2] fix: key may be int64

---
 general/general_object.go | 42 +++++++++++++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/general/general_object.go b/general/general_object.go
index f86c3dd..20666a0 100644
--- a/general/general_object.go
+++ b/general/general_object.go
@@ -3,7 +3,10 @@ package general
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"github.com/thrift-iterator/go/protocol"
+	"reflect"
+	"strconv"
 )
 
 var (
@@ -52,12 +55,13 @@ func (m Map) MarshalJSON() ([]byte, error) {
 	}
 	buf := bytes.NewBuffer([]byte("{"))
 	for k, v := range m {
-		buf.WriteString(`"`)
-		buf.WriteString(k.(string))
-		buf.WriteString(`":`)
-		b, err := json.Marshal(v)
-		if err != nil {
-			return nil, err
+		i, errKey := writeKey(buf, k)
+		if errKey != nil {
+			return i, errKey
+		}
+		b, errValue := json.Marshal(v)
+		if errValue != nil {
+			return nil, errValue
 		}
 		buf.Write(b)
 		buf.WriteString(",")
@@ -67,6 +71,32 @@ func (m Map) MarshalJSON() ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
+func writeKey(buf *bytes.Buffer, k interface{}) ([]byte, error) {
+	buf.WriteString(`"`)
+	switch k.(type) {
+	case string:
+		buf.WriteString(k.(string))
+	case byte:
+		buf.WriteString(strconv.FormatInt(int64(k.(byte)), 10))
+	case int64:
+		buf.WriteString(strconv.FormatInt(k.(int64), 10))
+	case int32:
+		buf.WriteString(strconv.FormatInt(int64(k.(int32)), 10))
+	case int:
+		buf.WriteString(strconv.FormatInt(int64(k.(int)), 10))
+	case int16:
+		buf.WriteString(strconv.FormatInt(int64(k.(int16)), 10))
+	case float64:
+		buf.WriteString(strconv.FormatFloat(k.(float64), 'f', -1, 64))
+	case bool:
+		buf.WriteString(strconv.FormatBool(k.(bool)))
+	default:
+		return nil, errors.New("unsupported map key type " + reflect.TypeOf(k).String())
+	}
+	buf.WriteString(`":`)
+	return nil, nil
+}
+
 type Struct map[protocol.FieldId]interface{}
 
 func (obj Struct) Get(path ...interface{}) interface{} {