charon_lib/ids/
index_map.rs

1//! A vector with custom index types.
2//!
3//! This data-structure is mostly meant to be used with the index types defined
4//! with [`crate::generate_index_type!`]: by using custom index types, we
5//! leverage the type checker to prevent us from mixing them.
6
7use index_vec::{Idx, IdxSliceIndex, IndexVec};
8use serde::{Deserialize, Serialize, Serializer};
9use serde_state::{DeserializeState, SerializeState};
10use std::{
11    iter::{FromIterator, IntoIterator},
12    ops::{ControlFlow, Index, IndexMut},
13};
14
15use derive_generic_visitor::*;
16
17/// Non-contiguous indexed vector.
18/// To prevent accidental id reuse, the vector supports reserving a slot to be filled later. Use
19/// `IndexVec` if this is not needed.
20#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
21#[charon::rename("IndexedMap")]
22pub struct IndexMap<I, T>
23where
24    I: Idx,
25{
26    vector: IndexVec<I, Option<T>>,
27    /// The number of non-`None` elements.
28    elem_count: usize,
29}
30
31impl<I, T> IndexMap<I, T>
32where
33    I: Idx,
34{
35    pub fn new() -> Self {
36        IndexMap {
37            vector: IndexVec::new(),
38            elem_count: 0,
39        }
40    }
41
42    pub fn with_capacity(capacity: usize) -> Self {
43        IndexMap {
44            vector: IndexVec::with_capacity(capacity),
45            elem_count: 0,
46        }
47    }
48
49    pub fn get(&self, i: I) -> Option<&T> {
50        self.vector.get(i).and_then(Option::as_ref)
51    }
52
53    pub fn get_mut(&mut self, i: I) -> Option<&mut T> {
54        self.vector.get_mut(i).and_then(Option::as_mut)
55    }
56
57    pub fn is_empty(&self) -> bool {
58        self.elem_count == 0
59    }
60
61    /// The number of elements stored in the vector.
62    pub fn elem_count(&self) -> usize {
63        self.elem_count
64    }
65
66    /// The number of slots allocated in the vector (empty or not).
67    pub fn slot_count(&self) -> usize {
68        self.vector.len()
69    }
70
71    /// Reserve a spot in the vector.
72    pub fn reserve_slot(&mut self) -> I {
73        // Push a `None` to ensure we don't reuse the id.
74        self.vector.push(None)
75    }
76    /// Ensure there's a slot for this id.
77    fn ensure_slot_for(&mut self, id: I) {
78        if id.index() >= self.vector.len() {
79            self.vector.resize_with(id.index() + 1, || None);
80        }
81    }
82
83    /// Fill the reserved slot.
84    pub fn set_slot(&mut self, id: I, x: T) {
85        assert!(self.vector[id].is_none());
86        self.vector[id] = Some(x);
87        self.elem_count += 1;
88    }
89    /// Fill the given slot even if it hadn't been reserved before. Panics if the slot already has
90    /// a value.
91    pub fn set_slot_extend(&mut self, id: I, x: T) {
92        self.ensure_slot_for(id);
93        self.set_slot(id, x);
94    }
95    /// Fill the given slot even if it hadn't been reserved before. Returns the old value.
96    pub fn insert(&mut self, id: I, x: T) -> Option<T> {
97        self.ensure_slot_for(id);
98        let old = self.vector[id].replace(x);
99        if old.is_none() {
100            self.elem_count += 1;
101        }
102        old
103    }
104
105    /// Remove the value from this slot, leaving other ids unchanged.
106    pub fn remove(&mut self, id: I) -> Option<T> {
107        if self.vector[id].is_some() {
108            self.elem_count -= 1;
109        }
110        self.vector[id].take()
111    }
112
113    /// Remove the value from this slot, shifting other ids as needed.
114    pub fn remove_and_shift_ids(&mut self, id: I) -> Option<T> {
115        if id.index() >= self.slot_count() {
116            return None;
117        }
118        if self.vector[id].is_some() {
119            self.elem_count -= 1;
120        }
121        self.vector.remove(id)
122    }
123
124    /// Remove the last slot.
125    pub fn pop(&mut self) -> Option<T> {
126        if self.vector.last().is_some() {
127            self.elem_count -= 1;
128        }
129        self.vector.pop().flatten()
130    }
131
132    pub fn push(&mut self, x: T) -> I {
133        self.elem_count += 1;
134        self.vector.push(Some(x))
135    }
136
137    pub fn push_with(&mut self, f: impl FnOnce(I) -> T) -> I {
138        let id = self.reserve_slot();
139        let x = f(id);
140        self.set_slot(id, x);
141        id
142    }
143
144    pub fn extend_from_other(&mut self, other: Self) {
145        self.vector.extend(other.vector);
146        self.elem_count += other.elem_count;
147    }
148    pub fn clone_extend_from_other(&mut self, other: &Self)
149    where
150        T: Clone,
151    {
152        self.vector.extend_from_slice(&other.vector);
153        self.elem_count += other.elem_count;
154    }
155
156    /// Insert a value at that index, shifting all the values with equal or larger indices.
157    pub fn insert_and_shift_ids(&mut self, id: I, x: T) {
158        self.elem_count += 1;
159        self.vector.insert(id, Some(x))
160    }
161
162    /// Get a mutable reference into the ith element. If the vector is too short, extend it until
163    /// it has enough elements. If the element doesn't exist, use the provided function to
164    /// initialize it.
165    pub fn get_or_extend_and_insert(&mut self, id: I, f: impl FnOnce() -> T) -> &mut T {
166        self.ensure_slot_for(id);
167        self.vector[id].get_or_insert_with(f)
168    }
169
170    /// Get a mutable reference into the ith element. If the vector is too short, extend it until
171    /// it has enough elements. If the element doesn't exist, use the provided function to
172    /// initialize it.
173    pub fn get_or_extend_and_insert_default(&mut self, id: I) -> &mut T
174    where
175        T: Default,
176    {
177        self.get_or_extend_and_insert(id, Default::default)
178    }
179
180    /// Map each entry to a new one, keeping the same ids.
181    pub fn map<U>(self, mut f: impl FnMut(T) -> U) -> IndexMap<I, U> {
182        IndexMap {
183            vector: self
184                .vector
185                .into_iter()
186                .map(|x_opt| x_opt.map(&mut f))
187                .collect(),
188            elem_count: self.elem_count,
189        }
190    }
191
192    /// Map each entry to a new one, keeping the same ids.
193    pub fn map_ref<'a, U>(&'a self, mut f: impl FnMut(&'a T) -> U) -> IndexMap<I, U> {
194        IndexMap {
195            vector: self
196                .vector
197                .iter()
198                .map(|x_opt| x_opt.as_ref().map(&mut f))
199                .collect(),
200            elem_count: self.elem_count,
201        }
202    }
203
204    /// Map each entry to a new one, keeping the same ids.
205    pub fn map_ref_mut<'a, U>(&'a mut self, mut f: impl FnMut(&'a mut T) -> U) -> IndexMap<I, U> {
206        IndexMap {
207            vector: self
208                .vector
209                .iter_mut()
210                .map(|x_opt| x_opt.as_mut().map(&mut f))
211                .collect(),
212            elem_count: self.elem_count,
213        }
214    }
215
216    /// Map each entry to a new one, keeping the same ids.
217    pub fn map_indexed<U>(self, mut f: impl FnMut(I, T) -> U) -> IndexMap<I, U> {
218        IndexMap {
219            vector: self
220                .vector
221                .into_iter_enumerated()
222                .map(|(i, x_opt)| x_opt.map(|x| f(i, x)))
223                .collect(),
224            elem_count: self.elem_count,
225        }
226    }
227
228    /// Map each entry to a new one, keeping the same ids.
229    pub fn map_ref_indexed<'a, U>(&'a self, mut f: impl FnMut(I, &'a T) -> U) -> IndexMap<I, U> {
230        IndexMap {
231            vector: self
232                .vector
233                .iter_enumerated()
234                .map(|(i, x_opt)| x_opt.as_ref().map(|x| f(i, x)))
235                .collect(),
236            elem_count: self.elem_count,
237        }
238    }
239
240    /// Map each entry to a new one, keeping the same ids. Includes empty slots.
241    pub fn map_opt<U>(self, f: impl FnMut(Option<T>) -> Option<U>) -> IndexMap<I, U> {
242        IndexMap {
243            vector: self.vector.into_iter().map(f).collect(),
244            elem_count: self.elem_count,
245        }
246    }
247
248    /// Map each entry to a new one, keeping the same ids. Includes empty slots.
249    pub fn map_ref_opt<'a, U>(
250        &'a self,
251        mut f: impl FnMut(Option<&'a T>) -> Option<U>,
252    ) -> IndexMap<I, U> {
253        let mut ret = IndexMap {
254            vector: self.vector.iter().map(|x_opt| f(x_opt.as_ref())).collect(),
255            elem_count: self.elem_count,
256        };
257        ret.elem_count = ret.iter().count();
258        ret
259    }
260
261    /// Iter over the nonempty slots.
262    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &T> + Clone {
263        self.vector.iter().filter_map(|opt| opt.as_ref())
264    }
265
266    pub fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut T> {
267        self.vector.iter_mut().filter_map(|opt| opt.as_mut())
268    }
269
270    pub fn iter_enumerated(&self) -> impl Iterator<Item = (I, &T)> {
271        self.vector
272            .iter_enumerated()
273            .flat_map(|(i, opt)| Some((i, opt.as_ref()?)))
274    }
275    pub fn iter_indexed(&self) -> impl Iterator<Item = (I, &T)> {
276        self.iter_enumerated()
277    }
278    pub fn iter_indexed_values(&self) -> impl Iterator<Item = (I, &T)> {
279        self.iter_indexed()
280    }
281
282    pub fn iter_mut_enumerated(&mut self) -> impl Iterator<Item = (I, &mut T)> {
283        self.vector
284            .iter_mut_enumerated()
285            .flat_map(|(i, opt)| Some((i, opt.as_mut()?)))
286    }
287    pub fn iter_mut_indexed(&mut self) -> impl Iterator<Item = (I, &mut T)> {
288        self.iter_mut_enumerated()
289    }
290
291    pub fn into_iter_enumerated(self) -> impl Iterator<Item = (I, T)> {
292        self.vector
293            .into_iter_enumerated()
294            .flat_map(|(i, opt)| Some((i, opt?)))
295    }
296    pub fn into_iter_indexed(self) -> impl Iterator<Item = (I, T)> {
297        self.into_iter_enumerated()
298    }
299    pub fn into_iter_indexed_values(self) -> impl Iterator<Item = (I, T)> {
300        self.into_iter_indexed()
301    }
302
303    /// Iterate over all slots, even empty ones.
304    pub fn iter_all_slots(&self) -> impl Iterator<Item = &Option<T>> {
305        self.vector.iter()
306    }
307
308    pub fn iter_enumerated_all_slots(&self) -> impl Iterator<Item = (I, &Option<T>)> {
309        self.vector.iter_enumerated()
310    }
311    pub fn iter_indexed_all_slots(&self) -> impl Iterator<Item = (I, &Option<T>)> {
312        self.iter_enumerated_all_slots()
313    }
314
315    pub fn iter_indices(&self) -> impl Iterator<Item = I> + '_ {
316        // Reuse `iter_indexed` to filter only the filled indices.
317        self.iter_indexed().map(|(id, _)| id)
318    }
319
320    pub fn all_indices(&self) -> impl Iterator<Item = I> + use<I, T> {
321        self.vector.indices()
322    }
323
324    /// Remove matching items and return and iterator over the removed items. This is lazy: items
325    /// are only removed as the iterator is consumed.
326    pub fn extract<'a, F: FnMut(I, &mut T) -> bool>(
327        &'a mut self,
328        mut f: F,
329    ) -> impl Iterator<Item = (I, T)> + use<'a, I, T, F> {
330        let elem_count = &mut self.elem_count;
331        self.vector
332            .iter_mut_enumerated()
333            .filter_map(move |(i, opt)| {
334                if f(i, opt.as_mut()?) {
335                    *elem_count -= 1;
336                    let elem = opt.take()?;
337                    Some((i, elem))
338                } else {
339                    None
340                }
341            })
342    }
343
344    /// Remove the elements that don't match the predicate.
345    pub fn retain(&mut self, mut f: impl FnMut(I, &mut T) -> bool) {
346        self.extract(|i, x| !f(i, x)).for_each(drop);
347    }
348
349    /// Like `Vec::clear`.
350    pub fn clear(&mut self) {
351        self.vector.clear();
352        self.elem_count = 0;
353    }
354    /// Like `Vec::truncate`.
355    pub fn truncate(&mut self, at: usize) {
356        self.vector.truncate(at);
357        self.elem_count = self.iter().count();
358    }
359    /// Like `Vec::split_off`.
360    pub fn split_off(&mut self, at: usize) -> Self {
361        let mut ret = Self {
362            vector: self.vector.split_off(I::from_usize(at)),
363            elem_count: 0,
364        };
365        self.elem_count = self.iter().count();
366        ret.elem_count = ret.iter().count();
367        ret
368    }
369
370    pub fn make_contiguous(self) -> crate::ids::IndexVec<I, T> {
371        // Ensure that every slot is filled.
372        assert_eq!(
373            self.elem_count(),
374            self.slot_count(),
375            "`IndexMap` is not contiguous"
376        );
377        self.into_iter().collect()
378    }
379}
380
381impl<I: Idx, T> Default for IndexMap<I, T> {
382    fn default() -> Self {
383        Self::new()
384    }
385}
386
387impl<I: std::fmt::Debug, T: std::fmt::Debug> std::fmt::Debug for IndexMap<I, T>
388where
389    I: Idx,
390{
391    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
392        <IndexVec<_, _> as std::fmt::Debug>::fmt(&self.vector, f)
393    }
394}
395
396impl<I, R, T> Index<R> for IndexMap<I, T>
397where
398    I: Idx,
399    R: IdxSliceIndex<I, Option<T>, Output = Option<T>>,
400{
401    type Output = T;
402    fn index(&self, index: R) -> &Self::Output {
403        self.vector[index].as_ref().unwrap()
404    }
405}
406
407impl<I, R, T> IndexMut<R> for IndexMap<I, T>
408where
409    I: Idx,
410    R: IdxSliceIndex<I, Option<T>, Output = Option<T>>,
411{
412    fn index_mut(&mut self, index: R) -> &mut Self::Output {
413        self.vector[index].as_mut().unwrap()
414    }
415}
416
417impl<'a, I, T> IntoIterator for &'a IndexMap<I, T>
418where
419    I: Idx,
420{
421    type Item = &'a T;
422    type IntoIter = impl Iterator<Item = &'a T>;
423
424    fn into_iter(self) -> Self::IntoIter {
425        self.vector.iter().flat_map(|opt| opt.as_ref())
426    }
427}
428
429impl<'a, I, T> IntoIterator for &'a mut IndexMap<I, T>
430where
431    I: Idx,
432{
433    type Item = &'a mut T;
434    type IntoIter = impl Iterator<Item = &'a mut T>;
435
436    fn into_iter(self) -> Self::IntoIter {
437        self.vector.iter_mut().flat_map(|opt| opt.as_mut())
438    }
439}
440
441impl<I, T> IntoIterator for IndexMap<I, T>
442where
443    I: Idx,
444{
445    type Item = T;
446    type IntoIter = impl DoubleEndedIterator<Item = T>;
447
448    fn into_iter(self) -> Self::IntoIter {
449        self.vector.into_iter().flatten()
450    }
451}
452
453impl<I, T> FromIterator<T> for IndexMap<I, T>
454where
455    I: Idx,
456{
457    #[inline]
458    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> IndexMap<I, T> {
459        let mut elem_count = 0;
460        let vector = IndexVec::from_iter(iter.into_iter().inspect(|_| elem_count += 1).map(Some));
461        IndexMap { vector, elem_count }
462    }
463}
464
465impl<I: Idx, T: Serialize> Serialize for IndexMap<I, T> {
466    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
467    where
468        S: Serializer,
469    {
470        self.vector.serialize(serializer)
471    }
472}
473
474impl<I: Idx, State, T: SerializeState<State>> SerializeState<State> for IndexMap<I, T> {
475    fn serialize_state<S>(&self, state: &State, serializer: S) -> Result<S::Ok, S::Error>
476    where
477        S: Serializer,
478    {
479        self.vector.as_vec().serialize_state(state, serializer)
480    }
481}
482
483impl<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for IndexMap<I, T> {
484    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
485    where
486        D: serde::Deserializer<'de>,
487    {
488        let mut ret = Self {
489            vector: Deserialize::deserialize(deserializer)?,
490            elem_count: 0,
491        };
492        ret.elem_count = ret.iter().count();
493        Ok(ret)
494    }
495}
496
497impl<'de, I: Idx, State, T: DeserializeState<'de, State>> DeserializeState<'de, State>
498    for IndexMap<I, T>
499{
500    fn deserialize_state<D>(state: &State, deserializer: D) -> Result<Self, D::Error>
501    where
502        D: serde::Deserializer<'de>,
503    {
504        let vec: Vec<Option<_>> = DeserializeState::deserialize_state(state, deserializer)?;
505        let mut ret = Self {
506            vector: IndexVec::from(vec),
507            elem_count: 0,
508        };
509        ret.elem_count = ret.iter().count();
510        Ok(ret)
511    }
512}
513
514impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for IndexMap<I, T> {
515    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
516        for x in self {
517            v.visit(x)?;
518        }
519        Continue(())
520    }
521}
522impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for IndexMap<I, T> {
523    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
524        for x in self {
525            v.visit(x)?;
526        }
527        Continue(())
528    }
529}
charon_lib/ids/index_map.rs

charon_lib/ids/
index_map.rs
