Skip to main content

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    /// Shadow the `index_vec::IndexVec` method because it silently shifts ids.
67    pub fn remove(&mut self, _: I) -> Option<T>
68    where
69        // Make it not callable.
70        String: Copy,
71    {
72        panic!("remove")
73    }
74
75    pub fn remove_and_shift_ids(&mut self, id: I) -> Option<T> {
76        if self.vector.get(id).is_some() {
77            Some(self.vector.remove(id))
78        } else {
79            None
80        }
81    }
82
83    pub fn push_with(&mut self, f: impl FnOnce(I) -> T) -> I {
84        let id = self.next_idx();
85        let x = f(id);
86        self.push(x);
87        id
88    }
89
90    pub fn clone_extend_from_other(&mut self, other: &Self)
91    where
92        T: Clone,
93    {
94        self.vector.extend_from_slice(&other.vector);
95    }
96
97    /// Get a mutable reference into the ith element. If the vector is too short, extend it until
98    /// it has enough elements.
99    pub fn get_or_extend_and_insert(&mut self, id: I, f: impl FnMut() -> T) -> &mut T {
100        if id.index() >= self.vector.len() {
101            self.vector.resize_with(id.index() + 1, f);
102        }
103        &mut self.vector[id]
104    }
105
106    /// Map each entry to a new one, keeping the same ids.
107    pub fn map<U>(self, f: impl FnMut(T) -> U) -> IndexVec<I, U> {
108        IndexVec {
109            vector: self.vector.into_iter().map(f).collect(),
110        }
111    }
112    /// Map each entry to a new one, keeping the same ids.
113    pub fn map_ref<'a, U>(&'a self, f: impl FnMut(&'a T) -> U) -> IndexVec<I, U> {
114        IndexVec {
115            vector: self.vector.iter().map(f).collect(),
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_indexed<U>(self, mut f: impl FnMut(I, T) -> U) -> IndexVec<I, U> {
127        IndexVec {
128            vector: self
129                .vector
130                .into_iter_enumerated()
131                .map(|(i, x)| f(i, x))
132                .collect(),
133        }
134    }
135    /// Map each entry to a new one, keeping the same ids.
136    pub fn map_ref_indexed<'a, U>(&'a self, mut f: impl FnMut(I, &'a T) -> U) -> IndexVec<I, U> {
137        IndexVec {
138            vector: self
139                .vector
140                .iter_enumerated()
141                .map(|(i, x)| f(i, x))
142                .collect(),
143        }
144    }
145
146    pub fn into_iter_enumerated(self) -> impl Iterator<Item = (I, T)> {
147        self.vector.into_iter_enumerated()
148    }
149
150    /// Like `Vec::split_off`.
151    pub fn split_off(&mut self, at: usize) -> Self {
152        Self {
153            vector: self.vector.split_off(I::from_usize(at)),
154        }
155    }
156}
157
158impl<I: Idx, T> Default for IndexVec<I, T> {
159    fn default() -> Self {
160        Self::new()
161    }
162}
163
164impl<I, R, T> Index<R> for IndexVec<I, T>
165where
166    I: Idx,
167    R: IdxSliceIndex<I, T, Output = T>,
168{
169    type Output = T;
170    fn index(&self, index: R) -> &Self::Output {
171        &self.vector[index]
172    }
173}
174
175impl<I, R, T> IndexMut<R> for IndexVec<I, T>
176where
177    I: Idx,
178    R: IdxSliceIndex<I, T, Output = T>,
179{
180    fn index_mut(&mut self, index: R) -> &mut Self::Output {
181        &mut self.vector[index]
182    }
183}
184
185impl<'a, I, T> IntoIterator for &'a IndexVec<I, T>
186where
187    I: Idx,
188{
189    type Item = &'a T;
190    type IntoIter = impl Iterator<Item = &'a T>;
191
192    fn into_iter(self) -> Self::IntoIter {
193        self.vector.iter()
194    }
195}
196
197impl<'a, I, T> IntoIterator for &'a mut IndexVec<I, T>
198where
199    I: Idx,
200{
201    type Item = &'a mut T;
202    type IntoIter = impl Iterator<Item = &'a mut T>;
203
204    fn into_iter(self) -> Self::IntoIter {
205        self.vector.iter_mut()
206    }
207}
208
209impl<I, T> IntoIterator for IndexVec<I, T>
210where
211    I: Idx,
212{
213    type Item = T;
214    type IntoIter = impl DoubleEndedIterator<Item = T>;
215
216    fn into_iter(self) -> Self::IntoIter {
217        self.vector.into_iter()
218    }
219}
220
221impl<I, T> FromIterator<T> for IndexVec<I, T>
222where
223    I: Idx,
224{
225    #[inline]
226    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> IndexVec<I, T> {
227        IndexVec {
228            vector: index_vec::IndexVec::from_iter(iter),
229        }
230    }
231}
232
233impl<I: Idx, T: Serialize> Serialize for IndexVec<I, T> {
234    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
235    where
236        S: Serializer,
237    {
238        self.vector.serialize(serializer)
239    }
240}
241
242impl<I: Idx, State, T: SerializeState<State>> SerializeState<State> for IndexVec<I, T> {
243    fn serialize_state<S>(&self, state: &State, serializer: S) -> Result<S::Ok, S::Error>
244    where
245        S: Serializer,
246    {
247        self.vector.as_vec().serialize_state(state, serializer)
248    }
249}
250
251impl<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for IndexVec<I, T> {
252    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
253    where
254        D: serde::Deserializer<'de>,
255    {
256        Ok(Self {
257            vector: Deserialize::deserialize(deserializer)?,
258        })
259    }
260}
261
262impl<'de, I: Idx, State, T: DeserializeState<'de, State>> DeserializeState<'de, State>
263    for IndexVec<I, T>
264{
265    fn deserialize_state<D>(state: &State, deserializer: D) -> Result<Self, D::Error>
266    where
267        D: serde::Deserializer<'de>,
268    {
269        let vec: Vec<_> = DeserializeState::deserialize_state(state, deserializer)?;
270        Ok(Self {
271            vector: index_vec::IndexVec::from(vec),
272        })
273    }
274}
275
276impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for IndexVec<I, T> {
277    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
278        for x in self {
279            v.visit(x)?;
280        }
281        Continue(())
282    }
283}
284impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for IndexVec<I, T> {
285    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
286        for x in self {
287            v.visit(x)?;
288        }
289        Continue(())
290    }
291}
292
293impl<I, T> From<Vec<T>> for IndexVec<I, T>
294where
295    I: Idx,
296{
297    fn from(v: Vec<T>) -> Self {
298        v.into_iter().collect()
299    }
300}
301impl<I, T, const N: usize> From<[T; N]> for IndexVec<I, T>
302where
303    I: Idx,
304{
305    fn from(v: [T; N]) -> Self {
306        v.into_iter().collect()
307    }
308}