aghalg: ordered map

internal/aghalg/orderedmap.go

@@ -0,0 +1,55 @@
+package aghalg
+
+import (
+	"golang.org/x/exp/slices"
+)
+
+// OrderedMap is the implementation of the ordered map data structure.
+type OrderedMap[K comparable, T any] struct {
+	vals map[K]T
+	cmp  func(a, b K) int
+	keys []K
+}
+
+// NewOrderedMap initializes the new instance of ordered map.  cmp is a sort
+// function.
+func NewOrderedMap[K comparable, T any](cmp func(a, b K) int) OrderedMap[K, T] {
+	return OrderedMap[K, T]{
+		vals: make(map[K]T),
+		cmp:  cmp,
+	}
+}
+
+// Add adds val with key to the ordered map.
+func (m *OrderedMap[K, T]) Add(key K, val T) {
+	i, has := slices.BinarySearchFunc(m.keys, key, m.cmp)
+	if has {
+		m.keys[i] = key
+		m.vals[key] = val
+
+		return
+	}
+
+	m.keys = slices.Insert(m.keys, i, key)
+	m.vals[key] = val
+}
+
+// Del removes the value by key from the ordered map.
+func (m *OrderedMap[K, T]) Del(key K) {
+	i, has := slices.BinarySearchFunc(m.keys, key, m.cmp)
+	if !has {
+		return
+	}
+
+	m.keys = slices.Delete(m.keys, i, 1)
+	delete(m.vals, key)
+}
+
+// Range calls cb for each element of the map.  If cb returns false it stops.
+func (m *OrderedMap[K, T]) Range(cb func(K, T) bool) {
+	for _, k := range m.keys {
+		if !cb(k, m.vals[k]) {
+			return
+		}
+	}
+}

internal/aghalg/orderedmap_test.go

@@ -0,0 +1,55 @@
+package aghalg
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewOrderedMap(t *testing.T) {
+	var m OrderedMap[string, int]
+
+	letters := []string{}
+	for i := 0; i < 10; i++ {
+		r := string('a' + rune(i))
+		letters = append(letters, r)
+	}
+
+	t.Run("create_and_fill", func(t *testing.T) {
+		m = NewOrderedMap[string, int](strings.Compare)
+
+		nums := []int{}
+		for i, r := range letters {
+			m.Add(r, i)
+			nums = append(nums, i)
+		}
+
+		gotLetters := []string{}
+		gotNums := []int{}
+		m.Range(func(k string, v int) bool {
+			gotLetters = append(gotLetters, k)
+			gotNums = append(gotNums, v)
+
+			return true
+		})
+
+		assert.Equal(t, letters, gotLetters)
+		assert.Equal(t, nums, gotNums)
+	})
+
+	t.Run("clear", func(t *testing.T) {
+		for _, r := range letters {
+			m.Del(r)
+		}
+
+		gotLetters := []string{}
+		m.Range(func(k string, v int) bool {
+			gotLetters = append(gotLetters, k)
+
+			return true
+		})
+
+		assert.Len(t, gotLetters, 0)
+	})
+}