improvement: optimization of readmap function

Signed-off-by: thomas-dg <dyhshome@126.com>
apache · Jul 3, 2024 · 9fa6865 · 9fa6865
1 parent 67073b7
commit 9fa6865
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/gremlin-go/driver/graphBinary.go b/gremlin-go/driver/graphBinary.go
@@ -992,7 +992,12 @@ func readMap(data *[]byte, i *int) (interface{}, error) {
 		} else if reflect.TypeOf(k).Comparable() {
 			mapData[k] = v
 		} else {
-			mapData[fmt.Sprint(k)] = v
+			switch reflect.TypeOf(k).Kind() {
+			case reflect.Map:
+				mapData[&k] = v
+			default:
+				mapData[fmt.Sprint(k)] = v
+			}
 		}
 	}
 	return mapData, nil

diff --git a/gremlin-go/driver/graphBinary_test.go b/gremlin-go/driver/graphBinary_test.go
@@ -232,12 +232,12 @@ func TestGraphBinaryV1(t *testing.T) {
 			assert.Nil(t, err)
 			assert.Equal(t, fmt.Sprintf("%v", source), fmt.Sprintf("%v", res))
 		})
-		t.Run("read incomparable map", func(t *testing.T) {
+		t.Run("read incomparable map: a map value as the key", func(t *testing.T) {
 			// prepare test data
 			var buf = &bytes.Buffer{}
 			typeSerializer := &graphBinaryTypeSerializer{}
 			// write the size of map
-			err := binary.Write(buf, binary.BigEndian, uint32(2))
+			err := binary.Write(buf, binary.BigEndian, uint32(1))
 			if err != nil {
 				t.Fatalf("Failed to write data: %v", err)
 			}
@@ -252,6 +252,31 @@ func TestGraphBinaryV1(t *testing.T) {
 			if err != nil {
 				t.Fatalf("Failed to encode data: %v", err)
 			}
+
+			data := buf.Bytes()
+			i := 0
+			result, err := readMap(&data, &i)
+			if err != nil {
+				t.Fatalf("readMap failed: %v", err)
+			}
+			mResult, ok := result.(map[interface{}]interface{})
+			if !ok {
+				t.Fatalf("readMap result not map[interface{}]interface{}")
+			}
+			for k, v := range mResult {
+				assert.Equal(t, reflect.Ptr, reflect.TypeOf(k).Kind())
+				assert.Equal(t, "value1", v)
+			}
+		})
+		t.Run("read incomparable map: a slice value as the key", func(t *testing.T) {
+			// prepare test data
+			var buf = &bytes.Buffer{}
+			typeSerializer := &graphBinaryTypeSerializer{}
+			// write the size of map
+			err := binary.Write(buf, binary.BigEndian, uint32(1))
+			if err != nil {
+				t.Fatalf("Failed to write data: %v", err)
+			}
 			// write a slice value as the key
 			k2 := []int{1, 2, 3}
 			_, err = typeSerializer.write(reflect.ValueOf(k2).Interface(), buf)
@@ -271,8 +296,7 @@ func TestGraphBinaryV1(t *testing.T) {
 				t.Fatalf("readMap failed: %v", err)
 			}
 			expected := map[interface{}]interface{}{
-				"map[key:value]": "value1",
-				"[1 2 3]":        "value2",
+				"[1 2 3]": "value2",
 			}
 			if !reflect.DeepEqual(result, expected) {
 				t.Errorf("Expected %v, but got %v", expected, result)