1use std::num::NonZeroUsize;
11
12use crate::{DenseMap, EntryKey};
13
14pub trait DenseMapCollectionKey {
20 const VARIANT_COUNT: NonZeroUsize;
22
23 fn get_variant(&self) -> usize;
30
31 fn get_id(&self) -> usize;
33}
34
35impl<O> EntryKey for O
36where
37 O: DenseMapCollectionKey,
38{
39 fn get_key_index(&self) -> usize {
40 <O as DenseMapCollectionKey>::get_id(self)
41 }
42}
43
44pub struct VacantEntry<'a, K, T> {
46 entry: crate::VacantEntry<'a, K, T>,
47 count: &'a mut usize,
48}
49
50impl<'a, K, T> VacantEntry<'a, K, T> {
51 pub fn insert(self, value: T) -> &'a mut T
54 where
55 K: EntryKey,
56 {
57 let Self { entry, count } = self;
58 *count += 1;
59 entry.insert(value)
60 }
61
62 pub fn key(&self) -> &K {
65 self.entry.key()
66 }
67
68 pub fn into_key(self) -> K {
70 self.entry.into_key()
71 }
72
73 pub(crate) fn map_key<X, F>(self, f: F) -> VacantEntry<'a, X, T>
81 where
82 K: EntryKey,
83 X: EntryKey,
84 F: FnOnce(K) -> X,
85 {
86 let Self { entry, count } = self;
87 VacantEntry { entry: entry.map_key(f), count }
88 }
89}
90
91#[derive(Debug)]
93pub struct OccupiedEntry<'a, K, T> {
94 entry: crate::OccupiedEntry<'a, K, T>,
95 count: &'a mut usize,
96}
97
98impl<'a, K: EntryKey, T> OccupiedEntry<'a, K, T> {
99 pub fn key(&self) -> &K {
101 self.entry.key()
102 }
103
104 pub fn into_key(self) -> K {
106 self.entry.into_key()
107 }
108
109 pub fn get(&self) -> &T {
111 self.entry.get()
112 }
113
114 pub fn get_mut(&mut self) -> &mut T {
119 self.entry.get_mut()
120 }
121
122 pub fn into_mut(self) -> &'a mut T {
128 self.entry.into_mut()
129 }
130
131 pub fn insert(&mut self, value: T) -> T {
133 self.entry.insert(value)
134 }
135
136 pub fn remove(self) -> T {
138 let Self { entry, count } = self;
139 *count -= 1;
140 entry.remove()
141 }
142
143 pub(crate) fn map_key<X, F>(self, f: F) -> OccupiedEntry<'a, X, T>
151 where
152 K: EntryKey,
153 X: EntryKey,
154 F: FnOnce(K) -> X,
155 {
156 let Self { entry, count } = self;
157 OccupiedEntry { entry: entry.map_key(f), count }
158 }
159}
160
161pub enum Entry<'a, K, T> {
163 Vacant(VacantEntry<'a, K, T>),
165 Occupied(OccupiedEntry<'a, K, T>),
167}
168
169impl<'a, K: EntryKey, T> Entry<'a, K, T> {
170 pub fn key(&self) -> &K {
172 match self {
173 Entry::Occupied(e) => e.key(),
174 Entry::Vacant(e) => e.key(),
175 }
176 }
177
178 pub fn or_insert(self, default: T) -> &'a mut T
181 where
182 K: EntryKey,
183 {
184 self.or_insert_with(|| default)
185 }
186
187 pub fn or_insert_with<F: FnOnce() -> T>(self, f: F) -> &'a mut T {
190 match self {
191 Entry::Occupied(e) => e.into_mut(),
192 Entry::Vacant(e) => e.insert(f()),
193 }
194 }
195
196 pub fn or_default(self) -> &'a mut T
199 where
200 T: Default,
201 K: EntryKey,
202 {
203 self.or_insert_with(<T as Default>::default)
204 }
205
206 pub fn and_modify<F: FnOnce(&mut T)>(self, f: F) -> Self {
209 match self {
210 Entry::Occupied(mut e) => {
211 f(e.get_mut());
212 Entry::Occupied(e)
213 }
214 Entry::Vacant(e) => Entry::Vacant(e),
215 }
216 }
217
218 pub fn remove(self) -> Option<T> {
222 match self {
223 Entry::Vacant(_) => None,
224 Entry::Occupied(e) => Some(e.remove()),
225 }
226 }
227}
228
229struct SizeAugmentedIterator<I> {
234 wrapped: I,
235 remaining: usize,
236}
237
238impl<I: Iterator> Iterator for SizeAugmentedIterator<I> {
239 type Item = I::Item;
240
241 fn next(&mut self) -> Option<Self::Item> {
242 let Self { wrapped, remaining } = self;
243 match wrapped.next() {
244 Some(v) => {
245 *remaining -= 1;
246 Some(v)
247 }
248 None => {
249 assert_eq!(remaining, &0);
250 None
251 }
252 }
253 }
254
255 fn size_hint(&self) -> (usize, Option<usize>) {
256 (self.remaining, Some(self.remaining))
257 }
258}
259
260impl<I: Iterator> ExactSizeIterator for SizeAugmentedIterator<I> {}
261
262pub struct DenseMapCollection<K: DenseMapCollectionKey, T> {
268 data: Vec<DenseMap<T>>,
273 count: usize,
274 _marker: core::marker::PhantomData<K>,
275}
276
277impl<K: DenseMapCollectionKey, T> DenseMapCollection<K, T> {
278 pub fn new() -> Self {
280 let mut data = Vec::new();
281 data.resize_with(K::VARIANT_COUNT.get(), DenseMap::default);
282 Self { data, count: 0, _marker: core::marker::PhantomData }
283 }
284
285 fn get_map(&self, key: &K) -> &DenseMap<T> {
286 &self.data[key.get_variant()]
287 }
288
289 fn get_entry(&mut self, key: &K) -> Entry<'_, usize, T> {
290 let Self { data, count, _marker } = self;
291 match data[key.get_variant()].entry(key.get_id()) {
292 crate::Entry::Occupied(entry) => Entry::Occupied(OccupiedEntry { entry, count }),
293 crate::Entry::Vacant(entry) => Entry::Vacant(VacantEntry { entry, count }),
294 }
295 }
296
297 pub fn is_empty(&self) -> bool {
299 let Self { count, data: _, _marker } = self;
300 *count == 0
301 }
302
303 pub fn get(&self, key: &K) -> Option<&T> {
306 self.get_map(key).get(key.get_id())
307 }
308
309 pub fn get_mut(&mut self, key: &K) -> Option<&mut T> {
312 match self.get_entry(key) {
313 Entry::Occupied(e) => Some(e.into_mut()),
314 Entry::Vacant(_) => None,
315 }
316 }
317
318 pub fn remove(&mut self, key: &K) -> Option<T> {
322 match self.get_entry(key) {
323 Entry::Occupied(e) => Some(e.remove()),
324 Entry::Vacant(_) => None,
325 }
326 }
327
328 pub fn insert(&mut self, key: &K, item: T) -> Option<T> {
333 match self.get_entry(key) {
334 Entry::Occupied(mut e) => Some(e.insert(item)),
335 Entry::Vacant(e) => {
336 let _: &mut T = e.insert(item);
337 None
338 }
339 }
340 }
341
342 pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
344 let Self { data, count, _marker } = self;
345 SizeAugmentedIterator {
346 wrapped: data.iter().flat_map(|m| m.key_ordered_iter()).map(|(_, v)| v),
347 remaining: *count,
348 }
349 }
350
351 pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
353 let Self { data, count, _marker } = self;
354 SizeAugmentedIterator {
355 wrapped: data.iter_mut().flat_map(|m| m.key_ordered_iter_mut()).map(|(_, v)| v),
356 remaining: *count,
357 }
358 }
359
360 pub fn iter_maps(&self) -> impl Iterator<Item = &DenseMap<T>> {
362 let Self { data, count: _, _marker } = self;
363 data.iter()
364 }
365
366 pub fn entry(&mut self, key: K) -> Entry<'_, K, T> {
369 match self.get_entry(&key) {
370 Entry::Occupied(e) => Entry::Occupied(e.map_key(|_| key)),
371 Entry::Vacant(e) => Entry::Vacant(e.map_key(|_| key)),
372 }
373 }
374
375 pub fn push_entry(&mut self, make_key: fn(usize) -> K, value: T) -> OccupiedEntry<'_, K, T> {
382 let Self { count, data, _marker } = self;
383 let variant = make_key(0).get_variant();
384 let entry = data[variant].push_entry(value);
385 *count += 1;
386
387 let entry = entry.map_key(make_key);
388
389 let entry_variant = entry.key().get_variant();
390 assert_eq!(
391 entry_variant, variant,
392 "key variant is inconsistent; got both {variant} and {entry_variant}"
393 );
394 OccupiedEntry { entry, count }
395 }
396}
397
398impl<K: DenseMapCollectionKey, T> Default for DenseMapCollection<K, T> {
399 fn default() -> Self {
400 Self::new()
401 }
402}
403
404#[cfg(test)]
405mod tests {
406 use std::collections::HashSet;
407
408 use super::*;
409 use crate::testutil::assert_empty;
410
411 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
412 enum FakeVariants {
413 A,
414 B,
415 C,
416 }
417
418 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
419 struct FakeKey {
420 id: usize,
421 var: FakeVariants,
422 }
423
424 impl FakeKey {
425 const fn new(id: usize, var: FakeVariants) -> Self {
426 Self { id, var }
427 }
428 }
429
430 impl DenseMapCollectionKey for FakeKey {
431 const VARIANT_COUNT: NonZeroUsize = NonZeroUsize::new(3).unwrap();
432
433 fn get_variant(&self) -> usize {
434 match self.var {
435 FakeVariants::A => 0,
436 FakeVariants::B => 1,
437 FakeVariants::C => 2,
438 }
439 }
440
441 fn get_id(&self) -> usize {
442 self.id
443 }
444 }
445
446 type TestCollection = DenseMapCollection<FakeKey, i32>;
447
448 const KEY_A: FakeKey = FakeKey::new(0, FakeVariants::A);
449 const KEY_B: FakeKey = FakeKey::new(2, FakeVariants::B);
450 const KEY_C: FakeKey = FakeKey::new(4, FakeVariants::C);
451
452 #[test]
453 fn test_insert_and_get() {
454 let mut t = TestCollection::new();
455 let DenseMapCollection { data, count, _marker } = &t;
456 assert_empty(data[0].key_ordered_iter());
457 assert_empty(data[1].key_ordered_iter());
458 assert_empty(data[2].key_ordered_iter());
459 assert_eq!(count, &0);
460
461 assert_eq!(t.insert(&KEY_A, 1), None);
462 let DenseMapCollection { data, count, _marker } = &t;
463 assert!(!data[0].is_empty());
464 assert_eq!(count, &1);
465
466 assert_eq!(t.insert(&KEY_B, 2), None);
467 let DenseMapCollection { data, count, _marker } = &t;
468 assert!(!data[1].is_empty());
469 assert_eq!(count, &2);
470
471 assert_eq!(*t.get(&KEY_A).unwrap(), 1);
472 assert_eq!(t.get(&KEY_C), None);
473
474 *t.get_mut(&KEY_B).unwrap() = 3;
475 assert_eq!(*t.get(&KEY_B).unwrap(), 3);
476 }
477
478 #[test]
479 fn test_remove() {
480 let mut t = TestCollection::new();
481 assert_eq!(t.insert(&KEY_B, 15), None);
482 assert_eq!(t.remove(&KEY_B).unwrap(), 15);
483 let DenseMapCollection { data: _, count, _marker } = &t;
484 assert_eq!(count, &0);
485
486 assert_eq!(t.remove(&KEY_B), None);
487 }
488
489 #[test]
490 fn test_iter() {
491 let mut t = TestCollection::new();
492 assert_eq!(t.insert(&KEY_A, 15), None);
493 assert_eq!(t.insert(&KEY_B, -5), None);
494 assert_eq!(t.insert(&KEY_C, -10), None);
495 let mut c = 0;
496 let mut sum = 0;
497 for i in t.iter() {
498 c += 1;
499 sum += *i;
500 }
501 assert_eq!(c, 3);
502 assert_eq!(sum, 0);
503 }
504
505 #[test]
506 fn test_iter_len() {
507 let mut t = TestCollection::new();
508 assert_eq!(t.insert(&KEY_A, 1), None);
509 assert_eq!(t.insert(&KEY_B, 1), None);
510 assert_eq!(t.insert(&KEY_C, 1), None);
511 assert_eq!(t.iter().len(), 3);
512 assert_eq!(t.remove(&KEY_A), Some(1));
513 assert_eq!(t.iter().len(), 2);
514 }
515
516 #[test]
517 fn test_is_empty() {
518 let mut t = TestCollection::new();
519 assert!(t.is_empty());
520 assert_eq!(t.insert(&KEY_B, 15), None);
521 assert!(!t.is_empty());
522 }
523
524 #[test]
525 fn test_iter_mut() {
526 let mut t = TestCollection::new();
527 assert_eq!(t.insert(&KEY_A, 15), None);
528 assert_eq!(t.insert(&KEY_B, -5), None);
529 assert_eq!(t.insert(&KEY_C, -10), None);
530 for i in t.iter_mut() {
531 *i *= 2;
532 }
533 assert_eq!(*t.get(&KEY_A).unwrap(), 30);
534 assert_eq!(*t.get(&KEY_B).unwrap(), -10);
535 assert_eq!(*t.get(&KEY_C).unwrap(), -20);
536 assert_eq!(t.iter_mut().len(), 3);
537 }
538
539 #[test]
540 fn test_entry() {
541 let mut t = TestCollection::new();
542 assert_eq!(*t.entry(KEY_A).or_insert(2), 2);
543 assert_eq!(
544 *t.entry(KEY_A)
545 .and_modify(|v| {
546 *v = 10;
547 })
548 .or_insert(5),
549 10
550 );
551 assert_eq!(
552 *t.entry(KEY_B)
553 .and_modify(|v| {
554 *v = 10;
555 })
556 .or_insert(5),
557 5
558 );
559 assert_eq!(*t.entry(KEY_C).or_insert_with(|| 7), 7);
560
561 assert_eq!(*t.entry(KEY_C).key(), KEY_C);
562 assert_eq!(*t.get(&KEY_A).unwrap(), 10);
563 assert_eq!(*t.get(&KEY_B).unwrap(), 5);
564 assert_eq!(*t.get(&KEY_C).unwrap(), 7);
565 }
566
567 #[test]
568 fn push_entry_valid() {
569 let mut t = TestCollection::new();
570 assert_eq!(t.insert(&KEY_A, 0), None);
571 assert_eq!(t.insert(&KEY_B, 1), None);
572 assert_eq!(t.insert(&KEY_C, 2), None);
573
574 let make_key = |index| FakeKey { id: index, var: FakeVariants::A };
575
576 {
577 let entry = t.push_entry(make_key, 30);
578 assert_eq!(entry.key(), &FakeKey { id: 1, var: FakeVariants::A });
579 assert_eq!(entry.get(), &30);
580 }
581
582 {
583 let entry = t.push_entry(make_key, 20);
584 assert_eq!(entry.key(), &FakeKey { id: 2, var: FakeVariants::A });
585 assert_eq!(entry.get(), &20);
586 }
587
588 assert_eq!(t.iter().collect::<HashSet<_>>(), HashSet::from([&0, &1, &2, &30, &20]));
589 }
590
591 #[test]
592 #[should_panic(expected = "variant is inconsistent")]
593 fn push_entry_invalid_key_fn() {
594 let mut t = TestCollection::new();
595 assert_eq!(t.insert(&KEY_A, 0), None);
596
597 let bad_make_key = |index| FakeKey {
598 id: index,
599 var: if index % 2 == 0 { FakeVariants::A } else { FakeVariants::B },
600 };
601
602 let _ = t.push_entry(bad_make_key, 1);
603 let _ = t.push_entry(bad_make_key, 2);
604 }
605}