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 [ids::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    ops::{ControlFlow, Deref, Index, IndexMut},
15};
16
17use derive_generic_visitor::*;
18
19/// Indexed vector.
20/// To prevent accidental id reuse, the vector supports reserving a slot to be filled later.
21#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
22pub struct Vector<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 Vector<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
40pub struct ReservedSlot<I: Idx>(I);
41
42impl<I, T> Vector<I, T>
43where
44    I: Idx,
45{
46    pub fn new() -> Self {
47        Vector {
48            vector: IndexVec::new(),
49            elem_count: 0,
50        }
51    }
52
53    pub fn get(&self, i: I) -> Option<&T> {
54        self.vector.get(i).map(Option::as_ref).flatten()
55    }
56
57    pub fn get_mut(&mut self, i: I) -> Option<&mut T> {
58        self.vector.get_mut(i).map(Option::as_mut).flatten()
59    }
60
61    pub fn is_empty(&self) -> bool {
62        self.elem_count == 0
63    }
64
65    /// The number of elements stored in the vector.
66    pub fn elem_count(&self) -> usize {
67        self.elem_count
68    }
69
70    /// The number of slots allocated in the vector (empty or not).
71    pub fn slot_count(&self) -> usize {
72        self.vector.len()
73    }
74
75    /// Gets the value of the next available id. Avoid if possible; use `reserve_slot` instead.
76    pub fn next_id(&self) -> I {
77        self.vector.next_idx()
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.
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    pub fn push(&mut self, x: T) -> I {
102        self.elem_count += 1;
103        self.vector.push(Some(x))
104    }
105
106    pub fn push_with(&mut self, f: impl FnOnce(I) -> T) -> I {
107        let id = self.reserve_slot();
108        let x = f(id);
109        self.set_slot(id, x);
110        id
111    }
112
113    pub fn push_all<It>(&mut self, it: It) -> impl Iterator<Item = I> + use<'_, I, T, It>
114    where
115        It: Iterator<Item = T>,
116    {
117        it.map(move |x| self.push(x))
118    }
119
120    pub fn extend<It>(&mut self, it: It)
121    where
122        It: Iterator<Item = T>,
123    {
124        self.push_all(it).for_each(|_| ())
125    }
126
127    pub fn extend_from_slice(&mut self, other: &Self)
128    where
129        T: Clone,
130    {
131        self.vector.extend_from_slice(&other.vector);
132        self.elem_count += other.elem_count;
133    }
134
135    /// Get a mutable reference into the ith element. If the vector is too short, extend it until
136    /// it has enough elements. If the element doesn't exist, use the provided function to
137    /// initialize it.
138    pub fn get_or_extend_and_insert(&mut self, id: I, f: impl FnOnce() -> T) -> &mut T {
139        if id.index() >= self.vector.len() {
140            self.vector.resize_with(id.index() + 1, || None);
141        }
142        self.vector[id].get_or_insert_with(f)
143    }
144
145    /// Map each entry to a new one, keeping the same ids.
146    pub fn map<U>(self, mut f: impl FnMut(T) -> U) -> Vector<I, U> {
147        Vector {
148            vector: self
149                .vector
150                .into_iter()
151                .map(|x_opt| x_opt.map(&mut f))
152                .collect(),
153            elem_count: self.elem_count,
154        }
155    }
156
157    /// Map each entry to a new one, keeping the same ids.
158    pub fn map_ref<'a, U>(&'a self, mut f: impl FnMut(&'a T) -> U) -> Vector<I, U> {
159        Vector {
160            vector: self
161                .vector
162                .iter()
163                .map(|x_opt| x_opt.as_ref().map(&mut f))
164                .collect(),
165            elem_count: self.elem_count,
166        }
167    }
168
169    /// Map each entry to a new one, keeping the same ids.
170    pub fn map_ref_mut<'a, U>(&'a mut self, mut f: impl FnMut(&'a mut T) -> U) -> Vector<I, U> {
171        Vector {
172            vector: self
173                .vector
174                .iter_mut()
175                .map(|x_opt| x_opt.as_mut().map(&mut f))
176                .collect(),
177            elem_count: self.elem_count,
178        }
179    }
180
181    /// Map each entry to a new one, keeping the same ids.
182    pub fn map_ref_indexed<'a, U>(&'a self, mut f: impl FnMut(I, &'a T) -> U) -> Vector<I, U> {
183        Vector {
184            vector: self
185                .vector
186                .iter_enumerated()
187                .map(|(i, x_opt)| x_opt.as_ref().map(|x| f(i, x)))
188                .collect(),
189            elem_count: self.elem_count,
190        }
191    }
192
193    /// Map each entry to a new one, keeping the same ids. Includes empty slots.
194    pub fn map_opt<U>(self, f: impl FnMut(Option<T>) -> Option<U>) -> Vector<I, U> {
195        Vector {
196            vector: self.vector.into_iter().map(f).collect(),
197            elem_count: self.elem_count,
198        }
199    }
200
201    /// Map each entry to a new one, keeping the same ids. Includes empty slots.
202    pub fn map_ref_opt<'a, U>(
203        &'a self,
204        mut f: impl FnMut(Option<&'a T>) -> Option<U>,
205    ) -> Vector<I, U> {
206        let mut ret = Vector {
207            vector: self.vector.iter().map(|x_opt| f(x_opt.as_ref())).collect(),
208            elem_count: self.elem_count,
209        };
210        ret.elem_count = ret.iter().count();
211        ret
212    }
213
214    /// Iter over the nonempty slots.
215    pub fn iter(&self) -> impl Iterator<Item = &T> + DoubleEndedIterator + Clone {
216        self.vector.iter().filter_map(|opt| opt.as_ref())
217    }
218
219    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> + DoubleEndedIterator {
220        self.vector.iter_mut().filter_map(|opt| opt.as_mut())
221    }
222
223    pub fn iter_indexed(&self) -> impl Iterator<Item = (I, &T)> {
224        self.vector
225            .iter_enumerated()
226            .flat_map(|(i, opt)| Some((i, opt.as_ref()?)))
227    }
228
229    pub fn iter_mut_indexed(&mut self) -> impl Iterator<Item = (I, &mut T)> {
230        self.vector
231            .iter_mut_enumerated()
232            .flat_map(|(i, opt)| Some((i, opt.as_mut()?)))
233    }
234
235    pub fn into_iter_indexed(self) -> impl Iterator<Item = (I, T)> {
236        self.vector
237            .into_iter_enumerated()
238            .flat_map(|(i, opt)| Some((i, opt?)))
239    }
240
241    pub fn iter_indexed_values(&self) -> impl Iterator<Item = (I, &T)> {
242        self.iter_indexed()
243    }
244
245    pub fn into_iter_indexed_values(self) -> impl Iterator<Item = (I, T)> {
246        self.into_iter_indexed()
247    }
248
249    /// Iterate over all slots, even empty ones.
250    pub fn iter_all_slots(&self) -> impl Iterator<Item = &Option<T>> {
251        self.vector.iter()
252    }
253
254    pub fn iter_indexed_all_slots(&self) -> impl Iterator<Item = (I, &Option<T>)> {
255        self.vector.iter_enumerated()
256    }
257
258    pub fn iter_indices(&self) -> impl Iterator<Item = I> + '_ {
259        // Reuse `iter_indexed` to filter only the filled indices.
260        self.iter_indexed().map(|(id, _)| id)
261    }
262
263    pub fn all_indices(&self) -> impl Iterator<Item = I> {
264        self.vector.indices()
265    }
266
267    pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
268        self.vector.iter_mut().for_each(|opt| {
269            if let Some(x) = opt {
270                if !f(x) {
271                    *opt = None;
272                    self.elem_count -= 1;
273                }
274            }
275        });
276    }
277
278    /// Like `Vec::split_off`.
279    pub fn split_off(&mut self, at: usize) -> Self {
280        let mut ret = Self {
281            vector: self.vector.split_off(I::from_usize(at)),
282            elem_count: 0,
283        };
284        self.elem_count = self.iter().count();
285        ret.elem_count = ret.iter().count();
286        ret
287    }
288}
289
290impl<I: Idx, T> Default for Vector<I, T> {
291    fn default() -> Self {
292        Self::new()
293    }
294}
295
296impl<I: Idx> Deref for ReservedSlot<I> {
297    type Target = I;
298    fn deref(&self) -> &Self::Target {
299        &self.0
300    }
301}
302
303impl<I, R, T> Index<R> for Vector<I, T>
304where
305    I: Idx,
306    R: IdxSliceIndex<I, Option<T>, Output = Option<T>>,
307{
308    type Output = T;
309    fn index(&self, index: R) -> &Self::Output {
310        self.vector[index].as_ref().unwrap()
311    }
312}
313
314impl<I, R, T> IndexMut<R> for Vector<I, T>
315where
316    I: Idx,
317    R: IdxSliceIndex<I, Option<T>, Output = Option<T>>,
318{
319    fn index_mut(&mut self, index: R) -> &mut Self::Output {
320        self.vector[index].as_mut().unwrap()
321    }
322}
323
324impl<'a, I, T> IntoIterator for &'a Vector<I, T>
325where
326    I: Idx,
327{
328    type Item = &'a T;
329    type IntoIter = impl Iterator<Item = &'a T>;
330
331    fn into_iter(self) -> Self::IntoIter {
332        self.vector.iter().flat_map(|opt| opt.as_ref())
333    }
334}
335
336impl<'a, I, T> IntoIterator for &'a mut Vector<I, T>
337where
338    I: Idx,
339{
340    type Item = &'a mut T;
341    type IntoIter = impl Iterator<Item = &'a mut T>;
342
343    fn into_iter(self) -> Self::IntoIter {
344        self.vector.iter_mut().flat_map(|opt| opt.as_mut())
345    }
346}
347
348impl<I, T> IntoIterator for Vector<I, T>
349where
350    I: Idx,
351{
352    type Item = T;
353    type IntoIter = impl Iterator<Item = T>;
354
355    fn into_iter(self) -> Self::IntoIter {
356        self.vector.into_iter().flat_map(|opt| opt)
357    }
358}
359
360// FIXME: this impl is a footgun
361impl<I, T> FromIterator<T> for Vector<I, T>
362where
363    I: Idx,
364{
365    #[inline]
366    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> Vector<I, T> {
367        let mut elem_count = 0;
368        let vector = IndexVec::from_iter(iter.into_iter().inspect(|_| elem_count += 1).map(Some));
369        Vector { vector, elem_count }
370    }
371}
372
373// FIXME: this impl is a footgun
374impl<I, T> From<Vec<T>> for Vector<I, T>
375where
376    I: Idx,
377{
378    fn from(v: Vec<T>) -> Self {
379        v.into_iter().collect()
380    }
381}
382
383impl<I: Idx, T: Serialize> Serialize for Vector<I, T> {
384    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
385    where
386        S: Serializer,
387    {
388        self.vector.serialize(serializer)
389    }
390}
391
392impl<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for Vector<I, T> {
393    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
394    where
395        D: serde::Deserializer<'de>,
396    {
397        let mut ret = Self {
398            vector: Deserialize::deserialize(deserializer)?,
399            elem_count: 0,
400        };
401        ret.elem_count = ret.iter().count();
402        Ok(ret)
403    }
404}
405
406impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for Vector<I, T> {
407    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
408        for x in self {
409            v.visit(x)?;
410        }
411        Continue(())
412    }
413}
414impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for Vector<I, T> {
415    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
416        for x in self {
417            v.visit(x)?;
418        }
419        Continue(())
420    }
421}
charon_lib/ids/vector.rs

charon_lib/ids/
vector.rs
