go/types/mapsets: add OrderedMap
Oh look, more Go generics.
This is a pretty standard Ordered/SortedMap type structure, with a
little twist: it implements having a subordinate 'CycleIterator' which
allows to iterate over it in a round-robin fashion, even as its
underlying OrderedMap is mutated.
This will be used, unsurprisingly, in a little TCP load balancer
implementation.
Change-Id: Ic8ab6a7fbcd6e1faf6224fc93bcc997c09364190
Reviewed-on: https://review.monogon.dev/c/monogon/+/1334
Tested-by: Jenkins CI
Reviewed-by: Lorenz Brun <lorenz@monogon.tech>
diff --git a/go.mod b/go.mod
index a69c12c..ff2d8ed 100644
--- a/go.mod
+++ b/go.mod
@@ -121,6 +121,7 @@
go.etcd.io/etcd/tests/v3 v3.5.4
go.uber.org/multierr v1.8.0
golang.org/x/crypto v0.0.0-20220517005047-85d78b3ac167
+ golang.org/x/exp v0.0.0-20220428152302-39d4317da171
golang.org/x/mod v0.6.0-dev.0.20220106191415-9b9b3d81d5e3
golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
@@ -377,7 +378,6 @@
go.uber.org/atomic v1.9.0 // indirect
go.uber.org/zap v1.19.1 // indirect
golang.org/x/arch v0.0.0-20190927153633-4e8777c89be4 // indirect
- golang.org/x/exp v0.0.0-20220428152302-39d4317da171 // indirect
golang.org/x/oauth2 v0.0.0-20220411215720-9780585627b5 // indirect
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
golang.org/x/xerrors v0.0.0-20220411194840-2f41105eb62f // indirect
diff --git a/go/types/mapsets/BUILD.bazel b/go/types/mapsets/BUILD.bazel
new file mode 100644
index 0000000..0df82e9
--- /dev/null
+++ b/go/types/mapsets/BUILD.bazel
@@ -0,0 +1,15 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+ name = "mapsets",
+ srcs = ["orderedmap.go"],
+ importpath = "source.monogon.dev/go/types/mapsets",
+ visibility = ["//visibility:public"],
+ deps = ["@org_golang_x_exp//constraints"],
+)
+
+go_test(
+ name = "mapsets_test",
+ srcs = ["orderedmap_test.go"],
+ embed = [":mapsets"],
+)
diff --git a/go/types/mapsets/orderedmap.go b/go/types/mapsets/orderedmap.go
new file mode 100644
index 0000000..bbcbedd
--- /dev/null
+++ b/go/types/mapsets/orderedmap.go
@@ -0,0 +1,231 @@
+package mapsets
+
+import (
+ "sort"
+
+ "golang.org/x/exp/constraints"
+)
+
+// OrderedMap is a map from K to V which provides total ordering as defined by
+// the keys used.
+//
+// The K type used must implement total ordering and comparability per the Key
+// interface, and must not be mutated after being inserted (or at least, must
+// keep the same ordering expressed as part of its Key implementation). A string
+// (or string wrapping type) is a good key, as strings in Go are immutable and
+// act like pointers.
+//
+// The values used will be copied by the map implementation. For structured
+// types, defining V to be a pointer to a structure is probably the correct
+// choice.
+//
+// An empty OrderedMap is ready to use and will contain no elements.
+//
+// It is not safe to be used by multiple goroutines. It can be locked behind a
+// sync.RWMutex, with a read lock taken for any method that does not mutate the
+// map, and with a write lock taken for any method that does mutate it.
+type OrderedMap[K Key, V any] struct {
+ keys []K
+ values map[K]V
+}
+
+// Key must be implemented by keys used by OrderedMap. Most 'typical' key types
+// (string, integers, etc.) already implement it.
+type Key interface {
+ comparable
+ constraints.Ordered
+}
+
+func (s *OrderedMap[K, V]) sort() {
+ sort.Slice(s.keys, func(i, j int) bool {
+ return s.keys[i] < s.keys[j]
+ })
+}
+
+// Insert a given value at a given key. If a value at this key already exists, it
+// will be overwritten with the new value.
+//
+// This method mutates the map.
+func (s *OrderedMap[K, V]) Insert(k K, v V) {
+ if s.values == nil {
+ s.values = make(map[K]V)
+ }
+
+ if _, ok := s.values[k]; !ok {
+ s.keys = append(s.keys, k)
+ s.sort()
+ }
+ s.values[k] = v
+}
+
+// Get a value at a given key. If there is no value for the given key, an empty V
+// and false will be returned.
+//
+// This returns a copy of the stored value.
+func (s *OrderedMap[K, V]) Get(k K) (V, bool) {
+ var zero V
+
+ if s.values == nil {
+ return zero, false
+ }
+
+ v, ok := s.values[k]
+ return v, ok
+}
+
+// Keys returns a copy of the keys of this map, ordered according to the Key
+// implementation used.
+func (s *OrderedMap[K, V]) Keys() []K {
+ if s.values == nil {
+ return nil
+ }
+
+ keys := make([]K, len(s.keys))
+ for i, k := range s.keys {
+ keys[i] = k
+ }
+ return keys
+}
+
+// Values returns a copy of all the keys and values of this map, ordered
+// according to the key implementation used.
+func (s *OrderedMap[K, V]) Values() []KeyValue[K, V] {
+ var res []KeyValue[K, V]
+ for _, k := range s.keys {
+ res = append(res, KeyValue[K, V]{
+ Key: k,
+ Value: s.values[k],
+ })
+ }
+ return res
+}
+
+// KeyValue represents a value V at a key K of an OrderedMap.
+type KeyValue[K Key, V any] struct {
+ Key K
+ Value V
+}
+
+// Delete the value at the given key, if present.
+//
+// This method mutates the map.
+func (s *OrderedMap[K, V]) Delete(k K) {
+ // Short-circuit delete in empty map.
+ if s.values == nil {
+ return
+ }
+
+ // Key not set? Just return.
+ if _, ok := s.values[k]; !ok {
+ return
+ }
+
+ // Iterate over keys to find index.
+ toRemove := -1
+ for i, k2 := range s.keys {
+ if k2 == k {
+ toRemove = i
+ break
+ }
+ }
+ if toRemove == -1 {
+ panic("programming error: keys and values out of sync in OrderedMap")
+ }
+
+ // No need to re-sort, as the keys were already sorted.
+ left := s.keys[0:toRemove]
+ right := s.keys[toRemove+1:]
+ s.keys = append(left, right...)
+ delete(s.values, k)
+}
+
+// Clone (perform a deep copy) of the OrderedMap, copying all keys and values.
+func (s *OrderedMap[K, V]) Clone() OrderedMap[K, V] {
+ // Short-circuit clone of empty map.
+ if s.values == nil {
+ return OrderedMap[K, V]{}
+ }
+
+ keys := make([]K, len(s.keys))
+ values := make(map[K]V)
+ for i, k := range s.keys {
+ keys[i] = k
+ values[k] = s.values[k]
+ }
+
+ return OrderedMap[K, V]{
+ keys: keys,
+ values: values,
+ }
+}
+
+// Replace all contents of this map with the contents of another map. The other
+// map must not be concurrently accessed.
+//
+// This method mutates the map.
+func (s *OrderedMap[K, V]) Replace(o *OrderedMap[K, V]) {
+ if s.values == nil {
+ s.values = make(map[K]V)
+ }
+
+ cloned := o.Clone()
+ s.keys = cloned.keys
+ s.values = cloned.values
+}
+
+// Cycle returns a CycleIterator referencing this OrderedMap.
+//
+// This iterator will cycle through all the keys/values of the OrderedMap, then
+// wrap around again when it reaches the end of the OrderedMap. This behaviour is
+// designed to be used in round-robin mechanisms.
+//
+// As the OrderedMap changes, the returned CycleIterator will also change:
+// - If a key is removed or added, the CycleIterator is guaranteed to return it
+// within one entire cycle through all the values.
+// - If a value is changed, the new value is guaranteed to be returned on the
+// next iteration when the previous value would have been returned.
+//
+// Generally, it is safe to hold on to OrderedMap pointers and mutate the
+// underlying OrderedMap.
+//
+// However, concurrent access to the OrderedMap and CycleIterator from different
+// goroutines is not safe. Access to the CycleIterator should be guarded behind
+// the same mutual exclusion mechanism as access to the OrderedMap.
+func (s *OrderedMap[K, V]) Cycle() CycleIterator[K, V] {
+ return CycleIterator[K, V]{
+ set: s,
+ }
+}
+
+// CycleIterator for an OrderedMap. See OrderedMap.Cycle for more details.
+type CycleIterator[K Key, V any] struct {
+ set *OrderedMap[K, V]
+ ix int
+}
+
+// Reset the iteration to the beginning of the ordered map.
+func (c *CycleIterator[K, V]) Reset() {
+ c.ix = 0
+}
+
+// Next returns the next key and value of the OrderedMap and a bool describing
+// the iteration state. If the OrderedMap is empty (and thus can't provide any
+// keys/values, including for this iteration), false will be returned in the
+// third value.
+func (c *CycleIterator[K, V]) Next() (K, V, bool) {
+ var zeroK K
+ var zeroV V
+
+ if len(c.set.keys) == 0 {
+ return zeroK, zeroV, false
+ }
+
+ if c.ix >= len(c.set.keys) {
+ c.ix = 0
+ }
+
+ key := c.set.keys[c.ix]
+ val := c.set.values[key]
+ c.ix++
+ return key, val, true
+}
diff --git a/go/types/mapsets/orderedmap_test.go b/go/types/mapsets/orderedmap_test.go
new file mode 100644
index 0000000..bb93dfd
--- /dev/null
+++ b/go/types/mapsets/orderedmap_test.go
@@ -0,0 +1,160 @@
+package mapsets
+
+import (
+ "strings"
+ "testing"
+)
+
+// TestOrderedMap exercises the basics of an OrderedMap.
+func TestOrderedMap(t *testing.T) {
+ m := OrderedMap[string, string]{}
+
+ if want, got := 0, len(m.Keys()); want != got {
+ t.Errorf("empty map should have %d keys, got %d", want, got)
+ }
+
+ // Test insertion.
+ m.Insert("a", "foo")
+ if got, ok := m.Get("a"); !ok || got != "foo" {
+ t.Errorf("Wanted value 'foo', got %q", got)
+ }
+ // Test update.
+ m.Insert("a", "bar")
+ if got, ok := m.Get("a"); !ok || got != "bar" {
+ t.Errorf("Wanted value 'bar', got %q", got)
+ }
+
+ // Helper to ensure that the keys returned by the map are the same as keystring
+ // split by commas.
+ keysEq := func(m *OrderedMap[string, string], keystring string) {
+ t.Helper()
+ keys := strings.Split(keystring, ",")
+ if want, got := keys, m.Keys(); len(want) == len(got) {
+ for i, k := range want {
+ if got[i] != k {
+ t.Errorf("Wanted key %d %q, got %q", i, want[i], got[i])
+ }
+ }
+ } else {
+ t.Errorf("Wanted keys %v, got %v", want, got)
+ }
+ }
+
+ // Test keys. They should get sorted.
+ m.Insert("c", "baz")
+ m.Insert("b", "barfoo")
+ keysEq(&m, "a,b,c")
+ // Test deleting and re-inserting a key.
+ m.Delete("a")
+ keysEq(&m, "b,c")
+ m.Insert("a", "a")
+ keysEq(&m, "a,b,c")
+ // Test clone.
+ n := m.Clone()
+ // Inserting into the original OrderedMap shouldn't impact the clone.
+ m.Insert("z", "zz")
+ keysEq(&m, "a,b,c,z")
+ keysEq(&n, "a,b,c")
+ // Test replace.
+ n.Replace(&m)
+ keysEq(&n, "a,b,c,z")
+ // Test Values.
+ values := n.Values()
+ if want, got := 4, len(values); want != got {
+ t.Errorf("Expected %d values, got %d", want, got)
+ } else {
+ for i, k := range []string{"a", "b", "c", "z"} {
+ if want, got := k, values[i].Key; want != got {
+ t.Errorf("Key %d should've been %q, got %q", i, want, got)
+ }
+ }
+ for i, v := range []string{"a", "barfoo", "baz", "zz"} {
+ if want, got := v, values[i].Value; want != got {
+ t.Errorf("Value %d should've been %q, got %q", i, want, got)
+ }
+ }
+ }
+}
+
+// TestCycleIterator exercises the CycleIterator.
+func TestCycleIterator(t *testing.T) {
+ m := OrderedMap[string, string]{}
+ it := m.Cycle()
+
+ // Helper to ensure that the iterator returns a list of (key, values) string
+ // pairs, as defined by keystring and valuestring, split by commas.
+ valuesEq := func(keystring, valuestring string) {
+ t.Helper()
+
+ wantKeys := strings.Split(keystring, ",")
+ if keystring == "" {
+ wantKeys = nil
+ }
+ wantValues := strings.Split(valuestring, ",")
+ if valuestring == "" {
+ wantValues = nil
+ }
+ if len(wantKeys) != len(wantValues) {
+ // Test programming error.
+ t.Fatalf("keystring length != valuestring length")
+ }
+
+ var ix int
+ for {
+ // Out of test elements? We're done.
+ if ix >= len(wantKeys) {
+ return
+ }
+
+ k, v, ok := it.Next()
+ if !ok {
+ // Iterator empty while test cases present? Fail.
+ if ix < len(wantKeys) {
+ t.Errorf("Iterator empty")
+ }
+ break
+ }
+ if want, got := wantKeys[ix], k; want != got {
+ t.Errorf("Iterator returned key %q at %d, wanted %q", got, ix, want)
+ }
+ if want, got := wantValues[ix], v; want != got {
+ t.Errorf("Iterator returned value %q at %d, wanted %q", got, ix, want)
+ }
+ ix++
+ }
+ }
+
+ // Empty iterator should fail immediately.
+ it.Reset()
+ valuesEq("", "")
+
+ // Expect keys/values in order.
+ m.Insert("z", "foo")
+ m.Insert("a", "bar")
+ m.Insert("c", "baz")
+ valuesEq("a,c,z", "bar,baz,foo")
+
+ // The iterator should cycle.
+ valuesEq("a,c,z,a,c,z,a,c,z", "bar,baz,foo,bar,baz,foo,bar,baz,foo")
+
+ // Reset iterator and read only first value.
+ it.Reset()
+ valuesEq("a", "bar")
+ // Now insert a new key, the iterator should handle that.
+ m.Insert("b", "hello")
+ valuesEq("b,c,z", "hello,baz,foo")
+
+ // Reset iterator and only read first value.
+ it.Reset()
+ valuesEq("a", "bar")
+ // Now update an existing key. The iterator should continue.
+ m.Insert("b", "goodbye")
+ valuesEq("b,c,z", "goodbye,baz,foo")
+
+ // Reset iterator and only read first value.
+ it.Reset()
+ valuesEq("a", "bar")
+ // Now remove all values. The iterator should fail immediately.
+ m.Replace(&OrderedMap[string, string]{})
+ valuesEq("", "")
+}