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    fn index(&self, index: R) -> &Self::Output {
175        &self.vector[index]
176    }
177}
178
179impl<I, R, T> IndexMut<R> for IndexVec<I, T>
180where
181    I: Idx,
182    R: IdxSliceIndex<I, T, Output = T>,
183{
184    fn index_mut(&mut self, index: R) -> &mut Self::Output {
185        &mut self.vector[index]
186    }
187}
188
189impl<'a, I, T> IntoIterator for &'a IndexVec<I, T>
190where
191    I: Idx,
192{
193    type Item = &'a T;
194    type IntoIter = <&'a index_vec::IndexVec<I, T> as IntoIterator>::IntoIter;
195
196    fn into_iter(self) -> Self::IntoIter {
197        self.vector.iter()
198    }
199}
200
201impl<'a, I, T> IntoIterator for &'a mut IndexVec<I, T>
202where
203    I: Idx,
204{
205    type Item = &'a mut T;
206    type IntoIter = <&'a mut index_vec::IndexVec<I, T> as IntoIterator>::IntoIter;
207
208    fn into_iter(self) -> Self::IntoIter {
209        self.vector.iter_mut()
210    }
211}
212
213impl<I, T> IntoIterator for IndexVec<I, T>
214where
215    I: Idx,
216{
217    type Item = T;
218    type IntoIter = <index_vec::IndexVec<I, T> as IntoIterator>::IntoIter;
219
220    fn into_iter(self) -> Self::IntoIter {
221        self.vector.into_iter()
222    }
223}
224
225impl<I, T> FromIterator<T> for IndexVec<I, T>
226where
227    I: Idx,
228{
229    #[inline]
230    fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> IndexVec<I, T> {
231        IndexVec {
232            vector: index_vec::IndexVec::from_iter(iter),
233        }
234    }
235}
236
237impl<I: Idx, T: Serialize> Serialize for IndexVec<I, T> {
238    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
239    where
240        S: Serializer,
241    {
242        self.vector.serialize(serializer)
243    }
244}
245
246impl<I: Idx, State, T: SerializeState<State>> SerializeState<State> for IndexVec<I, T> {
247    fn serialize_state<S>(&self, state: &State, serializer: S) -> Result<S::Ok, S::Error>
248    where
249        S: Serializer,
250    {
251        self.vector.as_vec().serialize_state(state, serializer)
252    }
253}
254
255impl<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for IndexVec<I, T> {
256    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
257    where
258        D: serde::Deserializer<'de>,
259    {
260        Ok(Self {
261            vector: Deserialize::deserialize(deserializer)?,
262        })
263    }
264}
265
266impl<'de, I: Idx, State, T: DeserializeState<'de, State>> DeserializeState<'de, State>
267    for IndexVec<I, T>
268{
269    fn deserialize_state<D>(state: &State, deserializer: D) -> Result<Self, D::Error>
270    where
271        D: serde::Deserializer<'de>,
272    {
273        let vec: Vec<_> = DeserializeState::deserialize_state(state, deserializer)?;
274        Ok(Self {
275            vector: index_vec::IndexVec::from(vec),
276        })
277    }
278}
279
280impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for IndexVec<I, T> {
281    fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
282        for x in self {
283            v.visit(x)?;
284        }
285        Continue(())
286    }
287}
288impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for IndexVec<I, T> {
289    fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
290        for x in self {
291            v.visit(x)?;
292        }
293        Continue(())
294    }
295}
296
297impl<I, T> From<Vec<T>> for IndexVec<I, T>
298where
299    I: Idx,
300{
301    fn from(v: Vec<T>) -> Self {
302        v.into_iter().collect()
303    }
304}
305impl<I, T, const N: usize> From<[T; N]> for IndexVec<I, T>
306where
307    I: Idx,
308{
309    fn from(v: [T; N]) -> Self {
310        v.into_iter().collect()
311    }
312}