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

charon_lib/ids/
vector.rs
