charon_lib/ids/
index_vec.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
7pub use index_vec::Idx;
8use index_vec::IdxSliceIndex;
9use serde::{Deserialize, Serialize, Serializer};
10use serde_state::{DeserializeState, SerializeState};
11use std::{
12    iter::{FromIterator, IntoIterator},
13    ops::{ControlFlow, Deref, DerefMut, Index, IndexMut},
14};
15
16use derive_generic_visitor::*;
17
18/// Contiguous indexed vector.
19///
20// A thin wrapper around `IndexVec` that adds some methods and impos.
21#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
22pub struct IndexVec<I, T>
23where
24    I: Idx,
25{
26    vector: index_vec::IndexVec<I, T>,
27}
28
29impl<I: std::fmt::Debug, T: std::fmt::Debug> std::fmt::Debug for IndexVec<I, T>
30where
31    I: Idx,
32{
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        <index_vec::IndexVec<_, _> as std::fmt::Debug>::fmt(&self.vector, f)
35    }
36}
37
38impl<I: Idx, T> Deref for IndexVec<I, T> {
39    type Target = index_vec::IndexVec<I, T>;
40    fn deref(&self) -> &Self::Target {
41        &self.vector
42    }
43}
44impl<I: Idx, T> DerefMut for IndexVec<I, T> {
45    fn deref_mut(&mut self) -> &mut Self::Target {
46        &mut self.vector
47    }
48}
49
50impl<I, T> IndexVec<I, T>
51where
52    I: Idx,
53{
54    pub fn new() -> Self {
55        IndexVec {
56            vector: index_vec::IndexVec::new(),
57        }
58    }
59
60    pub fn with_capacity(capacity: usize) -> Self {
61        IndexVec {
62            vector: index_vec::IndexVec::with_capacity(capacity),
63        }
64    }
65
66    pub fn from_vec(v: Vec<T>) -> Self {
67        Self {
68            vector: index_vec::IndexVec::from_vec(v),
69        }
70    }
71
72    pub fn from_array<const N: usize>(v: [T; N]) -> Self {
73        v.into_iter().collect()
74    }
75
76    /// Shadow the `index_vec::IndexVec` method because it silently shifts ids.
77    pub fn remove(&mut self, _: I) -> Option<T>
78    where
79        // Make it not callable.
80        String: Copy,
81    {
82        panic!("remove")
83    }
84
85    pub fn push_with(&mut self, f: impl FnOnce(I) -> T) -> I {
86        let id = self.next_idx();
87        let x = f(id);
88        self.push(x);
89        id
90    }
91
92    pub fn extend_from_slice(&mut self, other: &Self)
93    where
94        T: Clone,
95    {
96        self.vector.extend_from_slice(&other.vector);
97    }
98
99    /// Insert a value at that index, shifting all the values with equal or larger indices.
100    pub fn insert_and_shift_ids(&mut self, id: I, x: T) {
101        self.vector.insert(id, x)
102    }
103
104    /// Map each entry to a new one, keeping the same ids.
105    pub fn map<U>(self, f: impl FnMut(T) -> U) -> IndexVec<I, U> {
106        IndexVec {
107            vector: self.vector.into_iter().map(f).collect(),
108        }
109    }
110
111    /// Map each entry to a new one, keeping the same ids.
112    pub fn map_ref<'a, U>(&'a self, f: impl FnMut(&'a T) -> U) -> IndexVec<I, U> {
113        IndexVec {
114            vector: self.vector.iter().map(f).collect(),
115        }
116    }
117
118    /// Map each entry to a new one, keeping the same ids.
119    pub fn map_ref_mut<'a, U>(&'a mut self, f: impl FnMut(&'a mut T) -> U) -> IndexVec<I, U> {
120        IndexVec {
121            vector: self.vector.iter_mut().map(f).collect(),
122        }
123    }
124
125    /// Map each entry to a new one, keeping the same ids.
126    pub fn map_ref_indexed<'a, U>(&'a self, mut f: impl FnMut(I, &'a T) -> U) -> IndexVec<I, U> {
127        IndexVec {
128            vector: self
129                .vector
130                .iter_enumerated()
131                .map(|(i, x)| f(i, x))
132                .collect(),
133        }
134    }
135
136    // TODO: rename once we've migrated from `IndexMap` completely.
137    pub fn iter_indexed(&self) -> impl Iterator<Item = (I, &T)> {
138        self.vector.iter_enumerated()
139    }
140
141    pub fn iter_mut_indexed(&mut self) -> impl Iterator<Item = (I, &mut T)> {
142        self.vector.iter_mut_enumerated()
143    }
144
145    pub fn into_iter_indexed(self) -> impl Iterator<Item = (I, T)> {
146        self.vector.into_iter_enumerated()
147    }
148
149    pub fn iter_indexed_values(&self) -> impl Iterator<Item = (I, &T)> {
150        self.iter_indexed()
151    }
152
153    pub fn into_iter_indexed_values(self) -> impl Iterator<Item = (I, T)> {
154        self.into_iter_indexed()
155    }
156
157    pub fn iter_indices(&self) -> impl Iterator<Item = I> + '_ {
158        self.vector.indices()
159    }
160
161    pub fn all_indices(&self) -> impl Iterator<Item = I> + use<I, T> {
162        self.vector.indices()
163    }
164
165    /// Like `Vec::split_off`.
166    pub fn split_off(&mut self, at: usize) -> Self {
167        Self {
168            vector: self.vector.split_off(I::from_usize(at)),
169        }
170    }
171}
172
173impl<I: Idx, T> Default for IndexVec<I, T> {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178
179impl<I, R, T> Index<R> for IndexVec<I, T>
180where
181    I: Idx,
182    R: IdxSliceIndex<I, T, Output = T>,
183{
184    type Output = T;
185    fn index(&self, index: R) -> &Self::Output {
186        &self.vector[index]
187    }
188}
189
190impl<I, R, T> IndexMut<R> for IndexVec<I, T>
191where
192    I: Idx,
193    R: IdxSliceIndex<I, T, Output = T>,
194{
195    fn index_mut(&mut self, index: R) -> &mut Self::Output {
196        &mut self.vector[index]
197    }
198}
199
200impl<'a, I, T> IntoIterator for &'a IndexVec<I, T>
201where
202    I: Idx,
203{
204    type Item = &'a T;
205    type IntoIter = impl Iterator<Item = &'a T>;
206
207    fn into_iter(self) -> Self::IntoIter {
208        self.vector.iter()
209    }
210}
211
212impl<'a, I, T> IntoIterator for &'a mut IndexVec<I, T>
213where
214    I: Idx,
215{
216    type Item = &'a mut T;
217    type IntoIter = impl Iterator<Item = &'a mut T>;
218
219    fn into_iter(self) -> Self::IntoIter {
220        self.vector.iter_mut()
221    }
222}
223
224impl<I, T> IntoIterator for IndexVec<I, T>
225where
226    I: Idx,
227{
228    type Item = T;
229    type IntoIter = impl Iterator<Item = T>;
230
231    fn into_iter(self) -> Self::IntoIter {
232        self.vector.into_iter()
233    }
234}
235
236// FIXME: this impl is a footgun
237impl<I, T> FromIterator<T> for IndexVec<I, T>
238where
239    I: Idx,
240{
241    #[inline]
242    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> IndexVec<I, T> {
243        IndexVec {
244            vector: index_vec::IndexVec::from_iter(iter),
245        }
246    }
247}
248
249impl<I: Idx, T: Serialize> Serialize for IndexVec<I, T> {
250    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
251    where
252        S: Serializer,
253    {
254        self.vector.serialize(serializer)
255    }
256}
257
258impl<I: Idx, State, T: SerializeState<State>> SerializeState<State> for IndexVec<I, T> {
259    fn serialize_state<S>(&self, state: &State, serializer: S) -> Result<S::Ok, S::Error>
260    where
261        S: Serializer,
262    {
263        self.vector.as_vec().serialize_state(state, serializer)
264    }
265}
266
267impl<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for IndexVec<I, T> {
268    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
269    where
270        D: serde::Deserializer<'de>,
271    {
272        Ok(Self {
273            vector: Deserialize::deserialize(deserializer)?,
274        })
275    }
276}
277
278impl<'de, I: Idx, State, T: DeserializeState<'de, State>> DeserializeState<'de, State>
279    for IndexVec<I, T>
280{
281    fn deserialize_state<D>(state: &State, deserializer: D) -> Result<Self, D::Error>
282    where
283        D: serde::Deserializer<'de>,
284    {
285        let vec: Vec<_> = DeserializeState::deserialize_state(state, deserializer)?;
286        Ok(Self {
287            vector: index_vec::IndexVec::from(vec),
288        })
289    }
290}
291
292impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for IndexVec<I, T> {
293    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
294        for x in self {
295            v.visit(x)?;
296        }
297        Continue(())
298    }
299}
300impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for IndexVec<I, T> {
301    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
302        for x in self {
303            v.visit(x)?;
304        }
305        Continue(())
306    }
307}