From ee9df321898b808d254b716228a6b91beb9aecdf Mon Sep 17 00:00:00 2001
From: mhmtszr <m.szr98@gmail.com>
Date: Mon, 22 Apr 2024 19:49:27 +0300
Subject: [PATCH] feat: add marshal unmarshal json support

---
 concurrent_swiss_map.go      | 23 +++++++++++++++++++++++
 concurrent_swiss_map_test.go | 21 +++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/concurrent_swiss_map.go b/concurrent_swiss_map.go
index f26a41e..1bf41fb 100644
--- a/concurrent_swiss_map.go
+++ b/concurrent_swiss_map.go
@@ -2,6 +2,7 @@ package csmap
 
 import (
 	"context"
+	"encoding/json"
 	"sync"
 
 	"github.com/mhmtszr/concurrent-swiss-map/maphash"
@@ -184,6 +185,28 @@ func (m *CsMap[K, V]) Range(f func(key K, value V) (stop bool)) {
 	listenCompleted.Wait()
 }
 
+func (m *CsMap[K, V]) MarshalJSON() ([]byte, error) {
+	tmp := make(map[K]V, m.Count())
+	m.Range(func(key K, value V) (stop bool) {
+		tmp[key] = value
+		return false
+	})
+	return json.Marshal(tmp)
+}
+
+func (m *CsMap[K, V]) UnmarshalJSON(b []byte) error {
+	tmp := make(map[K]V, m.Count())
+
+	if err := json.Unmarshal(b, &tmp); err != nil {
+		return err
+	}
+
+	for key, val := range tmp {
+		m.Store(key, val)
+	}
+	return nil
+}
+
 func (m *CsMap[K, V]) produce(ctx context.Context, ch chan Tuple[K, V]) {
 	var wg sync.WaitGroup
 	wg.Add(len(m.shards))
diff --git a/concurrent_swiss_map_test.go b/concurrent_swiss_map_test.go
index 4a9247c..6f9fb49 100644
--- a/concurrent_swiss_map_test.go
+++ b/concurrent_swiss_map_test.go
@@ -217,6 +217,27 @@ func TestDeleteFromRange(t *testing.T) {
 	}
 }
 
+func TestMarshal(t *testing.T) {
+	myMap := csmap.Create[string, int](
+		csmap.WithSize[string, int](1024),
+	)
+
+	myMap.Store("aaa", 10)
+	myMap.Store("aab", 11)
+
+	b, _ := myMap.MarshalJSON()
+
+	newMap := csmap.Create[string, int](
+		csmap.WithSize[string, int](1024),
+	)
+
+	_ = newMap.UnmarshalJSON(b)
+
+	if myMap.Count() != 2 || !myMap.Has("aaa") || !myMap.Has("aab") {
+		t.Fatal("count should be 2 after unmarshal")
+	}
+}
+
 func TestBasicConcurrentWriteDeleteCount(t *testing.T) {
 	myMap := csmap.Create[int, string](
 		csmap.WithShardCount[int, string](32),