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 std::{
13    iter::{FromIterator, IntoIterator},
14    mem,
15    ops::{ControlFlow, Deref, Index, IndexMut},
16};
17
18use derive_generic_visitor::*;
19
20/// Indexed vector.
21/// To prevent accidental id reuse, the vector supports reserving a slot to be filled later.
22#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
23pub struct Vector<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: std::fmt::Debug, T: std::fmt::Debug> std::fmt::Debug for Vector<I, T>
33where
34    I: Idx,
35{
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        <IndexVec<_, _> as std::fmt::Debug>::fmt(&self.vector, f)
38    }
39}
40
41pub struct ReservedSlot<I: Idx>(I);
42
43impl<I, T> Vector<I, T>
44where
45    I: Idx,
46{
47    pub fn new() -> Self {
48        Vector {
49            vector: IndexVec::new(),
50            elem_count: 0,
51        }
52    }
53
54    pub fn with_capacity(capacity: usize) -> Self {
55        Vector {
56            vector: IndexVec::with_capacity(capacity),
57            elem_count: 0,
58        }
59    }
60
61    pub fn get(&self, i: I) -> Option<&T> {
62        self.vector.get(i).map(Option::as_ref).flatten()
63    }
64
65    pub fn get_mut(&mut self, i: I) -> Option<&mut T> {
66        self.vector.get_mut(i).map(Option::as_mut).flatten()
67    }
68
69    pub fn is_empty(&self) -> bool {
70        self.elem_count == 0
71    }
72
73    /// The number of elements stored in the vector.
74    pub fn elem_count(&self) -> usize {
75        self.elem_count
76    }
77
78    /// The number of slots allocated in the vector (empty or not).
79    pub fn slot_count(&self) -> usize {
80        self.vector.len()
81    }
82
83    /// Gets the value of the next available id. Avoid if possible; use `reserve_slot` instead.
84    pub fn next_id(&self) -> I {
85        self.vector.next_idx()
86    }
87
88    /// Reserve a spot in the vector.
89    pub fn reserve_slot(&mut self) -> I {
90        // Push a `None` to ensure we don't reuse the id.
91        self.vector.push(None)
92    }
93
94    /// Fill the reserved slot.
95    pub fn set_slot(&mut self, id: I, x: T) {
96        assert!(self.vector[id].is_none());
97        self.vector[id] = Some(x);
98        self.elem_count += 1;
99    }
100
101    /// Remove the value from this slot, leaving other ids unchanged.
102    pub fn remove(&mut self, id: I) -> Option<T> {
103        if self.vector[id].is_some() {
104            self.elem_count -= 1;
105        }
106        self.vector[id].take()
107    }
108
109    /// Remove the value from this slot, shifting other ids as needed.
110    pub fn remove_and_shift_ids(&mut self, id: I) -> Option<T> {
111        if id.index() >= self.slot_count() {
112            return None;
113        }
114        if self.vector[id].is_some() {
115            self.elem_count -= 1;
116        }
117        self.vector.remove(id)
118    }
119
120    /// Remove the last slot.
121    pub fn pop(&mut self) -> Option<T> {
122        if self.vector.last().is_some() {
123            self.elem_count -= 1;
124        }
125        self.vector.pop().flatten()
126    }
127
128    pub fn push(&mut self, x: T) -> I {
129        self.elem_count += 1;
130        self.vector.push(Some(x))
131    }
132
133    pub fn push_with(&mut self, f: impl FnOnce(I) -> T) -> I {
134        let id = self.reserve_slot();
135        let x = f(id);
136        self.set_slot(id, x);
137        id
138    }
139
140    pub fn push_all<It>(&mut self, it: It) -> impl Iterator<Item = I> + use<'_, I, T, It>
141    where
142        It: IntoIterator<Item = T>,
143    {
144        it.into_iter().map(move |x| self.push(x))
145    }
146
147    pub fn extend<It>(&mut self, it: It)
148    where
149        It: IntoIterator<Item = T>,
150    {
151        self.push_all(it).for_each(|_| ())
152    }
153
154    pub fn extend_from_slice(&mut self, other: &Self)
155    where
156        T: Clone,
157    {
158        self.vector.extend_from_slice(&other.vector);
159        self.elem_count += other.elem_count;
160    }
161
162    /// Insert a value at that index, shifting all the values with equal or larger indices.
163    pub fn insert_and_shift_ids(&mut self, id: I, x: T) {
164        self.elem_count += 1;
165        self.vector.insert(id, Some(x))
166    }
167
168    /// Get a mutable reference into the ith element. If the vector is too short, extend it until
169    /// it has enough elements. If the element doesn't exist, use the provided function to
170    /// initialize it.
171    pub fn get_or_extend_and_insert(&mut self, id: I, f: impl FnOnce() -> T) -> &mut T {
172        if id.index() >= self.vector.len() {
173            self.vector.resize_with(id.index() + 1, || None);
174        }
175        self.vector[id].get_or_insert_with(f)
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_extend_and_insert_default(&mut self, id: I) -> &mut T
182    where
183        T: Default,
184    {
185        self.get_or_extend_and_insert(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) -> Vector<I, U> {
190        Vector {
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) -> Vector<I, U> {
202        Vector {
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) -> Vector<I, U> {
214        Vector {
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) -> Vector<I, U> {
226        Vector {
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) -> Vector<I, U> {
238        Vector {
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>) -> Vector<I, U> {
250        Vector {
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    ) -> Vector<I, U> {
261        let mut ret = Vector {
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 Iterator<Item = &T> + DoubleEndedIterator + Clone {
271        self.vector.iter().filter_map(|opt| opt.as_ref())
272    }
273
274    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> + DoubleEndedIterator {
275        self.vector.iter_mut().filter_map(|opt| opt.as_mut())
276    }
277
278    pub fn iter_indexed(&self) -> impl Iterator<Item = (I, &T)> {
279        self.vector
280            .iter_enumerated()
281            .flat_map(|(i, opt)| Some((i, opt.as_ref()?)))
282    }
283
284    pub fn iter_mut_indexed(&mut self) -> impl Iterator<Item = (I, &mut T)> {
285        self.vector
286            .iter_mut_enumerated()
287            .flat_map(|(i, opt)| Some((i, opt.as_mut()?)))
288    }
289
290    pub fn into_iter_indexed(self) -> impl Iterator<Item = (I, T)> {
291        self.vector
292            .into_iter_enumerated()
293            .flat_map(|(i, opt)| Some((i, opt?)))
294    }
295
296    pub fn iter_indexed_values(&self) -> impl Iterator<Item = (I, &T)> {
297        self.iter_indexed()
298    }
299
300    pub fn into_iter_indexed_values(self) -> impl Iterator<Item = (I, T)> {
301        self.into_iter_indexed()
302    }
303
304    /// Iterate over all slots, even empty ones.
305    pub fn iter_all_slots(&self) -> impl Iterator<Item = &Option<T>> {
306        self.vector.iter()
307    }
308
309    pub fn iter_indexed_all_slots(&self) -> impl Iterator<Item = (I, &Option<T>)> {
310        self.vector.iter_enumerated()
311    }
312
313    pub fn iter_indices(&self) -> impl Iterator<Item = I> + '_ {
314        // Reuse `iter_indexed` to filter only the filled indices.
315        self.iter_indexed().map(|(id, _)| id)
316    }
317
318    pub fn all_indices(&self) -> impl Iterator<Item = I> + use<I, T> {
319        self.vector.indices()
320    }
321
322    /// Remove matching items and return and iterator over the removed items. This is lazy: items
323    /// are only removed as the iterator is consumed.
324    pub fn extract<'a, F: FnMut(&mut T) -> bool>(
325        &'a mut self,
326        mut f: F,
327    ) -> impl Iterator<Item = (I, T)> + use<'a, I, T, F> {
328        let elem_count = &mut self.elem_count;
329        self.vector
330            .iter_mut_enumerated()
331            .filter_map(move |(i, opt)| {
332                if f(opt.as_mut()?) {
333                    *elem_count -= 1;
334                    let elem = mem::replace(opt, None)?;
335                    Some((i, elem))
336                } else {
337                    None
338                }
339            })
340    }
341
342    /// Remove the elements that don't match the predicate.
343    pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
344        self.extract(|x| !f(x)).for_each(drop);
345    }
346
347    /// Like `Vec::clear`.
348    pub fn clear(&mut self) {
349        self.vector.clear();
350        self.elem_count = 0;
351    }
352    /// Like `Vec::truncate`.
353    pub fn truncate(&mut self, at: usize) {
354        self.vector.truncate(at);
355        self.elem_count = self.iter().count();
356    }
357    /// Like `Vec::split_off`.
358    pub fn split_off(&mut self, at: usize) -> Self {
359        let mut ret = Self {
360            vector: self.vector.split_off(I::from_usize(at)),
361            elem_count: 0,
362        };
363        self.elem_count = self.iter().count();
364        ret.elem_count = ret.iter().count();
365        ret
366    }
367}
368
369impl<I: Idx, T> Default for Vector<I, T> {
370    fn default() -> Self {
371        Self::new()
372    }
373}
374
375impl<I: Idx> Deref for ReservedSlot<I> {
376    type Target = I;
377    fn deref(&self) -> &Self::Target {
378        &self.0
379    }
380}
381
382impl<I, R, T> Index<R> for Vector<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 Vector<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 Vector<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 Vector<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 Vector<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 Vector<I, T>
441where
442    I: Idx,
443{
444    #[inline]
445    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> Vector<I, T> {
446        let mut elem_count = 0;
447        let vector = IndexVec::from_iter(iter.into_iter().inspect(|_| elem_count += 1).map(Some));
448        Vector { vector, elem_count }
449    }
450}
451
452// FIXME: this impl is a footgun
453impl<I, T> From<Vec<T>> for Vector<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 Vector<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 Vector<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<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for Vector<I, T> {
482    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
483    where
484        D: serde::Deserializer<'de>,
485    {
486        let mut ret = Self {
487            vector: Deserialize::deserialize(deserializer)?,
488            elem_count: 0,
489        };
490        ret.elem_count = ret.iter().count();
491        Ok(ret)
492    }
493}
494
495impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for Vector<I, T> {
496    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
497        for x in self {
498            v.visit(x)?;
499        }
500        Continue(())
501    }
502}
503impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for Vector<I, T> {
504    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
505        for x in self {
506            v.visit(x)?;
507        }
508        Continue(())
509    }
510}
charon_lib/ids/vector.rs

charon_lib/ids/
vector.rs
