Add support for JSON serialization for ordered map

Signed-off-by: Dušan Mitrović <[email protected]>

Author: dusnm

types/ordered_map.go

@@ -5,20 +5,20 @@ import (
 )
 
 type (
-	OrderedMap[T comparable, V any] struct {
-		insertionOrder []T
-		values         map[T]V
-		cmpFunc        func(a, b T) int
+	OrderedMap[K comparable, V any] struct {
+		insertionOrder []K
+		values         map[K]V
+		cmpFunc        func(a, b K) int
 	}
 )
 
 // NewOrderedMap
 //
 // Initializes a new instance of the OrderedMap.
-func NewOrderedMap[T comparable, V any](length int) OrderedMap[T, V] {
-	return OrderedMap[T, V]{
-		insertionOrder: make([]T, 0, length),
-		values:         make(map[T]V, length),
+func NewOrderedMap[K comparable, V any](length int) OrderedMap[K, V] {
+	return OrderedMap[K, V]{
+		insertionOrder: make([]K, 0, length),
+		values:         make(map[K]V, length),
 		cmpFunc:        nil,
 	}
 }
@@ -27,28 +27,28 @@ func NewOrderedMap[T comparable, V any](length int) OrderedMap[T, V] {
 //
 // Initializes a new instance of the OrderedMap with custom ordering
 // dictated by the cmpFunc comparison function.
-func NewOrderedMapWithCompareFunc[T comparable, V any](
+func NewOrderedMapWithCompareFunc[K comparable, V any](
 	length int,
-	cmpFunc func(a, b T) int,
-) OrderedMap[T, V] {
-	return OrderedMap[T, V]{
-		insertionOrder: make([]T, 0, length),
-		values:         make(map[T]V, length),
+	cmpFunc func(a, b K) int,
+) OrderedMap[K, V] {
+	return OrderedMap[K, V]{
+		insertionOrder: make([]K, 0, length),
+		values:         make(map[K]V, length),
 		cmpFunc:        cmpFunc,
 	}
 }
 
 // SetCompareFunc
 //
 // Sets the comparison function to determine ordering.
-func (om *OrderedMap[T, V]) SetCompareFunc(cmpFunc func(a, b T) int) {
+func (om *OrderedMap[K, V]) SetCompareFunc(cmpFunc func(a, b K) int) {
 	om.cmpFunc = cmpFunc
 }
 
 // UnsetCompareFunc
 //
 // Unsets the comparison function.
-func (om *OrderedMap[T, V]) UnsetCompareFunc() {
+func (om *OrderedMap[K, V]) UnsetCompareFunc() {
 	om.cmpFunc = nil
 }
 
@@ -57,14 +57,14 @@ func (om *OrderedMap[T, V]) UnsetCompareFunc() {
 // Gets the value associated with a key.
 // If the key doesn't exist in the map the zero value
 // for the type is returned and the boolean is set to false.
-func (om OrderedMap[T, V]) Get(key T) (V, bool) {
+func (om OrderedMap[K, V]) Get(key K) (V, bool) {
 	v, ok := om.values[key]
 
 	return v, ok
 }
 
 // Inserts a new value into the map while tracking the order.
-func (om *OrderedMap[T, V]) Set(key T, value V) {
+func (om *OrderedMap[K, V]) Set(key K, value V) {
 	_, exists := om.values[key]
 	if !exists {
 		om.insertionOrder = append(om.insertionOrder, key)
@@ -76,8 +76,8 @@ func (om *OrderedMap[T, V]) Set(key T, value V) {
 // Unset
 //
 // Deletes the key from the map.
-func (om *OrderedMap[T, V]) Unset(key T) {
-	newInsertionOrder := make([]T, 0, len(om.insertionOrder)-1)
+func (om *OrderedMap[K, V]) Unset(key K) {
+	newInsertionOrder := make([]K, 0, len(om.insertionOrder)-1)
 	for _, k := range om.insertionOrder {
 		if k != key {
 			newInsertionOrder = append(newInsertionOrder, k)
@@ -92,22 +92,22 @@ func (om *OrderedMap[T, V]) Unset(key T) {
 // Reset
 //
 // Resets all state including the comparison function.
-func (om *OrderedMap[T, V]) Reset() {
-	om.insertionOrder = make([]T, 0)
-	om.values = make(map[T]V, 0)
+func (om *OrderedMap[K, V]) Reset() {
+	om.insertionOrder = make([]K, 0)
+	om.values = make(map[K]V, 0)
 	om.cmpFunc = nil
 }
 
 // Iter iterates through the map.
 // If the CompareFunc is unset, it iterates in insertion order.
-// Otherewise, it iterates in the order dictated by CompareFunc.
+// Otherwise, it iterates in the order dictated by CompareFunc.
 //
 // Either use NewOrderedMapWithCompareFunc or
 // call SetCompareFunc to set the comparison function.
-func (om OrderedMap[T, V]) Iter(yield func(key T, value V) bool) {
+func (om OrderedMap[K, V]) Iter(yield func(key K, value V) bool) {
 	order := om.insertionOrder
 	if om.cmpFunc != nil {
-		order = make([]T, len(om.insertionOrder))
+		order = make([]K, len(om.insertionOrder))
 		copy(order, om.insertionOrder)
 
 		slices.SortFunc(order, om.cmpFunc)

types/ordered_map_serialize.go

@@ -0,0 +1,122 @@
+package types
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"slices"
+)
+
+var (
+	ErrInvalidJSON    = errors.New("expected a JSON object")
+	ErrInvalidJSONKey = errors.New("expected a string key in JSON")
+)
+
+// UnmarshalJSON
+//
+// Unserializes JSON objects into an ordered map
+// preserving the order in which keys appear in
+// the input datastream.
+//
+// Doesn't support unmarshaling JSON arrays as those
+// are already ordered structures for which you
+// shouldn't use this ordered map.
+func (om *OrderedMap[K, V]) UnmarshalJSON(data []byte) error {
+	decoder := json.NewDecoder(bytes.NewReader(data))
+	token, err := decoder.Token()
+	if err != nil {
+		return err
+	}
+
+	delimiter, ok := token.(json.Delim)
+	if !ok || delimiter != '{' {
+		return ErrInvalidJSON
+	}
+
+	if om.insertionOrder == nil {
+		om.insertionOrder = make([]K, 0)
+	}
+
+	if om.values == nil {
+		om.values = make(map[K]V, 0)
+	}
+
+	for decoder.More() {
+		token, err = decoder.Token()
+		if err != nil {
+			return err
+		}
+
+		key, ok := token.(string)
+		if !ok {
+			return ErrInvalidJSONKey
+		}
+
+		kJson, err := json.Marshal(key)
+		if err != nil {
+			return err
+		}
+
+		var k K
+		if err := json.Unmarshal(kJson, &k); err != nil {
+			return err
+		}
+
+		var v V
+		if err := decoder.Decode(&v); err != nil {
+			return err
+		}
+
+		om.Set(k, v)
+	}
+
+	_, err = decoder.Token()
+	return err
+}
+
+// MarshalJSON
+//
+// Serializes the ordered map into a JSON stream.
+//
+// If the CompareFunc is unset, it serializes in insertion order.
+// Otherwise, it serializes in the order dictated by CompareFunc.
+//
+// Either use NewOrderedMapWithCompareFunc or
+// call SetCompareFunc to set the comparison function.
+func (om OrderedMap[K, V]) MarshalJSON() ([]byte, error) {
+	if len(om.values) == 0 {
+		return []byte("{}"), nil
+	}
+
+	b := new(bytes.Buffer)
+	order := om.insertionOrder
+	if om.cmpFunc != nil {
+		order = make([]K, len(om.insertionOrder))
+		copy(order, om.insertionOrder)
+
+		slices.SortFunc(order, om.cmpFunc)
+	}
+
+	b.WriteRune('{')
+	for i, k := range order {
+		key, err := json.Marshal(k)
+		if err != nil {
+			return nil, err
+		}
+
+		value, err := json.Marshal(om.values[k])
+		if err != nil {
+			return nil, err
+		}
+
+		b.Write(key)
+		b.WriteRune(':')
+		b.Write(value)
+		if i < len(om.values)-1 {
+			b.WriteRune(',')
+		}
+	}
+	b.WriteRune('}')
+
+	return b.Bytes(), nil
+}

types/ordered_map_serialize_test.go

@@ -0,0 +1,71 @@
+package types
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMarshalJSONPreservesOrder(t *testing.T) {
+	t.Parallel()
+
+	om := NewOrderedMap[string, any](10)
+	om.Set("id", 123)
+	om.Set("name", "Anamarija")
+	om.Set("city", "Belgrade")
+
+	data, err := json.Marshal(om)
+	require.NoError(t, err)
+
+	require.JSONEq(t,
+		`{"id":123,"name":"Anamarija","city":"Belgrade"}`,
+		string(data),
+	)
+}
+
+func TestUnmarshalJSONPopulatesValuesAndOrder(t *testing.T) {
+	t.Parallel()
+
+	input := []byte(`{"z":1,"a":2,"x":5}`)
+	om := NewOrderedMap[string, int](10)
+	err := json.Unmarshal(input, &om)
+	require.NoError(t, err)
+
+	require.Len(t, om.values, 3)
+	require.Contains(t, om.values, "z")
+	require.Contains(t, om.values, "a")
+	require.Contains(t, om.values, "x")
+
+	require.Equal(t, []string{"z", "a", "x"}, om.insertionOrder)
+}
+
+func TestMarshalJSONWithEmptyMap(t *testing.T) {
+	t.Parallel()
+
+	om := NewOrderedMap[string, int](0)
+	data, err := json.Marshal(om)
+	require.NoError(t, err)
+	require.Equal(t, "{}", string(data))
+}
+
+func TestMarshalJSONWithCustomSort(t *testing.T) {
+	t.Parallel()
+	om := NewOrderedMapWithCompareFunc[string, int](10, func(i, j string) int {
+		if i < j {
+			return -1
+		}
+		if i > j {
+			return 1
+		}
+		return 0
+	})
+
+	om.Set("b", 2)
+	om.Set("a", 1)
+
+	data, err := json.Marshal(om)
+	require.NoError(t, err)
+
+	require.JSONEq(t, `{"a":1,"b":2}`, string(data))
+}