diff --git a/libs/ds/ordered_map.go b/libs/ds/ordered_map.go index e7be78ebf..fa1297c6c 100644 --- a/libs/ds/ordered_map.go +++ b/libs/ds/ordered_map.go @@ -68,12 +68,15 @@ func (m *OrderedMap[T, V]) Delete(key T) { if !ok { return } - i++ - for ; i < len(m.values); i++ { - m.values[i-1] = m.values[i] - } delete(m.keys, key) + + m.values = append(m.values[:i], m.values[i+1:]...) m.len-- + for k, v := range m.keys { + if v > i { + m.keys[k] = v - 1 + } + } } // Values returns all values in the map diff --git a/libs/ds/ordered_map_test.go b/libs/ds/ordered_map_test.go index 3ed15647a..a1c011e8f 100644 --- a/libs/ds/ordered_map_test.go +++ b/libs/ds/ordered_map_test.go @@ -5,6 +5,7 @@ import ( "sync" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -59,3 +60,25 @@ func TestOrderedMapMultithread(t *testing.T) { wg.Wait() } + +func TestOrderedMapDelete(t *testing.T) { + m := NewOrderedMap[int, int]() + m.Put(1, 1) + m.Put(2, 2) + m.Put(3, 3) + m.Delete(2) + keys := m.Keys() + if len(keys) != 2 { + t.Errorf("Expected 2 keys, got %d", len(keys)) + } + if keys[0] != 1 && keys[1] != 3 { + t.Errorf("Expected keys [1, 3], got %v", keys) + } + v1, ok := m.Get(1) + assert.Equal(t, v1, 1) + assert.True(t, ok) + + v3, ok := m.Get(3) + assert.Equal(t, v3, 3) + assert.True(t, ok) +}