keyed_priority_queue/
mediator.rs

1use std::borrow::Borrow;
2use std::hash::{BuildHasher, Hash};
3
4use indexmap::map::{IndexMap, OccupiedEntry as IMOccupiedEntry, VacantEntry as IMVacantEntry};
5
6use crate::editable_binary_heap::HeapIndex;
7
8/// Wrapper around possible outer vec index
9/// Used to avoid mux up with heap index
10/// And to make sure that `Mediator` indexed only with MediatorIndex
11#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug, Hash)]
12pub(crate) struct MediatorIndex(pub(crate) usize);
13
14/// This is wrapper over over indexmap that uses `MediatorIndex` as index.
15/// Also it centralized checking for panics
16#[derive(Clone, Debug)]
17pub(crate) struct Mediator<TKey: Hash + Eq, S: BuildHasher> {
18    map: IndexMap<TKey, HeapIndex, S>,
19}
20
21#[inline(always)]
22fn with_copied_heap_index<'a, T>((k, &i): (&'a T, &HeapIndex)) -> (&'a T, HeapIndex) {
23    (k, i)
24}
25
26pub(crate) struct VacantEntry<'a, TKey: 'a + Hash + Eq, S: BuildHasher> {
27    internal: IMVacantEntry<'a, TKey, HeapIndex>,
28    // look `insert` definition for this
29    map: *mut Mediator<TKey, S>,
30}
31pub(crate) struct OccupiedEntry<'a, TKey: 'a + Hash + Eq, S: BuildHasher> {
32    internal: IMOccupiedEntry<'a, TKey, HeapIndex>,
33    // look `insert` definition for this
34    map: *mut Mediator<TKey, S>,
35}
36
37pub(crate) enum MediatorEntry<'a, TKey: 'a + Hash + Eq, S: BuildHasher> {
38    Vacant(VacantEntry<'a, TKey, S>),
39    Occupied(OccupiedEntry<'a, TKey, S>),
40}
41
42impl<TKey, S> Mediator<TKey, S>
43where
44    TKey: Hash + Eq,
45    S: BuildHasher,
46{
47    pub(crate) fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self {
48        Self {
49            map: IndexMap::with_capacity_and_hasher(capacity, hasher),
50        }
51    }
52
53    #[inline(always)]
54    pub(crate) fn reserve(&mut self, additional: usize) {
55        self.map.reserve(additional)
56    }
57
58    #[inline(always)]
59    pub(crate) fn len(&self) -> usize {
60        self.map.len()
61    }
62
63    #[inline(always)]
64    pub(crate) fn is_empty(&self) -> bool {
65        self.map.is_empty()
66    }
67
68    #[inline(always)]
69    pub(crate) fn clear(&mut self) {
70        self.map.clear()
71    }
72
73    #[inline(always)]
74    pub(crate) fn get_index(&self, MediatorIndex(position): MediatorIndex) -> (&TKey, HeapIndex) {
75        self.map
76            .get_index(position)
77            .map(with_copied_heap_index)
78            .expect("All mediator indexes must be valid")
79    }
80
81    #[inline(always)]
82    pub(crate) fn entry(&mut self, key: TKey) -> MediatorEntry<TKey, S> {
83        // Pointer dereferenced only after internal entry dropped
84        // This unsafe pointer dark magic is required because you cannot handle
85        // enum that keep either Entry or Map inside:
86        // Example: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=aee5555275572e385350a786127f91ed
87        //
88        // We need to acquire mutable reference to Mediator in KeyedPriorityQueue Entry API implementation to keep consistency
89        // but keeping multiple mutable references in entry disallowed by borrow checker.
90        // We keep IndexMap entry in Mediator entry and pointer to Mediator, and allow use second one only after dropping first.
91        // This references used in 3 places:
92        // 1. keyed_priority_queue::OccupiedEntry::set_priority
93        // 2. keyed_priority_queue::OccupiedEntry::remove
94        // 3. keyed_priority_queue::VacantEntry::set_priority
95        //
96        // Also after stabilisation of `polonius` (https://rust-lang.github.io/polonius/) we would be able
97        // to remove this pointer hack from OccupiedEntry and keep reference to Mediator and index in it instead.
98        let map = self as *mut _;
99        match self.map.entry(key) {
100            indexmap::map::Entry::Occupied(internal) => {
101                MediatorEntry::Occupied(OccupiedEntry { internal, map })
102            }
103            indexmap::map::Entry::Vacant(internal) => {
104                MediatorEntry::Vacant(VacantEntry { internal, map })
105            }
106        }
107    }
108
109    #[inline(always)]
110    pub(crate) fn get<Q>(&self, key: &Q) -> Option<HeapIndex>
111    where
112        TKey: Borrow<Q>,
113        Q: Hash + Eq + ?Sized,
114    {
115        self.map.get(key).copied()
116    }
117
118    #[inline(always)]
119    pub(crate) fn get_full<'a, Q>(&'a self, key: &Q) -> Option<(MediatorIndex, &'a TKey, HeapIndex)>
120    where
121        TKey: Borrow<Q>,
122        Q: Hash + Eq + ?Sized,
123    {
124        self.map
125            .get_full(key)
126            .map(|(idx, key, &val)| (MediatorIndex(idx), key, val))
127    }
128
129    #[inline(always)]
130    pub(crate) fn swap_remove_index(
131        &mut self,
132        MediatorIndex(index): MediatorIndex,
133    ) -> (TKey, HeapIndex) {
134        self.map
135            .swap_remove_index(index)
136            .expect("All mediator indexes must be valid")
137    }
138
139    #[inline(always)]
140    pub(crate) fn get_index_mut(&mut self, MediatorIndex(index): MediatorIndex) -> &mut HeapIndex {
141        self.map
142            .get_index_mut(index)
143            .expect("All mediator indexes must be valid")
144            .1
145    }
146
147    #[cfg(test)]
148    pub(crate) fn iter(&self) -> impl Iterator<Item = (&TKey, HeapIndex)> {
149        self.map.iter().map(with_copied_heap_index)
150    }
151}
152
153impl<'a, TKey: 'a + Hash + Eq, S: BuildHasher> VacantEntry<'a, TKey, S> {
154    // Safety: make sure that nobody uses original mutable reference to mediator
155    // when returned pointer are used
156    // And the pointer never available longer than `Mediator` instance which created the VacantEntry
157    // See `Mediator::entry` and KeyedPriorityQueue's `remove` and `set_priority` entry methods.
158    #[inline]
159    pub(crate) unsafe fn insert(
160        self,
161        value: HeapIndex,
162    ) -> (&'a mut Mediator<TKey, S>, MediatorIndex) {
163        let map = self.map;
164        let result_index = MediatorIndex(self.internal.index());
165        {
166            self.internal.insert(value);
167        }
168        let mediator = map.as_mut().expect("Validated in entry method");
169        (mediator, result_index)
170    }
171
172    #[inline]
173    pub(crate) fn get_key(&self) -> &TKey {
174        self.internal.key()
175    }
176
177    #[inline]
178    pub(crate) fn index(&self) -> MediatorIndex {
179        MediatorIndex(self.internal.index())
180    }
181}
182
183impl<'a, TKey: 'a + Hash + Eq, S: BuildHasher> OccupiedEntry<'a, TKey, S> {
184    #[inline]
185    pub(crate) fn get_heap_idx(&self) -> HeapIndex {
186        *self.internal.get()
187    }
188
189    #[inline]
190    pub(crate) fn get_key(&self) -> &TKey {
191        self.internal.key()
192    }
193
194    // Safety: make sure that nobody uses original mutable reference to mediator
195    // when returned reference are used
196    // And the pointer never available longer than `Mediator` instance which created the VacantEntry
197    // See `Mediator::entry` and KeyedPriorityQueue's `set_priority` entry method.
198    #[inline]
199    pub(crate) unsafe fn transform_to_map(self) -> &'a mut Mediator<TKey, S> {
200        let map = self.map;
201        std::mem::drop(self);
202        let mediator = map.as_mut().expect("Validated in entry method");
203        mediator
204    }
205}