part: Fix MarshalJSON and UnmarshalYAML with singletons

Neither of these properly handled the singleton cases.
Fix and add tests.

Signed-off-by: Jussi Maki <[email protected]>

Author: joamaki

part/map.go

@@ -259,12 +259,23 @@ func (m Map[K, V]) Len() int {
 }
 
 func (m Map[K, V]) MarshalJSON() ([]byte, error) {
-	if m.tree == nil {
+	if m.tree == nil && m.singleton == nil {
 		return []byte("[]"), nil
 	}
 
 	var b bytes.Buffer
 	b.WriteRune('[')
+
+	if m.singleton != nil {
+		bs, err := json.Marshal(*m.singleton)
+		if err != nil {
+			return nil, err
+		}
+		b.Write(bs)
+		b.WriteRune(']')
+		return b.Bytes(), nil
+	}
+
 	iter := m.tree.Iterator()
 	_, kv, ok := iter.Next()
 	for ok {
@@ -293,17 +304,32 @@ func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
 	if d, ok := t.(json.Delim); !ok || d != '[' {
 		return fmt.Errorf("%T.UnmarshalJSON: expected '[' got %v", m, t)
 	}
+	if !dec.More() {
+		return nil
+	}
+
+	var kv mapKVPair[K, V]
+	err = dec.Decode(&kv)
+	if err != nil {
+		return err
+	}
+
+	if !dec.More() {
+		m.singleton = &kv
+		return nil
+	}
+
 	m.ensureTree()
 	txn := m.tree.Txn()
+	txn.Insert(m.keyToBytes(kv.Key), kv)
 	for dec.More() {
 		var kv mapKVPair[K, V]
 		err := dec.Decode(&kv)
 		if err != nil {
 			return err
 		}
-		txn.Insert(m.keyToBytes(kv.Key), mapKVPair[K, V]{kv.Key, kv.Value})
+		txn.Insert(m.keyToBytes(kv.Key), kv)
 	}
-
 	t, err = dec.Token()
 	if err != nil {
 		return err
@@ -312,12 +338,6 @@ func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
 		return fmt.Errorf("%T.UnmarshalJSON: expected ']' got %v", m, t)
 	}
 	m.tree = txn.CommitOnly()
-
-	if m.tree.size == 1 {
-		_, kv, _ := m.tree.Iterator().Next()
-		m.singleton = &kv
-		m.tree = nil
-	}
 	return nil
 }
 
@@ -335,9 +355,18 @@ func (m *Map[K, V]) UnmarshalYAML(value *yaml.Node) error {
 	if value.Kind != yaml.SequenceNode {
 		return fmt.Errorf("%T.UnmarshalYAML: expected sequence", m)
 	}
-	if len(value.Content) == 0 {
+	switch len(value.Content) {
+	case 0:
+		return nil
+	case 1:
+		var kv mapKVPair[K, V]
+		if err := value.Content[0].Decode(&kv); err != nil {
+			return err
+		}
+		m.singleton = &kv
 		return nil
 	}
+
 	m.ensureTree()
 	txn := m.tree.Txn()
 	for _, e := range value.Content {

part/map_test.go

@@ -215,6 +215,18 @@ func TestSingletonMap(t *testing.T) {
 	v, found = m2.Get("two")
 	assert.True(t, found)
 	assert.Equal(t, 2, v)
+
+	var m3 part.Map[string, int]
+	bs, err := m.MarshalJSON()
+	assert.NoError(t, err)
+	assert.NoError(t, m3.UnmarshalJSON(bs))
+	assert.True(t, m.SlowEqual(m3))
+
+	m3 = part.Map[string, int]{}
+	bs, err = yaml.Marshal(m)
+	assert.NoError(t, err)
+	assert.NoError(t, yaml.Unmarshal(bs, &m3))
+	assert.True(t, m.SlowEqual(m3))
 }
 
 func TestUint64Map(t *testing.T) {

part/quick_test.go

@@ -1,6 +1,10 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Cilium
+
 package part_test
 
 import (
+	"encoding/json"
 	"fmt"
 	"maps"
 	"slices"
@@ -9,6 +13,7 @@ import (
 
 	"github.com/cilium/statedb/part"
 	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v3"
 )
 
 var quickConfig = &quick.Config{
@@ -233,6 +238,19 @@ func TestQuick_Map(t *testing.T) {
 			)
 			require.False(t, partMap.SlowEqual(newPartMap), "SlowEqual")
 		}
+
+		bs, err := json.Marshal(newPartMap)
+		require.NoError(t, err, "json.Marshal")
+		var m part.Map[uint8, int]
+		require.NoError(t, json.Unmarshal(bs, &m), "json.Unmarshal")
+		require.True(t, m.SlowEqual(newPartMap), "SlowEqual after json.Marshal")
+
+		m = part.Map[uint8, int]{}
+		bs, err = yaml.Marshal(newPartMap)
+		require.NoError(t, err)
+		require.NoError(t, yaml.Unmarshal(bs, &m), "yaml.Unmarshal")
+		require.True(t, m.SlowEqual(newPartMap), "SlowEqual after yaml.Marshal")
+
 		partMap = newPartMap
 		return
 	}