1use 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#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
23pub struct Vector<I, T>
24where
25 I: Idx,
26{
27 vector: IndexVec<I, Option<T>>,
28 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 with_capacity(capacity: usize) -> Self {
55 Vector {
56 vector: IndexVec::with_capacity(capacity),
57 elem_count: 0,
58 }
59 }
60
61 pub fn get(&self, i: I) -> Option<&T> {
62 self.vector.get(i).map(Option::as_ref).flatten()
63 }
64
65 pub fn get_mut(&mut self, i: I) -> Option<&mut T> {
66 self.vector.get_mut(i).map(Option::as_mut).flatten()
67 }
68
69 pub fn is_empty(&self) -> bool {
70 self.elem_count == 0
71 }
72
73 pub fn elem_count(&self) -> usize {
75 self.elem_count
76 }
77
78 pub fn slot_count(&self) -> usize {
80 self.vector.len()
81 }
82
83 pub fn next_id(&self) -> I {
85 self.vector.next_idx()
86 }
87
88 pub fn reserve_slot(&mut self) -> I {
90 self.vector.push(None)
92 }
93
94 pub fn set_slot(&mut self, id: I, x: T) {
96 assert!(self.vector[id].is_none());
97 self.vector[id] = Some(x);
98 self.elem_count += 1;
99 }
100
101 pub fn remove(&mut self, id: I) -> Option<T> {
103 if self.vector[id].is_some() {
104 self.elem_count -= 1;
105 }
106 self.vector[id].take()
107 }
108
109 pub fn remove_and_shift_ids(&mut self, id: I) -> Option<T> {
111 if id.index() >= self.slot_count() {
112 return None;
113 }
114 if self.vector[id].is_some() {
115 self.elem_count -= 1;
116 }
117 self.vector.remove(id)
118 }
119
120 pub fn pop(&mut self) -> Option<T> {
122 if self.vector.last().is_some() {
123 self.elem_count -= 1;
124 }
125 self.vector.pop().flatten()
126 }
127
128 pub fn push(&mut self, x: T) -> I {
129 self.elem_count += 1;
130 self.vector.push(Some(x))
131 }
132
133 pub fn push_with(&mut self, f: impl FnOnce(I) -> T) -> I {
134 let id = self.reserve_slot();
135 let x = f(id);
136 self.set_slot(id, x);
137 id
138 }
139
140 pub fn push_all<It>(&mut self, it: It) -> impl Iterator<Item = I> + use<'_, I, T, It>
141 where
142 It: IntoIterator<Item = T>,
143 {
144 it.into_iter().map(move |x| self.push(x))
145 }
146
147 pub fn extend<It>(&mut self, it: It)
148 where
149 It: IntoIterator<Item = T>,
150 {
151 self.push_all(it).for_each(|_| ())
152 }
153
154 pub fn extend_from_slice(&mut self, other: &Self)
155 where
156 T: Clone,
157 {
158 self.vector.extend_from_slice(&other.vector);
159 self.elem_count += other.elem_count;
160 }
161
162 pub fn insert_and_shift_ids(&mut self, id: I, x: T) {
164 self.elem_count += 1;
165 self.vector.insert(id, Some(x))
166 }
167
168 pub fn get_or_extend_and_insert(&mut self, id: I, f: impl FnOnce() -> T) -> &mut T {
172 if id.index() >= self.vector.len() {
173 self.vector.resize_with(id.index() + 1, || None);
174 }
175 self.vector[id].get_or_insert_with(f)
176 }
177
178 pub fn map<U>(self, mut f: impl FnMut(T) -> U) -> Vector<I, U> {
180 Vector {
181 vector: self
182 .vector
183 .into_iter()
184 .map(|x_opt| x_opt.map(&mut f))
185 .collect(),
186 elem_count: self.elem_count,
187 }
188 }
189
190 pub fn map_ref<'a, U>(&'a self, mut f: impl FnMut(&'a T) -> U) -> Vector<I, U> {
192 Vector {
193 vector: self
194 .vector
195 .iter()
196 .map(|x_opt| x_opt.as_ref().map(&mut f))
197 .collect(),
198 elem_count: self.elem_count,
199 }
200 }
201
202 pub fn map_ref_mut<'a, U>(&'a mut self, mut f: impl FnMut(&'a mut T) -> U) -> Vector<I, U> {
204 Vector {
205 vector: self
206 .vector
207 .iter_mut()
208 .map(|x_opt| x_opt.as_mut().map(&mut f))
209 .collect(),
210 elem_count: self.elem_count,
211 }
212 }
213
214 pub fn map_indexed<U>(self, mut f: impl FnMut(I, T) -> U) -> Vector<I, U> {
216 Vector {
217 vector: self
218 .vector
219 .into_iter_enumerated()
220 .map(|(i, x_opt)| x_opt.map(|x| f(i, x)))
221 .collect(),
222 elem_count: self.elem_count,
223 }
224 }
225
226 pub fn map_ref_indexed<'a, U>(&'a self, mut f: impl FnMut(I, &'a T) -> U) -> Vector<I, U> {
228 Vector {
229 vector: self
230 .vector
231 .iter_enumerated()
232 .map(|(i, x_opt)| x_opt.as_ref().map(|x| f(i, x)))
233 .collect(),
234 elem_count: self.elem_count,
235 }
236 }
237
238 pub fn map_opt<U>(self, f: impl FnMut(Option<T>) -> Option<U>) -> Vector<I, U> {
240 Vector {
241 vector: self.vector.into_iter().map(f).collect(),
242 elem_count: self.elem_count,
243 }
244 }
245
246 pub fn map_ref_opt<'a, U>(
248 &'a self,
249 mut f: impl FnMut(Option<&'a T>) -> Option<U>,
250 ) -> Vector<I, U> {
251 let mut ret = Vector {
252 vector: self.vector.iter().map(|x_opt| f(x_opt.as_ref())).collect(),
253 elem_count: self.elem_count,
254 };
255 ret.elem_count = ret.iter().count();
256 ret
257 }
258
259 pub fn iter(&self) -> impl Iterator<Item = &T> + DoubleEndedIterator + Clone {
261 self.vector.iter().filter_map(|opt| opt.as_ref())
262 }
263
264 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> + DoubleEndedIterator {
265 self.vector.iter_mut().filter_map(|opt| opt.as_mut())
266 }
267
268 pub fn iter_indexed(&self) -> impl Iterator<Item = (I, &T)> {
269 self.vector
270 .iter_enumerated()
271 .flat_map(|(i, opt)| Some((i, opt.as_ref()?)))
272 }
273
274 pub fn iter_mut_indexed(&mut self) -> impl Iterator<Item = (I, &mut T)> {
275 self.vector
276 .iter_mut_enumerated()
277 .flat_map(|(i, opt)| Some((i, opt.as_mut()?)))
278 }
279
280 pub fn into_iter_indexed(self) -> impl Iterator<Item = (I, T)> {
281 self.vector
282 .into_iter_enumerated()
283 .flat_map(|(i, opt)| Some((i, opt?)))
284 }
285
286 pub fn iter_indexed_values(&self) -> impl Iterator<Item = (I, &T)> {
287 self.iter_indexed()
288 }
289
290 pub fn into_iter_indexed_values(self) -> impl Iterator<Item = (I, T)> {
291 self.into_iter_indexed()
292 }
293
294 pub fn iter_all_slots(&self) -> impl Iterator<Item = &Option<T>> {
296 self.vector.iter()
297 }
298
299 pub fn iter_indexed_all_slots(&self) -> impl Iterator<Item = (I, &Option<T>)> {
300 self.vector.iter_enumerated()
301 }
302
303 pub fn iter_indices(&self) -> impl Iterator<Item = I> + '_ {
304 self.iter_indexed().map(|(id, _)| id)
306 }
307
308 pub fn all_indices(&self) -> impl Iterator<Item = I> + use<I, T> {
309 self.vector.indices()
310 }
311
312 pub fn extract<'a, F: FnMut(&mut T) -> bool>(
315 &'a mut self,
316 mut f: F,
317 ) -> impl Iterator<Item = (I, T)> + use<'a, I, T, F> {
318 let elem_count = &mut self.elem_count;
319 self.vector
320 .iter_mut_enumerated()
321 .filter_map(move |(i, opt)| {
322 if f(opt.as_mut()?) {
323 *elem_count -= 1;
324 let elem = mem::replace(opt, None)?;
325 Some((i, elem))
326 } else {
327 None
328 }
329 })
330 }
331
332 pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
334 self.extract(|x| !f(x)).for_each(drop);
335 }
336
337 pub fn clear(&mut self) {
339 self.vector.clear();
340 self.elem_count = 0;
341 }
342 pub fn truncate(&mut self, at: usize) {
344 self.vector.truncate(at);
345 self.elem_count = self.iter().count();
346 }
347 pub fn split_off(&mut self, at: usize) -> Self {
349 let mut ret = Self {
350 vector: self.vector.split_off(I::from_usize(at)),
351 elem_count: 0,
352 };
353 self.elem_count = self.iter().count();
354 ret.elem_count = ret.iter().count();
355 ret
356 }
357}
358
359impl<I: Idx, T> Default for Vector<I, T> {
360 fn default() -> Self {
361 Self::new()
362 }
363}
364
365impl<I: Idx> Deref for ReservedSlot<I> {
366 type Target = I;
367 fn deref(&self) -> &Self::Target {
368 &self.0
369 }
370}
371
372impl<I, R, T> Index<R> for Vector<I, T>
373where
374 I: Idx,
375 R: IdxSliceIndex<I, Option<T>, Output = Option<T>>,
376{
377 type Output = T;
378 fn index(&self, index: R) -> &Self::Output {
379 self.vector[index].as_ref().unwrap()
380 }
381}
382
383impl<I, R, T> IndexMut<R> for Vector<I, T>
384where
385 I: Idx,
386 R: IdxSliceIndex<I, Option<T>, Output = Option<T>>,
387{
388 fn index_mut(&mut self, index: R) -> &mut Self::Output {
389 self.vector[index].as_mut().unwrap()
390 }
391}
392
393impl<'a, I, T> IntoIterator for &'a Vector<I, T>
394where
395 I: Idx,
396{
397 type Item = &'a T;
398 type IntoIter = impl Iterator<Item = &'a T>;
399
400 fn into_iter(self) -> Self::IntoIter {
401 self.vector.iter().flat_map(|opt| opt.as_ref())
402 }
403}
404
405impl<'a, I, T> IntoIterator for &'a mut Vector<I, T>
406where
407 I: Idx,
408{
409 type Item = &'a mut T;
410 type IntoIter = impl Iterator<Item = &'a mut T>;
411
412 fn into_iter(self) -> Self::IntoIter {
413 self.vector.iter_mut().flat_map(|opt| opt.as_mut())
414 }
415}
416
417impl<I, T> IntoIterator for Vector<I, T>
418where
419 I: Idx,
420{
421 type Item = T;
422 type IntoIter = impl Iterator<Item = T>;
423
424 fn into_iter(self) -> Self::IntoIter {
425 self.vector.into_iter().flat_map(|opt| opt)
426 }
427}
428
429impl<I, T> FromIterator<T> for Vector<I, T>
431where
432 I: Idx,
433{
434 #[inline]
435 fn from_iter<It: IntoIterator<Item = T>>(iter: It) -> Vector<I, T> {
436 let mut elem_count = 0;
437 let vector = IndexVec::from_iter(iter.into_iter().inspect(|_| elem_count += 1).map(Some));
438 Vector { vector, elem_count }
439 }
440}
441
442impl<I, T> From<Vec<T>> for Vector<I, T>
444where
445 I: Idx,
446{
447 fn from(v: Vec<T>) -> Self {
448 v.into_iter().collect()
449 }
450}
451
452impl<I, T, const N: usize> From<[T; N]> for Vector<I, T>
454where
455 I: Idx,
456{
457 fn from(v: [T; N]) -> Self {
458 v.into_iter().collect()
459 }
460}
461
462impl<I: Idx, T: Serialize> Serialize for Vector<I, T> {
463 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
464 where
465 S: Serializer,
466 {
467 self.vector.serialize(serializer)
468 }
469}
470
471impl<'de, I: Idx, T: Deserialize<'de>> Deserialize<'de> for Vector<I, T> {
472 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
473 where
474 D: serde::Deserializer<'de>,
475 {
476 let mut ret = Self {
477 vector: Deserialize::deserialize(deserializer)?,
478 elem_count: 0,
479 };
480 ret.elem_count = ret.iter().count();
481 Ok(ret)
482 }
483}
484
485impl<'s, I: Idx, T, V: Visit<'s, T>> Drive<'s, V> for Vector<I, T> {
486 fn drive_inner(&'s self, v: &mut V) -> ControlFlow<V::Break> {
487 for x in self {
488 v.visit(x)?;
489 }
490 Continue(())
491 }
492}
493impl<'s, I: Idx, T, V: VisitMut<'s, T>> DriveMut<'s, V> for Vector<I, T> {
494 fn drive_inner_mut(&'s mut self, v: &mut V) -> ControlFlow<V::Break> {
495 for x in self {
496 v.visit(x)?;
497 }
498 Continue(())
499 }
500}