add marshal/unmarshal for orderedmap without changing order

go.mod

@@ -56,6 +56,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/mattn/go-runewidth v0.0.14 // indirect
+	github.com/metalim/jsonmap v0.4.1 // indirect
 	github.com/mholt/archiver/v3 v3.5.1 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect

go.sum

@@ -133,6 +133,8 @@ github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m
 github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk=
 github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU=
 github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/metalim/jsonmap v0.4.1 h1:kvj9Q5oj+xne2APsoe34LmJGTUlrY3D9WYw2zwbVuVM=
+github.com/metalim/jsonmap v0.4.1/go.mod h1:Rlps8z72TXjyqKPAE7pttAsBfhiZ99FLn0qzxvT4jDs=
 github.com/mholt/archiver/v3 v3.5.1 h1:rDjOBX9JSF5BvoJGvjqK479aL70qh9DIpZCl+k7Clwo=
 github.com/mholt/archiver/v3 v3.5.1/go.mod h1:e3dqJ7H78uzsRSEACH1joayhuSyhnonssnDhppzS1L4=
 github.com/microcosm-cc/bluemonday v1.0.21/go.mod h1:ytNkv4RrDrLJ2pqlsSI46O6IVXmZOBBD4SaJyDwwTkM=

maps/ordered_map.go

@@ -1,11 +1,25 @@
 package mapsutil
 
 import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"reflect"
+
+	"github.com/projectdiscovery/utils/conversion"
 	sliceutil "github.com/projectdiscovery/utils/slice"
+	"github.com/tidwall/gjson"
 	"golang.org/x/exp/maps"
 )
 
+var (
+	_ json.Marshaler   = &OrderedMap[string, struct{}]{}
+	_ json.Unmarshaler = &OrderedMap[string, struct{}]{}
+)
+
 // OrderedMap is a map that preserves the order of elements
+// Note: Order is only guaranteed for current level of OrderedMap
+// nested values only have order preserved if they are also OrderedMap
 type OrderedMap[k comparable, v any] struct {
 	keys []k
 	m    map[k]v
@@ -84,6 +98,85 @@ func (o *OrderedMap[k, v]) Len() int {
 	return len(o.keys)
 }
 
+// MarshalJSON marshals the OrderedMap to JSON
+func (o OrderedMap[k, v]) MarshalJSON() ([]byte, error) {
+	var buf bytes.Buffer
+	buf.WriteByte('{')
+	for i, key := range o.keys {
+		if i > 0 {
+			buf.WriteByte(',')
+		}
+		// marshal key
+		keyBin, err := json.Marshal(key)
+		if err != nil {
+			return nil, fmt.Errorf("marshal key: %w", err)
+		}
+		if len(keyBin) > 0 && keyBin[len(keyBin)-1] != '"' {
+			buf.WriteByte('"')
+			buf.Write(keyBin)
+			buf.WriteByte('"')
+		} else {
+			buf.Write(keyBin)
+		}
+		buf.WriteByte(':')
+		// marshal value
+		valueBin, err := json.Marshal(o.m[key])
+		if err != nil {
+			return nil, fmt.Errorf("marshal value: %w", err)
+		}
+		buf.Write(valueBin)
+	}
+	buf.WriteByte('}')
+	return buf.Bytes(), nil
+}
+
+type tempStruct[k comparable] struct {
+	Key k
+}
+
+// UnmarshalJSON unmarshals the OrderedMap from JSON
+func (o *OrderedMap[k, v]) UnmarshalJSON(data []byte) error {
+	// init
+	o.m = map[k]v{}
+
+	// we are only concerned about current level of ordered map
+	// nested ordered maps are not supported or need to be supported
+	// via recursive use of OrderedMap
+	err := json.Unmarshal(data, &o.m)
+	if err != nil {
+		return err
+	}
+
+	// get type of k
+	var tmpKey k
+	keyKind := reflect.TypeOf(tmpKey).Kind()
+
+	o.keys = []k{}
+	// gjson is memory efficient and faster than encoding/json
+	// so it shouldn't have any performance impact ( might consume some cpu though )
+	result := gjson.Parse(conversion.String(data))
+	result.ForEach(func(key, value gjson.Result) bool {
+		if keyKind == reflect.Interface {
+			// heterogeneous keys use any and assign
+			o.keys = append(o.keys, any(key.Value()).(k))
+			return true
+		}
+		if keyKind == reflect.String {
+			o.keys = append(o.keys, any(key.String()).(k))
+			return true
+		}
+		// if not use tmpStruct to unmarshal
+		var temp tempStruct[k]
+		err = json.Unmarshal([]byte(`{"key":`+key.String()+`}`), &temp)
+		if err != nil {
+			return false
+		}
+		o.keys = append(o.keys, temp.Key)
+		return true
+	})
+	return nil
+}
+
 // NewOrderedMap creates a new OrderedMap
 func NewOrderedMap[k comparable, v any]() OrderedMap[k, v] {
 	return OrderedMap[k, v]{

maps/ordered_map_test.go

@@ -1,7 +1,9 @@
 package mapsutil
 
 import (
+	"encoding/json"
 	"fmt"
+	"reflect"
 	"strconv"
 	"testing"
 )
@@ -70,3 +72,72 @@ func TestOrderedMap(t *testing.T) {
 	}
 
 }
+
+func TestOrderedMapMarshalUnmarshal(t *testing.T) {
+	t.Run("TestSimpleStringToStringMapping", func(t *testing.T) {
+		orderedMap1 := NewOrderedMap[string, string]()
+		orderedMap1.Set("name", "John Doe")
+		orderedMap1.Set("occupation", "Software Developer")
+
+		marshaled1, err := json.Marshal(orderedMap1)
+		if err != nil {
+			t.Fatalf("Failed to marshal orderedMap1: %v", err)
+		}
+
+		unmarshaled1 := NewOrderedMap[string, string]()
+		err = json.Unmarshal(marshaled1, &unmarshaled1)
+		if err != nil {
+			t.Fatalf("Failed to unmarshal orderedMap1: %v", err)
+		}
+
+		if !reflect.DeepEqual(orderedMap1, unmarshaled1) {
+			t.Fatal("Unmarshaled map is not equal to the original map for orderedMap1")
+		}
+	})
+
+	t.Run("TestIntegerToStructMapping", func(t *testing.T) {
+		type Employee struct {
+			ID   int    `json:"id"`
+			Name string `json:"name"`
+		}
+		orderedMap2 := NewOrderedMap[int, Employee]()
+		orderedMap2.Set(1, Employee{ID: 1, Name: "Alice"})
+		orderedMap2.Set(2, Employee{ID: 2, Name: "Bob"})
+
+		marshaled2, err := json.Marshal(orderedMap2)
+		if err != nil {
+			t.Fatalf("Failed to marshal orderedMap2: %v", err)
+		}
+
+		unmarshaled2 := NewOrderedMap[int, Employee]()
+		err = json.Unmarshal(marshaled2, &unmarshaled2)
+		if err != nil {
+			t.Fatalf("Failed to unmarshal orderedMap2: %v", err)
+		}
+
+		if !reflect.DeepEqual(orderedMap2, unmarshaled2) {
+			t.Fatal("Unmarshaled map is not equal to the original map for orderedMap2")
+		}
+	})
+
+	t.Run("TestStringToSliceOfStringsMapping", func(t *testing.T) {
+		orderedMap3 := NewOrderedMap[string, []string]()
+		orderedMap3.Set("fruits", []string{"apple", "banana", "cherry"})
+		orderedMap3.Set("vegetables", []string{"tomato", "potato", "carrot"})
+
+		marshaled3, err := json.Marshal(orderedMap3)
+		if err != nil {
+			t.Fatalf("Failed to marshal orderedMap3: %v", err)
+		}
+
+		unmarshaled3 := NewOrderedMap[string, []string]()
+		err = json.Unmarshal(marshaled3, &unmarshaled3)
+		if err != nil {
+			t.Fatalf("Failed to unmarshal orderedMap3: %v", err)
+		}
+
+		if !reflect.DeepEqual(orderedMap3, unmarshaled3) {
+			t.Fatal("Unmarshaled map is not equal to the original map for orderedMap3")
+		}
+	})
+}