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

charon_lib/ids/
index_map.rs
