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