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

charon_lib/ids/
index_map.rs
