We had an ordered set type (set.Slice) already but we occasionally want to do the same thing with a map, preserving the order things were added, so add that too, as mapsx.OrderedMap[K, V], and then use in ipnext. Updates #12614 Change-Id: I85e6f5e11035571a28316441075e952aef9a0863 Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>main
Brad Fitzpatrick
12 months ago
committed by
Brad Fitzpatrick
parent
3bc10ea585
commit
dbf13976d3
@ -0,0 +1,111 @@ |
||||
// Copyright (c) Tailscale Inc & AUTHORS
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
// Package mapx contains extra map types and functions.
|
||||
package mapx |
||||
|
||||
import ( |
||||
"iter" |
||||
"slices" |
||||
) |
||||
|
||||
// OrderedMap is a map that maintains the order of its keys.
|
||||
//
|
||||
// It is meant for maps that only grow or that are small;
|
||||
// is it not optimized for deleting keys.
|
||||
//
|
||||
// The zero value is ready to use.
|
||||
//
|
||||
// Locking-wise, it has the same rules as a regular Go map:
|
||||
// concurrent reads are safe, but not writes.
|
||||
type OrderedMap[K comparable, V any] struct { |
||||
// m is the underlying map.
|
||||
m map[K]V |
||||
|
||||
// keys is the order of keys in the map.
|
||||
keys []K |
||||
} |
||||
|
||||
func (m *OrderedMap[K, V]) init() { |
||||
if m.m == nil { |
||||
m.m = make(map[K]V) |
||||
} |
||||
} |
||||
|
||||
// Set sets the value for the given key in the map.
|
||||
//
|
||||
// If the key already exists, it updates the value and keeps the order.
|
||||
func (m *OrderedMap[K, V]) Set(key K, value V) { |
||||
m.init() |
||||
len0 := len(m.keys) |
||||
m.m[key] = value |
||||
if len(m.m) > len0 { |
||||
// New key (not an update)
|
||||
m.keys = append(m.keys, key) |
||||
} |
||||
} |
||||
|
||||
// Get returns the value for the given key in the map.
|
||||
// If the key does not exist, it returns the zero value for V.
|
||||
func (m *OrderedMap[K, V]) Get(key K) V { |
||||
return m.m[key] |
||||
} |
||||
|
||||
// GetOk returns the value for the given key in the map
|
||||
// and whether it was present in the map.
|
||||
func (m *OrderedMap[K, V]) GetOk(key K) (_ V, ok bool) { |
||||
v, ok := m.m[key] |
||||
return v, ok |
||||
} |
||||
|
||||
// Contains reports whether the map contains the given key.
|
||||
func (m *OrderedMap[K, V]) Contains(key K) bool { |
||||
_, ok := m.m[key] |
||||
return ok |
||||
} |
||||
|
||||
// Delete removes the key from the map.
|
||||
//
|
||||
// The cost is O(n) in the number of keys in the map.
|
||||
func (m *OrderedMap[K, V]) Delete(key K) { |
||||
len0 := len(m.m) |
||||
delete(m.m, key) |
||||
if len(m.m) == len0 { |
||||
// Wasn't present; no need to adjust keys.
|
||||
return |
||||
} |
||||
was := m.keys |
||||
m.keys = m.keys[:0] |
||||
for _, k := range was { |
||||
if k != key { |
||||
m.keys = append(m.keys, k) |
||||
} |
||||
} |
||||
} |
||||
|
||||
// All yields all the keys and values, in the order they were inserted.
|
||||
func (m *OrderedMap[K, V]) All() iter.Seq2[K, V] { |
||||
return func(yield func(K, V) bool) { |
||||
for _, k := range m.keys { |
||||
if !yield(k, m.m[k]) { |
||||
return |
||||
} |
||||
} |
||||
} |
||||
} |
||||
|
||||
// Keys yields the map keys, in the order they were inserted.
|
||||
func (m *OrderedMap[K, V]) Keys() iter.Seq[K] { |
||||
return slices.Values(m.keys) |
||||
} |
||||
|
||||
// Values yields the map values, in the order they were inserted.
|
||||
func (m *OrderedMap[K, V]) Values() iter.Seq[V] { |
||||
return func(yield func(V) bool) { |
||||
for _, k := range m.keys { |
||||
if !yield(m.m[k]) { |
||||
return |
||||
} |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,56 @@ |
||||
// Copyright (c) Tailscale Inc & AUTHORS
|
||||
// SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
package mapx |
||||
|
||||
import ( |
||||
"fmt" |
||||
"slices" |
||||
"testing" |
||||
) |
||||
|
||||
func TestOrderedMap(t *testing.T) { |
||||
// Test the OrderedMap type and its methods.
|
||||
var m OrderedMap[string, int] |
||||
m.Set("d", 4) |
||||
m.Set("a", 1) |
||||
m.Set("b", 1) |
||||
m.Set("b", 2) |
||||
m.Set("c", 3) |
||||
m.Delete("d") |
||||
m.Delete("e") |
||||
|
||||
want := map[string]int{ |
||||
"a": 1, |
||||
"b": 2, |
||||
"c": 3, |
||||
"d": 0, |
||||
} |
||||
for k, v := range want { |
||||
if m.Get(k) != v { |
||||
t.Errorf("Get(%q) = %d, want %d", k, m.Get(k), v) |
||||
continue |
||||
} |
||||
got, ok := m.GetOk(k) |
||||
if got != v { |
||||
t.Errorf("GetOk(%q) = %d, want %d", k, got, v) |
||||
} |
||||
if ok != m.Contains(k) { |
||||
t.Errorf("GetOk and Contains don't agree for %q", k) |
||||
} |
||||
} |
||||
|
||||
if got, want := slices.Collect(m.Keys()), []string{"a", "b", "c"}; !slices.Equal(got, want) { |
||||
t.Errorf("Keys() = %q, want %q", got, want) |
||||
} |
||||
if got, want := slices.Collect(m.Values()), []int{1, 2, 3}; !slices.Equal(got, want) { |
||||
t.Errorf("Values() = %v, want %v", got, want) |
||||
} |
||||
var allGot []string |
||||
for k, v := range m.All() { |
||||
allGot = append(allGot, fmt.Sprintf("%s:%d", k, v)) |
||||
} |
||||
if got, want := allGot, []string{"a:1", "b:2", "c:3"}; !slices.Equal(got, want) { |
||||
t.Errorf("All() = %q, want %q", got, want) |
||||
} |
||||
} |
||||
Loading…
Reference in new issue