Merge pull request #3910 from VladPavliuk/json-add-int-key-map-support

Allow to `marshal` and `unmarshal` maps with int keys

3 files changed, 103 insertions, 41 deletions

diff --git a/core/encoding/json/marshal.odin b/core/encoding/json/marshal.odin
index fa86d3468..009bf7ade 100644
--- a/core/encoding/json/marshal.odin
+++ b/core/encoding/json/marshal.odin
@@ -100,38 +100,7 @@ marshal_to_writer :: proc(w: io.Writer, v: any, opt: ^Marshal_Options) -> (err:
 
 	case runtime.Type_Info_Integer:
 		buf: [40]byte
-		u: u128
-		switch i in a {
-		case i8:      u = u128(i)
-		case i16:     u = u128(i)
-		case i32:     u = u128(i)
-		case i64:     u = u128(i)
-		case i128:    u = u128(i)
-		case int:     u = u128(i)
-		case u8:      u = u128(i)
-		case u16:     u = u128(i)
-		case u32:     u = u128(i)
-		case u64:     u = u128(i)
-		case u128:    u = u128(i)
-		case uint:    u = u128(i)
-		case uintptr: u = u128(i)
-
-		case i16le:  u = u128(i)
-		case i32le:  u = u128(i)
-		case i64le:  u = u128(i)
-		case u16le:  u = u128(i)
-		case u32le:  u = u128(i)
-		case u64le:  u = u128(i)
-		case u128le: u = u128(i)
-
-		case i16be:  u = u128(i)
-		case i32be:  u = u128(i)
-		case i64be:  u = u128(i)
-		case u16be:  u = u128(i)
-		case u32be:  u = u128(i)
-		case u64be:  u = u128(i)
-		case u128be: u = u128(i)
-		}
+		u := cast_any_int_to_u128(a)
 
 		s: string
 
@@ -310,7 +279,12 @@ marshal_to_writer :: proc(w: io.Writer, v: any, opt: ^Marshal_Options) -> (err:
 							case cstring: name = string(s)
 							}
 							opt_write_key(w, opt, name) or_return
-
+						case runtime.Type_Info_Integer:
+							buf: [40]byte
+							u := cast_any_int_to_u128(ka)
+							name = strconv.append_bits_128(buf[:], u, 10, info.signed, 8*kti.size, "0123456789", nil)
+							
+							opt_write_key(w, opt, name) or_return
 						case: return .Unsupported_Type
 						}
 					}
@@ -662,3 +636,41 @@ opt_write_indentation :: proc(w: io.Writer, opt: ^Marshal_Options) -> (err: io.E
 
 	return
 }
+
+@(private)
+cast_any_int_to_u128 :: proc(any_int_value: any) -> u128 {
+	u: u128 = 0
+	switch i in any_int_value {
+	case i8:      u = u128(i)
+	case i16:     u = u128(i)
+	case i32:     u = u128(i)
+	case i64:     u = u128(i)
+	case i128:    u = u128(i)
+	case int:     u = u128(i)
+	case u8:      u = u128(i)
+	case u16:     u = u128(i)
+	case u32:     u = u128(i)
+	case u64:     u = u128(i)
+	case u128:    u = u128(i)
+	case uint:    u = u128(i)
+	case uintptr: u = u128(i)
+
+	case i16le:  u = u128(i)
+	case i32le:  u = u128(i)
+	case i64le:  u = u128(i)
+	case u16le:  u = u128(i)
+	case u32le:  u = u128(i)
+	case u64le:  u = u128(i)
+	case u128le: u = u128(i)
+
+	case i16be:  u = u128(i)
+	case i32be:  u = u128(i)
+	case i64be:  u = u128(i)
+	case u16be:  u = u128(i)
+	case u32be:  u = u128(i)
+	case u64be:  u = u128(i)
+	case u128be: u = u128(i)
+	}
+
+	return u
+}
\ No newline at end of file
diff --git a/core/encoding/json/unmarshal.odin b/core/encoding/json/unmarshal.odin
index 344d2973e..1c1801bcd 100644
--- a/core/encoding/json/unmarshal.odin
+++ b/core/encoding/json/unmarshal.odin
@@ -475,7 +475,7 @@ unmarshal_object :: proc(p: ^Parser, v: any, end_token: Token_Kind) -> (err: Unm
 		}
 		
 	case reflect.Type_Info_Map:
-		if !reflect.is_string(t.key) {
+		if !reflect.is_string(t.key) && !reflect.is_integer(t.key) {
 			return UNSUPPORTED_TYPE
 		}
 		raw_map := (^mem.Raw_Map)(v.data)
@@ -492,25 +492,39 @@ unmarshal_object :: proc(p: ^Parser, v: any, end_token: Token_Kind) -> (err: Unm
 			key, _ := parse_object_key(p, p.allocator)
 			unmarshal_expect_token(p, .Colon)
 			
-			
+
 			mem.zero_slice(elem_backing)
 			if uerr := unmarshal_value(p, map_backing_value); uerr != nil {
 				delete(key, p.allocator)
 				return uerr
 			}
 
-			key_ptr := rawptr(&key)
+			key_ptr: rawptr
 
-			key_cstr: cstring
-			if reflect.is_cstring(t.key) {
-				key_cstr = cstring(raw_data(key))
-				key_ptr = &key_cstr
+			#partial switch tk in t.key.variant {
+				case runtime.Type_Info_String:			
+					key_ptr = rawptr(&key)
+					key_cstr: cstring
+					if reflect.is_cstring(t.key) {
+						key_cstr = cstring(raw_data(key))
+						key_ptr = &key_cstr
+					}
+				case runtime.Type_Info_Integer:
+					i, ok := strconv.parse_i128(key)
+					if !ok	{ return UNSUPPORTED_TYPE }
+					key_ptr = rawptr(&i)
+				case: return UNSUPPORTED_TYPE
 			}
-			
+
 			set_ptr := runtime.__dynamic_map_set_without_hash(raw_map, t.map_info, key_ptr, map_backing_value.data)
 			if set_ptr == nil {
 				delete(key, p.allocator)
 			} 
+
+			// there's no need to keep string value on the heap, since it was copied into map 
+			if reflect.is_integer(t.key) {
+				delete(key, p.allocator)
+			}
 			
 			if parse_comma(p) {
 				break map_loop
diff --git a/tests/core/encoding/json/test_core_json.odin b/tests/core/encoding/json/test_core_json.odin
index a50dd7fe0..10e09df3b 100644
--- a/tests/core/encoding/json/test_core_json.odin
+++ b/tests/core/encoding/json/test_core_json.odin
@@ -3,6 +3,7 @@ package test_core_json
 import "core:encoding/json"
 import "core:testing"
 import "core:mem/virtual"
+import "base:runtime"
 
 @test
 parse_json :: proc(t: ^testing.T) {
@@ -389,4 +390,39 @@ struct_with_ignore_tags :: proc(t: ^testing.T) {
 	expected_json := `{}`
 
 	testing.expectf(t, expected_json == my_struct_json, "Expected `json.marshal` to return %s, got %s", expected_json, my_struct_json)
+}
+
+@test
+map_with_integer_keys :: proc(t: ^testing.T) {
+	my_map := make(map[i32]string)
+	defer delete_map(my_map)
+
+	my_map[-1] = "a"
+	my_map[0] = "b"
+	my_map[42] = "c"
+	my_map[99999999] = "d"
+
+	marshaled_data, marshal_err := json.marshal(my_map)
+	defer delete(marshaled_data)
+	
+	testing.expectf(t, marshal_err == nil, "Expected `json.marshal` to return nil error, got %v", marshal_err)
+
+	my_map2 := make(map[i32]string)
+	defer delete_map(my_map2)
+
+	unmarshal_err := json.unmarshal(marshaled_data, &my_map2)
+	defer for key, item in my_map2 {
+		runtime.delete_string(item)
+	}
+	testing.expectf(t, unmarshal_err == nil, "Expected `json.unmarshal` to return nil, got %v", unmarshal_err)
+
+	testing.expectf(t, len(my_map) == len(my_map2), "Expected %v map items to have been unmarshaled, got %v", len(my_map), len(my_map2))
+
+	for key, item in my_map {
+		testing.expectf(t, key in my_map2, "Expected key %v to be present in unmarshaled map", key)
+		
+		if key in my_map2 {
+			testing.expectf(t, runtime.string_eq(item, my_map2[key]), "Expected value %s to be present in unmarshaled map", key)
+		}
+	}
 }
\ No newline at end of file