Ordered map for JSON unmarshaling

ordmap.go

package ordmap

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"strings"
)

// Pair is a key-value pair.
type Pair[K ~string, V any] struct {
	Key   K
	Value V
}

func (p Pair[K, V]) String() string {
	return fmt.Sprintf("%v: %v", p.Key, p.Value)
}

func (p Pair[K, V]) GoString() string {
	var zk K
	var zv V
	return fmt.Sprintf(
		"ordmap.Pair[%T, %T]{Key: %#v, Value: %#v}",
		zk, zv, p.Key, p.Value,
	)
}

// OrderedMap is a map that keeps the order of the keys. It is made to be used
// with package encoding/json.
type OrderedMap[K ~string, V any] struct {
	vals map[K]ordMapValue[V]
	keys []K
}

type ordMapValue[V any] struct {
	value V
	index int
}

// NewOrderedMap creates a new OrderedMap.
func NewOrderedMap[K ~string, V any](cap int) *OrderedMap[K, V] {
	return &OrderedMap[K, V]{
		vals: make(map[K]ordMapValue[V], cap),
		keys: make([]K, 0, cap),
	}
}

// OrderedMapFromPairs creates a new OrderedMap from a slice of key-value pairs.
func OrderedMapFromPairs[K ~string, V any](pairs ...Pair[K, V]) *OrderedMap[K, V] {
	m := NewOrderedMap[K, V](len(pairs))
	for _, p := range pairs {
		m.Set(p.Key, p.Value)
	}
	return m
}

// Set sets the value for the given key.
func (m *OrderedMap[K, V]) Set(key K, value V) {
	_, ok := m.vals[key]
	if ok {
		return
	}

	m.vals[key] = ordMapValue[V]{
		value: value,
		index: len(m.keys),
	}
	m.keys = append(m.keys, key)
}

// Get returns the value for the given key, and a boolean indicating whether the
// key was found.
func (m *OrderedMap[K, V]) Get(key K) (V, bool) {
	val, ok := m.vals[key]
	return val.value, ok
}

// Delete removes the given key from the map.
func (m *OrderedMap[K, V]) Delete(key K) {
	val, ok := m.vals[key]
	if !ok {
		return
	}

	delete(m.vals, key)
	m.keys = append(m.keys[:val.index], m.keys[val.index+1:]...)
}

// Each iterates over the map in order, calling the given function for each
// key/value pair. If the function returns true, iteration stops.
func (m *OrderedMap[K, V]) Each(f func(K, V) bool) {
	for key, val := range m.vals {
		if f(key, val.value) {
			return
		}
	}
}

// Length returns the number of keys in the map.
func (m *OrderedMap[K, V]) Length() int {
	return len(m.keys)
}

// ToPairs returns a slice of key-value pairs.
func (m *OrderedMap[K, V]) ToPairs() []Pair[K, V] {
	pairs := make([]Pair[K, V], len(m.keys))
	for i, key := range m.keys {
		pairs[i] = Pair[K, V]{
			Key:   key,
			Value: m.vals[key].value,
		}
	}
	return pairs
}

func (m *OrderedMap[K, V]) MarshalJSON() ([]byte, error) {
	var b bytes.Buffer
	b.Grow(512)
	b.WriteString("{")
	m.Each(func(key K, value V) bool {
		if b.Len() != 0 {
			b.WriteString(",")
		}
		b.WriteString(fmt.Sprintf("%q:%v", key, value))
		return false
	})
	b.WriteString("}")
	return b.Bytes(), nil
}

func (m *OrderedMap[K, V]) UnmarshalJSON(b []byte) error {
	if !bytes.HasPrefix(b, []byte{'{'}) {
		return fmt.Errorf("expected '{' at start of JSON object")
	}

	dec := json.NewDecoder(bytes.NewReader(b))

	for dec.More() {
		// Read prop name
		t, err := dec.Token()
		if err != nil {
			log.Printf("Err: %v", err)
			break
		}

		name, ok := t.(string)
		if !ok {
			continue // May be a delimeter
		}

		// Read value:
		t, err = dec.Token()
		if err != nil {
			log.Printf("Err: %v", err)
			break
		}

		// Slow: remarshal to JSON, then unmarshal to our type.
		b, err := json.Marshal(t)
		if err != nil {
			return fmt.Errorf("OrderedMap: cannot marshal value: %v", err)
		}

		var val V
		if err := json.Unmarshal(b, &val); err != nil {
			return fmt.Errorf("OrderedMap: cannot unmarshal value: %v", err)
		}

		m.Set(K(name), val)
	}

	return nil
}

func (m *OrderedMap[K, V]) GoString() string {
	var s strings.Builder
	var kv K
	var vv V
	fmt.Fprintf(&s, "ordmap.OrderedMapFromPairs[%[1]T, %[2]T](", kv, vv)
	fmt.Fprintf(&s, "[]ordmap.Pair[%[1]T, %[2]T]{", kv, vv)
	var idx int
	m.Each(func(k K, v V) bool {
		if idx > 0 {
			s.WriteString(", ")
		}

		fmt.Fprintf(&s, "{Key: %#v, Value: %#v}", k, v)
		idx++
		return false
	})
	s.WriteString("})")
	return s.String()
}

ordmap_test.go

package ordmap

import (
	"encoding/json"
	"testing"
)

func TestOrderedMap(t *testing.T) {
	const j = `{
		"foo": "bar",
		"baz": "qux",
		"quux": "quuz"
	}`

	m := NewOrderedMap[string, string](3)
	if err := json.Unmarshal([]byte(j), &m); err != nil {
		t.Fatal(err)
	}

	values := m.ToPairs()
	t.Logf("%#v\n", m)
	t.Logf("%#v\n", values)

	expectedValues := []Pair[string, string]{
		{Key: "foo", Value: "bar"},
		{Key: "baz", Value: "qux"},
		{Key: "quux", Value: "quuz"},
	}

	if len(values) != len(expectedValues) {
		t.Fatalf("expected %d values, got %d", len(expectedValues), len(values))
	}

	for i, v := range values {
		if v != expectedValues[i] {
			t.Fatalf("expected %v, got %v", expectedValues[i], v)
		}
	}
}