map.go

package orderedmap

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
)

var (
	colon = []byte{':'}
	comma = []byte{','}
)

var (
	leftCurlyBrace  = json.Delim('{')
	rightCurlyBrace = json.Delim('}')
	leftBracket     = json.Delim('[')
	rightBracket    = json.Delim(']')
)

var (
	ErrEndOfJSON    = errors.New("end of JSON")
	ErrNestedObject = errors.New("detect nested object")
	ErrNestedArray  = errors.New("detect nested array")
)

type Pair[K comparable, 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]) QuotedKey() string {
	// FIXME: how to convert comparable to string
	return fmt.Sprintf("%q", fmt.Sprintf("%v", p.Key))
}

func NewPair[K comparable, V any](k K, v V) *Pair[K, V] {
	return &Pair[K, V]{
		Key:   k,
		Value: v,
	}
}

type OrderedMap[K comparable, V any] struct {
	pairs []Pair[K, V]
	index map[K]int
	pos   int
}

func (m *OrderedMap[K, V]) Set(key K, value V) {
	pair := NewPair(key, value)
	if i, ok := m.index[key]; ok {
		m.pairs[i] = *pair
		return
	}
	m.pairs = append(m.pairs, *NewPair(key, value))
	m.index[key] = m.pos
	m.pos += 1
}

func (m *OrderedMap[K, V]) Get(key K) (V, bool) {
	i, ok := m.index[key]
	if !ok {
		return *new(V), ok
	}
	pair := m.pairs[i]
	return pair.Value, ok
}

func (m *OrderedMap[K, V]) Pairs() []Pair[K, V] {
	return m.pairs
}

func (m *OrderedMap[K, V]) String() string {
	b, err := json.Marshal(m)
	if err != nil {
		panic(fmt.Errorf("invalid data: %w", err))
	}
	return string(b)
}

func WithCapacity[K comparable, V any](size int) *OrderedMap[K, V] {
	return &OrderedMap[K, V]{
		pairs: make([]Pair[K, V], 0, size),
		index: make(map[K]int, size),
	}
}

func New[K comparable, V any]() *OrderedMap[K, V] {
	return WithCapacity[K, V](0)
}

func NewFromMap[M ~map[K]V, K comparable, V any](m M) *OrderedMap[K, V] {
	om := WithCapacity[K, V](len(m))
	for k, v := range m {
		om.Set(k, v)
	}
	return om
}

func (m OrderedMap[K, V]) MarshalJSON() ([]byte, error) {
	buf := &bytes.Buffer{}
	buf.Write([]byte(leftCurlyBrace.String()))
	pairs := m.Pairs()
	length := len(pairs) - 1
	for i, pair := range pairs {
		v, err := json.Marshal(pair.Value)
		if err != nil {
			return nil, err
		}
		buf.Write([]byte(pair.QuotedKey()))
		buf.Write(colon)
		buf.Write(v)
		if i < length {
			buf.Write(comma)
		}
	}
	buf.Write([]byte(rightCurlyBrace.String()))
	return buf.Bytes(), nil
}

func (m *OrderedMap[K, V]) initialize() {
	m.pairs = make([]Pair[K, V], 0)
	m.index = make(map[K]int)
	m.pos = 0
}

func (m *OrderedMap[K, V]) decodeKey(decoder *json.Decoder) (K, error) {
	var zero K
	token, err := decoder.Token()
	if err != nil {
		return zero, fmt.Errorf("failed to get token as key: %w", err)
	}
	if delim, ok := token.(json.Delim); ok {
		switch delim {
		case leftCurlyBrace:
			return zero, ErrNestedObject
		case leftBracket:
			return zero, ErrNestedArray
		case rightCurlyBrace, rightBracket:
			return zero, ErrEndOfJSON
		}
	}
	return assertType[K](token)
}

func (m *OrderedMap[K, V]) decodeValue(decoder *json.Decoder) (V, error) {
	var zero V
	v, err := decodeValueAsAny(decoder)
	if err != nil {
		return zero, err
	}
	if v == nil {
		return zero, nil
	}
	return assertType[V](v)
}

func (m *OrderedMap[K, V]) decodeKeyAndValue(decoder *json.Decoder) error {
	key, err := m.decodeKey(decoder)
	if err != nil {
		return fmt.Errorf("failed to decode key: %w", err)
	}
	value, err := m.decodeValue(decoder)
	if err != nil {
		return fmt.Errorf("failed to decode value: %w", err)
	}
	m.Set(key, value)
	return nil
}

func (m *OrderedMap[K, V]) validateKeyValueType() error {
	var key K
	if reflect.ValueOf(key).Kind() != reflect.String {
		return fmt.Errorf("key type must be string")
	}

	var value V
	switch reflect.ValueOf(value).Kind() {
	case reflect.Pointer:
		return fmt.Errorf("unsupported pointer, use any")
	case reflect.Map:
		return fmt.Errorf("unsupported map, use any")
	case reflect.Struct:
		return fmt.Errorf("unsupported struct, use any")
	case reflect.Array:
		return fmt.Errorf("unsupported array, use slice")
	}
	return nil
}

func (m *OrderedMap[K, V]) UnmarshalJSON(b []byte) error {
	if err := m.validateKeyValueType(); err != nil {
		return fmt.Errorf("failed to validate type: %w", err)
	}
	m.initialize()
	decoder := json.NewDecoder(bytes.NewReader(b))
	token, err := decoder.Token()
	if err != nil {
		return fmt.Errorf("failed to get token: %w", err)
	}
	if delim, ok := token.(json.Delim); !(ok && isObject(delim)) {
		return fmt.Errorf("expects JSON object, but not: %s", delim)
	}
	for {
		if err := m.decodeKeyAndValue(decoder); err != nil {
			if errors.Is(err, ErrEndOfJSON) {
				return nil
			}
			return fmt.Errorf("failed to decode key/value: %w", err)
		}
	}
}

func ToMap[M ~map[K]V, K comparable, V any](om *OrderedMap[K, V]) M {
	pairs := om.Pairs()
	m := make(M, len(pairs))
	for _, p := range pairs {
		m[p.Key] = p.Value
	}
	return m
}