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

charon_lib/ids/
index_map.rs
