feat: add OrderedMap for issue #237

maputil/orderedmap.go

@@ -0,0 +1,431 @@
+package maputil
+
+import (
+	"container/list"
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"sort"
+	"strconv"
+	"sync"
+)
+
+// OrderedMap is a map that maintains the order of keys.
+type OrderedMap[K comparable, V any] struct {
+	mu sync.RWMutex
+
+	data  map[K]V
+	order *list.List
+	index map[K]*list.Element
+}
+
+// NewOrderedMap creates a new OrderedMap.
+func NewOrderedMap[K comparable, V any]() *OrderedMap[K, V] {
+	return &OrderedMap[K, V]{
+		data:  make(map[K]V),
+		order: list.New(),
+		index: make(map[K]*list.Element),
+	}
+}
+
+// Get returns the value for the given key.
+// Play: todo
+func (om *OrderedMap[K, V]) Set(key K, value V) {
+	om.mu.Lock()
+	defer om.mu.Unlock()
+
+	if elem, ok := om.index[key]; ok {
+		elem.Value = value
+		om.order.MoveToBack(elem)
+
+		return
+	}
+
+	om.data[key] = value
+
+	elem := om.order.PushBack(key)
+	om.index[key] = elem
+}
+
+// Get returns the value for the given key.
+// Play: todo
+func (om *OrderedMap[K, V]) Get(key K) (V, bool) {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	value, ok := om.data[key]
+
+	return value, ok
+}
+
+// Delete deletes the given key.
+// Play: todo
+func (om *OrderedMap[K, V]) Delete(key K) {
+	om.mu.Lock()
+	defer om.mu.Unlock()
+
+	if elem, ok := om.index[key]; ok {
+		om.order.Remove(elem)
+		delete(om.data, key)
+		delete(om.index, key)
+	}
+}
+
+// Clear clears the map.
+// Play: todo
+func (om *OrderedMap[K, V]) Clear() {
+	om.mu.Lock()
+	defer om.mu.Unlock()
+
+	om.data = make(map[K]V)
+	om.order.Init()
+	om.index = make(map[K]*list.Element)
+}
+
+// Front returns the first key-value pair.
+// Play: todo
+func (om *OrderedMap[K, V]) Front() (struct {
+	Key   K
+	Value V
+}, bool) {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	if elem := om.order.Front(); elem != nil {
+		key := elem.Value.(K)
+		value := om.data[key]
+
+		return struct {
+			Key   K
+			Value V
+		}{
+			Key:   key,
+			Value: value,
+		}, true
+	}
+
+	return struct {
+		Key   K
+		Value V
+	}{}, false
+}
+
+// Back returns the last key-value pair.
+// Play: todo
+func (om *OrderedMap[K, V]) Back() (struct {
+	Key   K
+	Value V
+}, bool) {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	if elem := om.order.Back(); elem != nil {
+		key := elem.Value.(K)
+		value := om.data[key]
+
+		return struct {
+			Key   K
+			Value V
+		}{
+			Key:   key,
+			Value: value,
+		}, true
+	}
+
+	return struct {
+		Key   K
+		Value V
+	}{}, false
+}
+
+// Range calls the given function for each key-value pair.
+// Play: todo
+func (om *OrderedMap[K, V]) Range(iteratee func(key K, value V) bool) {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	for elem := om.order.Front(); elem != nil; elem = elem.Next() {
+		key := elem.Value.(K)
+		value := om.data[key]
+
+		if !iteratee(key, value) {
+			break
+		}
+	}
+}
+
+// Keys returns the keys in order.
+// Play: todo
+func (om *OrderedMap[K, V]) Keys() []K {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	keys := make([]K, 0, len(om.data))
+
+	for elem := om.order.Front(); elem != nil; elem = elem.Next() {
+		keys = append(keys, elem.Value.(K))
+	}
+
+	return keys
+}
+
+// Values returns the values in order.
+// Play: todo
+func (om *OrderedMap[K, V]) Values() []V {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	values := make([]V, 0, len(om.data))
+
+	for elem := om.order.Front(); elem != nil; elem = elem.Next() {
+		key := elem.Value.(K)
+		values = append(values, om.data[key])
+	}
+
+	return values
+}
+
+// Len returns the number of key-value pairs.
+// Play: todo
+func (om *OrderedMap[K, V]) Len() int {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	return len(om.data)
+}
+
+// Contains returns true if the given key exists.
+// Play: todo
+func (om *OrderedMap[K, V]) Contains(key K) bool {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	_, ok := om.data[key]
+
+	return ok
+}
+
+// Elements returns the key-value pairs in order.
+// Play: todo
+func (om *OrderedMap[K, V]) Elements() []struct {
+	Key   K
+	Value V
+} {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	elements := make([]struct {
+		Key   K
+		Value V
+	}, 0, len(om.data))
+
+	for elem := om.order.Front(); elem != nil; elem = elem.Next() {
+		key := elem.Value.(K)
+		value := om.data[key]
+		elements = append(elements, struct {
+			Key   K
+			Value V
+		}{Key: key, Value: value})
+	}
+
+	return elements
+}
+
+// Iter returns a channel that yields key-value pairs in order.
+// Play: todo
+func (om *OrderedMap[K, V]) Iter() <-chan struct {
+	Key   K
+	Value V
+} {
+	ch := make(chan struct {
+		Key   K
+		Value V
+	})
+
+	go func() {
+		om.mu.RLock()
+		defer om.mu.RUnlock()
+		defer close(ch)
+
+		for elem := om.order.Front(); elem != nil; elem = elem.Next() {
+			key := elem.Value.(K)
+			value := om.data[key]
+
+			ch <- struct {
+				Key   K
+				Value V
+			}{Key: key, Value: value}
+		}
+	}()
+
+	return ch
+}
+
+// ReverseIter returns a channel that yields key-value pairs in reverse order.
+// Play: todo
+func (om *OrderedMap[K, V]) ReverseIter() <-chan struct {
+	Key   K
+	Value V
+} {
+	ch := make(chan struct {
+		Key   K
+		Value V
+	})
+
+	go func() {
+		om.mu.RLock()
+		defer om.mu.RUnlock()
+		defer close(ch)
+
+		for elem := om.order.Back(); elem != nil; elem = elem.Prev() {
+			key := elem.Value.(K)
+			value := om.data[key]
+
+			ch <- struct {
+				Key   K
+				Value V
+			}{Key: key, Value: value}
+		}
+	}()
+
+	return ch
+}
+
+// SortByValue sorts the map by key given less function.
+// Play: todo
+func (om *OrderedMap[K, V]) SortByKey(less func(a, b K) bool) {
+	om.mu.Lock()
+	defer om.mu.Unlock()
+
+	keys := make([]K, 0, om.order.Len())
+	for elem := om.order.Front(); elem != nil; elem = elem.Next() {
+		keys = append(keys, elem.Value.(K))
+	}
+
+	sort.Slice(keys, func(i, j int) bool {
+		return less(keys[i], keys[j])
+	})
+
+	om.order.Init()
+	om.index = make(map[K]*list.Element)
+	for _, key := range keys {
+		elem := om.order.PushBack(key)
+		om.index[key] = elem
+	}
+}
+
+// MarshalJSON implements the json.Marshaler interface.
+// Play: todo
+func (om *OrderedMap[K, V]) MarshalJSON() ([]byte, error) {
+	om.mu.RLock()
+	defer om.mu.RUnlock()
+
+	tempMap := make(map[string]V)
+	for e := om.order.Front(); e != nil; e = e.Next() {
+		key := e.Value.(K)
+		keyStr, err := keyToString(key)
+		if err != nil {
+			return nil, err
+		}
+		tempMap[keyStr] = om.data[key]
+	}
+
+	return json.Marshal(tempMap)
+}
+
+// UnmarshalJSON implements the json.Unmarshaler interface.
+// Play: todo
+func (om *OrderedMap[K, V]) UnmarshalJSON(data []byte) error {
+	om.mu.Lock()
+	defer om.mu.Unlock()
+
+	tempMap := make(map[string]V)
+	if err := json.Unmarshal(data, &tempMap); err != nil {
+		return err
+	}
+
+	om.data = make(map[K]V)
+	om.order.Init()
+	om.index = make(map[K]*list.Element)
+
+	for keyStr, value := range tempMap {
+		key, err := stringToKey[K](keyStr)
+		if err != nil {
+			return err
+		}
+		om.data[key] = value
+		elem := om.order.PushBack(key)
+		om.index[key] = elem
+	}
+
+	return nil
+}
+
+func keyToString[K any](key K) (string, error) {
+	switch v := any(key).(type) {
+	case int:
+		return strconv.Itoa(v), nil
+	case float64:
+		return strconv.FormatFloat(v, 'f', -1, 64), nil
+	case string:
+		return v, nil
+	default:
+		// 使用反射将未知类型转换为字符串
+		rv := reflect.ValueOf(key)
+		switch rv.Kind() {
+		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+			return strconv.FormatInt(rv.Int(), 10), nil
+		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+			return strconv.FormatUint(rv.Uint(), 10), nil
+		case reflect.Float32, reflect.Float64:
+			return strconv.FormatFloat(rv.Float(), 'f', -1, 64), nil
+		case reflect.String:
+			return rv.String(), nil
+		default:
+			return "", fmt.Errorf("unsupported key type: %T", key)
+		}
+	}
+}
+
+func stringToKey[K any](s string) (K, error) {
+	var zero K
+	switch any(zero).(type) {
+	case int:
+		value, err := strconv.Atoi(s)
+		return any(value).(K), err
+	case float64:
+		value, err := strconv.ParseFloat(s, 64)
+		return any(value).(K), err
+	case string:
+		return any(s).(K), nil
+	default:
+		// 使用反射恢复未知类型的键
+		rv := reflect.ValueOf(&zero).Elem()
+		switch rv.Kind() {
+		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+			val, err := strconv.ParseInt(s, 10, 64)
+			if err != nil {
+				return zero, err
+			}
+			rv.SetInt(val)
+		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+			val, err := strconv.ParseUint(s, 10, 64)
+			if err != nil {
+				return zero, err
+			}
+			rv.SetUint(val)
+		case reflect.Float32, reflect.Float64:
+			val, err := strconv.ParseFloat(s, 64)
+			if err != nil {
+				return zero, err
+			}
+			rv.SetFloat(val)
+		case reflect.String:
+			rv.SetString(s)
+		default:
+			return zero, fmt.Errorf("unsupported key type: %T", zero)
+		}
+
+		return rv.Interface().(K), nil
+	}
+}

maputil/orderedmap_test.go

@@ -0,0 +1,245 @@
+package maputil
+
+import (
+	"testing"
+
+	"github.com/duke-git/lancet/v2/internal"
+)
+
+func TestOrderedMap_Set_Get(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Set_Get")
+
+	val, ok := om.Get("a")
+	assert.Equal(1, val)
+	assert.Equal(true, ok)
+
+	val, ok = om.Get("d")
+	assert.Equal(false, ok)
+	assert.Equal(0, val)
+}
+
+func TestOrderedMap_Front_Back(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Front_Back")
+
+	frontElement, ok := om.Front()
+	assert.Equal("a", frontElement.Key)
+	assert.Equal(1, frontElement.Value)
+	assert.Equal(true, ok)
+
+	backElement, ok := om.Back()
+	assert.Equal("c", backElement.Key)
+	assert.Equal(3, backElement.Value)
+	assert.Equal(true, ok)
+}
+
+func TestOrderedMap_Delete_Clear(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Delete_Clear")
+
+	assert.Equal(3, om.Len())
+
+	om.Delete("b")
+	assert.Equal(2, om.Len())
+
+	om.Clear()
+	assert.Equal(0, om.Len())
+}
+
+func TestOrderedMap_Keys_Values(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Keys_Values")
+
+	assert.Equal([]string{"a", "b", "c"}, om.Keys())
+	assert.Equal([]int{1, 2, 3}, om.Values())
+}
+
+func TestOrderedMap_Contains(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Contains")
+
+	assert.Equal(true, om.Contains("a"))
+	assert.Equal(false, om.Contains("d"))
+}
+
+func TestOrderedMap_Eelements(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Eelements")
+
+	elements := []struct {
+		Key   string
+		Value int
+	}{
+		{"a", 1},
+		{"b", 2},
+		{"c", 3},
+	}
+
+	assert.Equal(elements, om.Elements())
+}
+
+func TestOrderedMap_Range(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+	om.Set("d", 4)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Range")
+
+	var keys []string
+	om.Range(func(key string, value int) bool {
+		keys = append(keys, key)
+		return key != "c"
+	})
+
+	assert.Equal([]string{"a", "b", "c"}, keys)
+}
+
+func TestOrderedMap_Iter(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_Iter")
+
+	var items []struct {
+		Key   string
+		Value int
+	}
+
+	iterCh := om.Iter()
+
+	for item := range iterCh {
+		items = append(items, item)
+	}
+
+	expected := []struct {
+		Key   string
+		Value int
+	}{
+		{"a", 1},
+		{"b", 2},
+		{"c", 3},
+	}
+
+	assert.Equal(expected, items)
+}
+
+func TestOrderedMap_ReverseIter(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_ReverseIter")
+
+	var items []struct {
+		Key   string
+		Value int
+	}
+
+	iterCh := om.ReverseIter()
+
+	for item := range iterCh {
+		items = append(items, item)
+	}
+
+	expected := []struct {
+		Key   string
+		Value int
+	}{
+		{"c", 3},
+		{"b", 2},
+		{"a", 1},
+	}
+
+	assert.Equal(expected, items)
+}
+
+func TestOrderedMap_SortByKey(t *testing.T) {
+	assert := internal.NewAssert(t, "TestOrderedMap_SortByKey")
+
+	om := NewOrderedMap[string, int]()
+
+	om.Set("d", 4)
+	om.Set("b", 2)
+	om.Set("c", 3)
+	om.Set("a", 1)
+
+	om.SortByKey(func(a, b string) bool {
+		return a < b
+	})
+
+	assert.Equal([]string{"a", "b", "c", "d"}, om.Keys())
+}
+
+func TestOrderedMap_MarshalJSON(t *testing.T) {
+	om := NewOrderedMap[string, int]()
+
+	om.Set("a", 1)
+	om.Set("b", 2)
+	om.Set("c", 3)
+
+	assert := internal.NewAssert(t, "TestOrderedMap_MarshalJSON")
+
+	jsonBytes, err := om.MarshalJSON()
+
+	if err != nil {
+		t.Errorf("MarshalJSON error: %v", err)
+	}
+
+	assert.Equal(`{"a":1,"b":2,"c":3}`, string(jsonBytes))
+}
+
+func TestOrderedMap_UnmarshalJSON(t *testing.T) {
+	assert := internal.NewAssert(t, "TestOrderedMap_UnmarshalJSON")
+
+	jsonStr := `{"a":1,"b":2,"c":3}`
+
+	om := NewOrderedMap[string, int]()
+	err := om.UnmarshalJSON([]byte(jsonStr))
+	if err != nil {
+		t.Errorf("MarshalJSON error: %v", err)
+	}
+
+	assert.Equal(3, om.Len())
+	assert.Equal(true, om.Contains("a"))
+	assert.Equal(true, om.Contains("b"))
+	assert.Equal(true, om.Contains("c"))
+}