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

charon_lib/ids/
vector.rs
