From dad8c0f4fd73fc31b5618a7fc590bbb2f063da73 Mon Sep 17 00:00:00 2001
From: thomas-dg <dyhshome@126.com>
Date: Tue, 2 Jul 2024 21:52:42 +0800
Subject: [PATCH] improvement: optimization of readmap function

Signed-off-by: thomas-dg <dyhshome@126.com>
---
 CHANGELOG.asciidoc                    |  4 ++
 gremlin-go/driver/graphBinary.go      |  6 +--
 gremlin-go/driver/graphBinary_test.go | 70 +++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc
index eada449247e..fad5620cf79 100644
--- a/CHANGELOG.asciidoc
+++ b/CHANGELOG.asciidoc
@@ -33,6 +33,10 @@ This release also includes changes from <<release-3-6-8, 3.6.8>>.
 * Fixed so that TrimGlobalStep and TrimLocalStep have the same character control handling as Ltrim and Rtrim
 * Fix a bug in MaxLocalStep, MinLocalStep, MeanLocalStep and SumLocalStep that it throws NoSuchElementException when encounters an empty iterator as input.
 
+==== Improvements
+
+* TINKERPOP-3093 optimization of readmap function
+
 [[release-3-7-2]]
 === TinkerPop 3.7.2 (April 8, 2024)
 
diff --git a/gremlin-go/driver/graphBinary.go b/gremlin-go/driver/graphBinary.go
index f9eb54850b2..7c95b01ebe6 100644
--- a/gremlin-go/driver/graphBinary.go
+++ b/gremlin-go/driver/graphBinary.go
@@ -989,14 +989,14 @@ func readMap(data *[]byte, i *int) (interface{}, error) {
 		}
 		if k == nil {
 			mapData[nil] = v
+		} else if reflect.TypeOf(k).Comparable() {
+			mapData[k] = v
 		} else {
 			switch reflect.TypeOf(k).Kind() {
 			case reflect.Map:
 				mapData[&k] = v
-			case reflect.Slice:
-				mapData[fmt.Sprint(k)] = v
 			default:
-				mapData[k] = v
+				mapData[fmt.Sprint(k)] = v
 			}
 		}
 	}
diff --git a/gremlin-go/driver/graphBinary_test.go b/gremlin-go/driver/graphBinary_test.go
index 5bfcf687fb8..eda45d6dd07 100644
--- a/gremlin-go/driver/graphBinary_test.go
+++ b/gremlin-go/driver/graphBinary_test.go
@@ -232,6 +232,76 @@ 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: 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(1))
+			if err != nil {
+				t.Fatalf("Failed to write data: %v", err)
+			}
+			// write a map value as the key
+			k1 := map[string]string{"key": "value"}
+			_, err = typeSerializer.write(reflect.ValueOf(k1).Interface(), buf)
+			if err != nil {
+				t.Fatalf("Failed to encode data: %v", err)
+			}
+			v1 := "value1"
+			_, err = typeSerializer.write(reflect.ValueOf(v1).Interface(), buf)
+			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)
+			if err != nil {
+				t.Fatalf("Failed to encode data: %v", err)
+			}
+			v2 := "value2"
+			_, err = typeSerializer.write(reflect.ValueOf(v2).Interface(), buf)
+			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)
+			}
+			expected := map[interface{}]interface{}{
+				"[1 2 3]": "value2",
+			}
+			if !reflect.DeepEqual(result, expected) {
+				t.Errorf("Expected %v, but got %v", expected, result)
+			}
+		})
 		t.Run("read-write time", func(t *testing.T) {
 			pos := 0
 			var buffer bytes.Buffer