1
//! Data structure `MultikeyIsolatedMap` indexed by multiple keys and an isolation
2

            
3
use std::collections::HashMap;
4
use std::hash::Hash;
5

            
6
use derive_more::{Deref, DerefMut};
7
use educe::Educe;
8
use slotmap::dense::DenseSlotMap;
9

            
10
use tor_circmgr::isolation::Isolation;
11

            
12
/// Data structure indexed by multiple keys and an isolation
13
///
14
/// This is a map keyed by:
15
///  * `K1`, a hashable key type
16
///  * `K2`, a possibly-non-hashable key type,
17
///  * [`Box<dyn Isolation>`](Isolation), a circuit/client isolation
18
///
19
/// Lookup/inserts will look for a suitable entry with a compatible isolation
20
/// and, if found, narrow that entry's isolation.
21
///
22
/// A lookup/insert yields a **table index** `I`,
23
/// which can later be used to do an O(1) lookup of the same entry.
24
/// (This is particularly useful given that an isolation cannot be simply compared;
25
/// a proper lookup might narrow the isolation of an existing record.)
26
///
27
/// `I` must implement `slotmap:Key`.
28
///
29
/// The values are `V`.
30
///
31
/// Internally, it looks like this:
32
///
33
/// ```text
34
///           index                                         table
35
///           HashMap           Vec_______________          SlotMap_____________________
36
///           |     | contains  | table_index    |  t._i.   | K2, isol, V  /  <vacant> |
37
///     K1 -> |  ---+---------> | table_index    | -------> | K2, isol, V  /  <vacant> |
38
///           |_____|           | table_index    | 1      1 | K2, isol, V  /  <vacant> |
39
///                             | table_index    |          | K2, isol, V  /  <vacant> |
40
///   K2, isol ---------------> | .............. |          | K2, isol, V  /  <vacant> |
41
///             linear search   |________________|          | ...             ....     |
42
/// ```                                                     |__________________________|
43
7
#[derive(Debug, Educe)]
44
#[educe(Default)]
45
pub(crate) struct MultikeyIsolatedMap<I, K1, K2, V>
46
where
47
    I: slotmap::Key,
48
    K1: Hash + Eq,
49
    K2: Eq,
50
{
51
    /// The first stage index, mapping `K1` to `I`
52
    index: HashMap<K1, Vec<I>>,
53

            
54
    /// Actual table containing the entries, including `K2` and the isolation, and `V`
55
    ///
56
    /// ### Invariant
57
    ///
58
    /// Entries in `table` and `index` correspond precisely, one-to-one:
59
    /// each `Vec` element in `index` refers to an (occupied) entry in `table`, and
60
    /// each (occupied) entry in `table` is referenced precisely once from `index`.
61
    table: DenseSlotMap<I, Record<K2, V>>,
62
}
63

            
64
/// Record within a `MultikeyIsolatedMap`
65
///
66
/// This contains `K2`, the isolation, and the value.
67
/// It derefs to the value, which you may mutate.
68
///
69
/// (You can't mutate the key parts; that would be wrong.)
70
#[derive(Debug, Deref, DerefMut)]
71
pub(crate) struct Record<K2, V>
72
where
73
    K2: Eq,
74
{
75
    /// K2 (part of the data structure key)
76
    k2: K2,
77
    /// Circuit isolation (part of the data structure key)
78
    isolation: Box<dyn Isolation>,
79
    /// Actual value
80
    #[deref]
81
    #[deref_mut]
82
    value: V,
83
}
84

            
85
impl<K2, V> Record<K2, V>
86
where
87
    K2: Eq,
88
{
89
    /// Obtain a reference to this record's 2nd-stage key `K2`
90
    #[allow(dead_code)] // TODO remove if and when we make this all pub
91
    pub(crate) fn k2(&self) -> &K2 {
92
        &self.k2
93
    }
94

            
95
    /// Obtain a reference to this record's isolation
96
    #[allow(dead_code)] // TODO remove if and when we make this all pub
97
    pub(crate) fn isolation(&self) -> &dyn Isolation {
98
        &*self.isolation
99
    }
100
}
101

            
102
impl<I, K1, K2, V> MultikeyIsolatedMap<I, K1, K2, V>
103
where
104
    I: slotmap::Key,
105
    K1: Hash + Eq,
106
    K2: Eq,
107
{
108
    /// Lookup, or insert, an entry
109
    ///
110
    /// Looks for an entry with keys `k1` and `k2` and a compatible isolation.
111
    /// If it finds one, narrows the isolation and returns the entry's index.
112
    ///
113
    /// If no entry is found, inserts a new value made by `create`.
114
    pub(crate) fn index_or_insert_with(
115
        &mut self,
116
        k1: &K1,
117
        k2: &K2,
118
        isolation: Box<dyn Isolation>,
119
        create: impl FnOnce() -> V,
120
    ) -> I
121
    where
122
        K1: Clone,
123
        K2: Clone,
124
    {
125
        let indices = self.index.entry(k1.clone()).or_default();
126

            
127
        match indices.iter().find_map(|&t_index| {
128
            // Deconstruct so that we can't accidentally fail to check some of the key fields
129
            let Record {
130
                k2: t_k2,
131
                isolation: t_isolation,
132
                value: _,
133
            } = self.table.get(t_index)
134
                // should be Some, unless data structure corrupted, but don't panic here
135
                    ?;
136
            (t_k2 == k2).then_some(())?;
137
            let new_isolation = t_isolation.join(&*isolation)?;
138
            Some((t_index, new_isolation))
139
        }) {
140
            Some((t_index, new_isolation)) => {
141
                self.table
142
                    .get_mut(t_index)
143
                    .expect("table entry disappeared")
144
                    .isolation = new_isolation;
145
                t_index
146
            }
147
            None => {
148
                let value = create();
149
                let record = Record {
150
                    k2: k2.clone(),
151
                    isolation,
152
                    value,
153
                };
154
                let table_index = self.table.insert(record);
155
                indices.push(table_index);
156
                table_index
157
            }
158
        }
159
    }
160

            
161
    /// Look up an existing entry by index
162
    ///
163
    /// If the entry was removed in the meantime, will return `None`
164
    #[allow(dead_code)] // TODO remove if and when we make this all pub
165
    pub(crate) fn by_index(&self, t_index: I) -> Option<&Record<K2, V>> {
166
        self.table.get(t_index)
167
    }
168

            
169
    /// Look up an existing entry by index (mutably)
170
    ///
171
    /// If the entry was removed in the meantime, will return `None`
172
    pub(crate) fn by_index_mut(&mut self, t_index: I) -> Option<&mut Record<K2, V>> {
173
        self.table.get_mut(t_index)
174
    }
175

            
176
    /// Keep only entries that match a predicate
177
    ///
178
    /// Each entry is passed to `test`, and removed unless `test` returned `true`.
179
    #[allow(dead_code)] // TODO HS remove
180
7
    pub(crate) fn retain(&mut self, mut test: impl FnMut(&K1, &Record<K2, V>, I) -> bool) {
181
7
        self.index.retain(|k1, indices| {
182
            indices.retain(|&t_index| {
183
                let record = match self.table.get(t_index) {
184
                    Some(record) => record,
185
                    None => return false, // shouldn't happen
186
                };
187
                let keep = test(k1, record, t_index);
188
                if !keep {
189
                    self.table.remove(t_index);
190
                }
191
                keep
192
            });
193
            !indices.is_empty()
194
7
        });
195
7
    }
196

            
197
    /// Checks that the structure is consistent
198
    ///
199
    /// # Panics
200
    ///
201
    /// If it is found not to be.
202
    #[cfg(test)]
203
    fn check_or_panic(&self) {
204
        let mut referenced = slotmap::SecondaryMap::default();
205

            
206
        for indices in self.index.values() {
207
            assert!(!indices.is_empty(), "empty Vec not GC'd");
208
            for (vi1, &ti1) in indices.iter().enumerate() {
209
                let rec1 = self.table.get(ti1).expect("dangling index");
210
                match referenced.entry(ti1) {
211
                    Some(slotmap::secondary::Entry::Vacant(ve)) => ve.insert(()),
212
                    _ => panic!("colliding references or something {ti1:?}"),
213
                };
214
                for &ti2 in &indices[vi1 + 1..] {
215
                    let rec2 = &self.table[ti2];
216
                    assert!(
217
                        !(rec1.k2 == rec2.k2 && rec1.isolation.compatible(&*rec2.isolation)),
218
                        "Vec contains entries that should have been merged",
219
                    );
220
                }
221
            }
222
        }
223

            
224
        for ti in self.table.keys() {
225
            let () = referenced.get(ti).expect("unreferenced entry");
226
        }
227
    }
228
}
229

            
230
// TODO HS: Currently tested via state.rs's stests, which do exercise this fairly well,
231
// but some more dedicated unit tests would be good.
232
#[cfg(test)]
233
mod test {
234
    // @@ begin test lint list maintained by maint/add_warning @@
235
    #![allow(clippy::bool_assert_comparison)]
236
    #![allow(clippy::clone_on_copy)]
237
    #![allow(clippy::dbg_macro)]
238
    #![allow(clippy::print_stderr)]
239
    #![allow(clippy::print_stdout)]
240
    #![allow(clippy::single_char_pattern)]
241
    #![allow(clippy::unwrap_used)]
242
    #![allow(clippy::unchecked_duration_subtraction)]
243
    #![allow(clippy::useless_vec)]
244
    #![allow(clippy::needless_pass_by_value)]
245
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
246
    use super::*;
247

            
248
    slotmap::new_key_type! {
249
        struct Idx;
250
    }
251

            
252
    use crate::state::test::NarrowableIsolation;
253

            
254
    fn mk_isol(s: impl Into<String>) -> Box<dyn Isolation> {
255
        NarrowableIsolation(s.into()).into()
256
    }
257

            
258
    fn mk() -> MultikeyIsolatedMap<Idx, u32, u16, String> {
259
        let mut out = MultikeyIsolatedMap::<Idx, u32, u16, String>::default();
260
        let ti = out.index_or_insert_with(&1, &22, mk_isol("a"), || "hi".into());
261
        assert_eq!(out.by_index(ti).unwrap().k2(), &22);
262
        out.check_or_panic();
263
        out
264
    }
265

            
266
    #[test]
267
    fn simple() {
268
        mk();
269
    }
270

            
271
    #[test]
272
    fn retain() {
273
        let mut m = mk();
274
        m.index_or_insert_with(&2, &22, mk_isol("ab"), || "22".into());
275
        m.check_or_panic();
276
        m.index_or_insert_with(&2, &23, mk_isol("ac"), || "23".into());
277
        m.check_or_panic();
278
        m.index_or_insert_with(&2, &24, mk_isol("dd"), || "24".into());
279
        m.check_or_panic();
280
        dbg!(&m);
281
        m.retain(|_k1, rec, _ti| (rec.k2 % 2) == 1);
282
        dbg!(&m);
283
        m.check_or_panic();
284
    }
285
}