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

charon_lib/ids/
index_map.rs
