基于泛型实现一个有序map

package main

import (
    "fmt"
    "sort"
    "sync"
)

type OrderedMap[K comparable, V any] struct {
    keys   []K
    values map[K]V
    mutex  sync.RWMutex
    less   func(K, K) bool
}

func NewOrderedMap[K comparable, V any](less func(K, K) bool) *OrderedMap[K, V] {
    return &OrderedMap[K, V]{
        keys:   make([]K, 0),
        values: make(map[K]V),
        less:   less,
    }
}

func (om *OrderedMap[K, V]) Set(key K, value V) {
    om.mutex.Lock()
    defer om.mutex.Unlock()

    if _, exists := om.values[key]; !exists {
        om.keys = append(om.keys, key)
        sort.Slice(om.keys, func(i, j int) bool {
            return om.less(om.keys[i], om.keys[j])
        })
    }
    om.values[key] = value
}

func (om *OrderedMap[K, V]) Get(key K) (V, bool) {
    om.mutex.RLock()
    defer om.mutex.RUnlock()

    value, exists := om.values[key]
    return value, exists
}

func (om *OrderedMap[K, V]) Keys() []K {
    om.mutex.RLock()
    defer om.mutex.RUnlock()

    return om.keys
}

func (om *OrderedMap[K, V]) Range(f func(key K, value V) bool) {
    om.mutex.RLock()
    defer om.mutex.RUnlock()

    for _, key := range om.keys {
        value, exists := om.values[key]
        if !exists {
            continue
        }
        if !f(key, value) {
            break
        }
    }
}

func main() {
    om := NewOrderedMap[int, int](func(a, b int) bool {
        return a < b
    })

    var wg sync.WaitGroup
    wg.Add(2)

    go func() {
        defer wg.Done()
        om.Set(4, 2)
        om.Set(5, 1)
    }()

    go func() {
        defer wg.Done()
        om.Set(9, 3)
        om.Set(7, 4)
    }()

    wg.Wait()

    //fmt.Println("Keys in order:")
    //for _, key := range om.Keys() {
    //    value, _ := om.Get(key)
    //    fmt.Printf("%v: %v\n", key, value)
    //}

    fmt.Println("Keys in order:")
    om.Range(func(key int, value int) bool {
        fmt.Printf("%v: %v\n", key, value)
        return true
    })
}

本作品采用《CC 协议》，转载必须注明作者和本文链接