1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

use core::hash::{BuildHasher, Hash};

use super::{
    Bucket, Entries, Entry, Equivalent, IndexMap, IndexedEntry, IterMut2, OccupiedEntry,
    VacantEntry,
};

/// Opt-in mutable access to [`IndexMap`] keys.
///
/// These methods expose `&mut K`, mutable references to the key as it is stored
/// in the map.
/// You are allowed to modify the keys in the map **if the modification
/// does not change the key’s hash and equality**.
///
/// If keys are modified erroneously, you can no longer look them up.
/// This is sound (memory safe) but a logical error hazard (just like
/// implementing `PartialEq`, `Eq`, or `Hash` incorrectly would be).
///
/// `use` this trait to enable its methods for `IndexMap`.
///
/// This trait is sealed and cannot be implemented for types outside this crate.
pub trait MutableKeys: private::Sealed {
    type Key;
    type Value;

    /// Return item index, mutable reference to key and value
    ///
    /// Computes in **O(1)** time (average).
    fn get_full_mut2<Q>(&mut self, key: &Q) -> Option<(usize, &mut Self::Key, &mut Self::Value)>
    where
        Q: ?Sized + Hash + Equivalent<Self::Key>;

    /// Return mutable reference to key and value at an index.
    ///
    /// Valid indices are *0 <= index < self.len()*
    ///
    /// Computes in **O(1)** time.
    fn get_index_mut2(&mut self, index: usize) -> Option<(&mut Self::Key, &mut Self::Value)>;

    /// Return an iterator over the key-value pairs of the map, in their order
    fn iter_mut2(&mut self) -> IterMut2<'_, Self::Key, Self::Value>;

    /// Scan through each key-value pair in the map and keep those where the
    /// closure `keep` returns `true`.
    ///
    /// The elements are visited in order, and remaining elements keep their
    /// order.
    ///
    /// Computes in **O(n)** time (average).
    fn retain2<F>(&mut self, keep: F)
    where
        F: FnMut(&mut Self::Key, &mut Self::Value) -> bool;
}

/// Opt-in mutable access to [`IndexMap`] keys.
///
/// See [`MutableKeys`] for more information.
impl<K, V, S> MutableKeys for IndexMap<K, V, S>
where
    S: BuildHasher,
{
    type Key = K;
    type Value = V;

    fn get_full_mut2<Q>(&mut self, key: &Q) -> Option<(usize, &mut K, &mut V)>
    where
        Q: ?Sized + Hash + Equivalent<K>,
    {
        if let Some(i) = self.get_index_of(key) {
            let entry = &mut self.as_entries_mut()[i];
            Some((i, &mut entry.key, &mut entry.value))
        } else {
            None
        }
    }

    fn get_index_mut2(&mut self, index: usize) -> Option<(&mut K, &mut V)> {
        self.as_entries_mut().get_mut(index).map(Bucket::muts)
    }

    fn iter_mut2(&mut self) -> IterMut2<'_, Self::Key, Self::Value> {
        IterMut2::new(self.as_entries_mut())
    }

    fn retain2<F>(&mut self, keep: F)
    where
        F: FnMut(&mut K, &mut V) -> bool,
    {
        self.core.retain_in_order(keep);
    }
}

/// Opt-in mutable access to [`Entry`] keys.
///
/// These methods expose `&mut K`, mutable references to the key as it is stored
/// in the map.
/// You are allowed to modify the keys in the map **if the modification
/// does not change the key’s hash and equality**.
///
/// If keys are modified erroneously, you can no longer look them up.
/// This is sound (memory safe) but a logical error hazard (just like
/// implementing `PartialEq`, `Eq`, or `Hash` incorrectly would be).
///
/// `use` this trait to enable its methods for `Entry`.
///
/// This trait is sealed and cannot be implemented for types outside this crate.
pub trait MutableEntryKey: private::Sealed {
    type Key;
    fn key_mut(&mut self) -> &mut Self::Key;
}

/// Opt-in mutable access to [`Entry`] keys.
///
/// See [`MutableEntryKey`] for more information.
impl<K, V> MutableEntryKey for Entry<'_, K, V> {
    type Key = K;

    /// Gets a mutable reference to the entry's key, either within the map if occupied,
    /// or else the new key that was used to find the entry.
    fn key_mut(&mut self) -> &mut Self::Key {
        match self {
            Entry::Occupied(e) => e.key_mut(),
            Entry::Vacant(e) => e.key_mut(),
        }
    }
}

/// Opt-in mutable access to [`OccupiedEntry`] keys.
///
/// See [`MutableEntryKey`] for more information.
impl<K, V> MutableEntryKey for OccupiedEntry<'_, K, V> {
    type Key = K;
    fn key_mut(&mut self) -> &mut Self::Key {
        self.key_mut()
    }
}

/// Opt-in mutable access to [`VacantEntry`] keys.
///
/// See [`MutableEntryKey`] for more information.
impl<K, V> MutableEntryKey for VacantEntry<'_, K, V> {
    type Key = K;
    fn key_mut(&mut self) -> &mut Self::Key {
        self.key_mut()
    }
}

/// Opt-in mutable access to [`IndexedEntry`] keys.
///
/// See [`MutableEntryKey`] for more information.
impl<K, V> MutableEntryKey for IndexedEntry<'_, K, V> {
    type Key = K;
    fn key_mut(&mut self) -> &mut Self::Key {
        self.key_mut()
    }
}

mod private {
    pub trait Sealed {}

    impl<K, V, S> Sealed for super::IndexMap<K, V, S> {}
    impl<K, V> Sealed for super::Entry<'_, K, V> {}
    impl<K, V> Sealed for super::OccupiedEntry<'_, K, V> {}
    impl<K, V> Sealed for super::VacantEntry<'_, K, V> {}
    impl<K, V> Sealed for super::IndexedEntry<'_, K, V> {}
}