1use super::indexes::{PageIndex, PageOffset, RowHash, RowPointer, SquashedOffset};
17use crate::{static_assert_size, MemoryUsage};
18use core::{hint, slice};
19use spacetimedb_data_structures::map::{
20 Entry,
21 IntMap, };
23
24#[derive(Clone, Copy, PartialEq, Eq, Debug)]
26struct ColliderSlotIndex(u32);
27
28impl MemoryUsage for ColliderSlotIndex {}
29
30impl ColliderSlotIndex {
31 fn new(idx: usize) -> Self {
33 Self(idx as u32)
34 }
35
36 fn idx(self) -> usize {
38 self.0 as usize
39 }
40}
41
42#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
46struct PtrOrCollider(RowPointer);
47
48impl MemoryUsage for PtrOrCollider {}
49
50enum MapSlotRef<'map> {
52 Pointer(&'map RowPointer),
54 Collider(ColliderSlotIndex),
57}
58
59enum MapSlotMut<'map> {
61 Pointer(&'map mut RowPointer),
63 Collider(ColliderSlotIndex),
66}
67
68#[inline]
72const fn ensure_ptr(rp: RowPointer) -> RowPointer {
73 rp.with_reserved_bit(false)
74}
75
76impl PtrOrCollider {
77 const fn ptr(rp: RowPointer) -> Self {
79 Self(ensure_ptr(rp))
80 }
81
82 const fn collider(c: ColliderSlotIndex) -> Self {
84 let pi = PageIndex(c.0 as u64);
86 Self(RowPointer::new(
87 true,
88 pi,
89 PageOffset::VAR_LEN_NULL,
90 SquashedOffset::COMMITTED_STATE,
91 ))
92 }
93
94 const fn is_ptr(&self) -> bool {
96 !self.0.reserved_bit()
97 }
98
99 const fn as_collider(&self) -> ColliderSlotIndex {
101 ColliderSlotIndex(self.0.page_index().0 as u32)
102 }
103
104 const fn unpack(&self) -> MapSlotRef<'_> {
106 if self.is_ptr() {
107 MapSlotRef::Pointer(&self.0)
108 } else {
109 MapSlotRef::Collider(self.as_collider())
110 }
111 }
112
113 fn unpack_mut(&mut self) -> MapSlotMut<'_> {
115 if self.is_ptr() {
116 MapSlotMut::Pointer(&mut self.0)
117 } else {
118 MapSlotMut::Collider(self.as_collider())
119 }
120 }
121}
122
123impl From<ColliderSlotIndex> for PtrOrCollider {
124 fn from(index: ColliderSlotIndex) -> Self {
125 Self::collider(index)
126 }
127}
128
129#[derive(Default, Clone, PartialEq, Eq, Debug)]
131pub struct PointerMap {
132 map: IntMap<RowHash, PtrOrCollider>,
136 colliders: Vec<Vec<RowPointer>>,
151 emptied_collider_slots: Vec<ColliderSlotIndex>,
154}
155
156impl MemoryUsage for PointerMap {
157 fn heap_usage(&self) -> usize {
158 let Self {
159 map,
160 colliders,
161 emptied_collider_slots,
162 } = self;
163 map.heap_usage() + colliders.heap_usage() + emptied_collider_slots.heap_usage()
164 }
165}
166
167static_assert_size!(PointerMap, 80);
168
169impl PointerMap {
171 pub fn num_collisions(&self) -> usize {
175 self.colliders.iter().map(|a| a.len()).sum()
176 }
177
178 pub fn num_non_collisions(&self) -> usize {
180 self.map.len() - (self.colliders.len() - self.emptied_collider_slots.len())
181 }
182
183 pub fn len(&self) -> usize {
186 self.num_collisions() + self.num_non_collisions()
187 }
188
189 pub fn is_empty(&self) -> bool {
191 self.map.is_empty()
192 }
193
194 pub fn pointers_for(&self, hash: RowHash) -> &[RowPointer] {
196 self.map.get(&hash).map_or(&[], |poc| match poc.unpack() {
197 MapSlotRef::Pointer(ro) => slice::from_ref(ro),
198 MapSlotRef::Collider(ci) => &self.colliders[ci.idx()],
199 })
200 }
201
202 pub fn pointers_for_mut(&mut self, hash: RowHash) -> &mut [RowPointer] {
207 self.map.get_mut(&hash).map_or(&mut [], |poc| match poc.unpack_mut() {
208 MapSlotMut::Pointer(ro) => slice::from_mut(ro),
209 MapSlotMut::Collider(ci) => &mut self.colliders[ci.idx()],
210 })
211 }
212
213 pub fn insert(&mut self, hash: RowHash, ptr: RowPointer) -> bool {
218 let mut was_in_map = false;
219
220 self.map
221 .entry(hash)
222 .and_modify(|v| match v.unpack() {
223 MapSlotRef::Pointer(existing) if *existing == ptr => was_in_map = true,
225 MapSlotRef::Pointer(existing) => {
227 let ptrs = [*existing, ptr].map(ensure_ptr);
228 let ci = match self.emptied_collider_slots.pop() {
229 None => {
231 let ci = ColliderSlotIndex::new(self.colliders.len());
232 self.colliders.push(ptrs.into());
233 ci
234 }
235 Some(ci) => {
237 self.colliders[ci.idx()].extend(ptrs);
238 ci
239 }
240 };
241 *v = PtrOrCollider::collider(ci);
242 }
243 MapSlotRef::Collider(ci) => {
245 let ptr = ensure_ptr(ptr);
246 let colliders = &mut self.colliders[ci.idx()];
247 if colliders.contains(&ptr) {
248 return was_in_map = true;
266 }
267 colliders.push(ptr);
268 }
269 })
270 .or_insert(PtrOrCollider::ptr(ptr));
272
273 was_in_map
274 }
275
276 pub fn remove(&mut self, hash: RowHash, ptr: RowPointer) -> bool {
280 let ret = 'fun: {
281 let Entry::Occupied(mut entry) = self.map.entry(hash) else {
282 break 'fun false;
283 };
284
285 match entry.get().unpack() {
286 MapSlotRef::Pointer(o) if *o == ptr => drop(entry.remove()),
288 MapSlotRef::Pointer(_) => break 'fun false,
289 MapSlotRef::Collider(ci) => {
290 let slot = &mut self.colliders[ci.idx()];
292 let Some(idx) = slot.iter().position(|o| *o == ptr) else {
293 break 'fun false;
294 };
295 slot.swap_remove(idx);
296
297 match slot.len() {
298 0 => unsafe { hint::unreachable_unchecked() },
300 1 => *entry.get_mut() = PtrOrCollider::ptr(slot.pop().unwrap()),
302 _ => break 'fun true,
303 }
304
305 self.emptied_collider_slots.push(ci);
307 }
308 }
309
310 true
311 };
312
313 ret
314 }
315}
316
317impl FromIterator<(RowHash, RowPointer)> for PointerMap {
318 fn from_iter<T: IntoIterator<Item = (RowHash, RowPointer)>>(iter: T) -> Self {
319 let mut map = PointerMap::default();
320 for (h, o) in iter {
321 let _ = map.insert(h, o);
322 }
323 map
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330 use core::hash::Hash;
331 use core::mem;
332 use itertools::Itertools;
333 use proptest::collection::vec;
334 use proptest::prelude::*;
335
336 type R = Result<(), TestCaseError>;
337
338 fn gen_row_pointer() -> impl Strategy<Value = RowPointer> {
339 (any::<PageOffset>(), any::<PageIndex>()).prop_map(|(po, pi)| RowPointer::new(false, pi, po, SquashedOffset(0)))
340 }
341
342 fn collect_entries(map: &PointerMap) -> Vec<(RowHash, PtrOrCollider)> {
343 map.map.iter().map(|(h, o)| (*h, *o)).collect::<Vec<_>>()
344 }
345
346 fn entry(hash: RowHash, ptr: RowPointer) -> (RowHash, PtrOrCollider) {
347 (hash, PtrOrCollider(ptr))
348 }
349
350 fn sorted<T: Ord + Copy>(xs: &[T]) -> Vec<T> {
351 xs.iter().copied().sorted().collect()
352 }
353
354 fn assert_ptrs_are(map: &mut PointerMap, hash: RowHash, ptrs: &[RowPointer]) -> R {
355 let ptrs = sorted(ptrs);
356 prop_assert_eq!(sorted(map.pointers_for(hash)), &*ptrs);
357 prop_assert_eq!(sorted(map.pointers_for_mut(hash)), ptrs);
358 Ok(())
359 }
360
361 fn assert_ptrs_and_len(map: &mut PointerMap, hash: RowHash, ptrs: &[RowPointer]) -> R {
362 assert_ptrs_are(map, hash, ptrs)?;
363 prop_assert_eq!(map.len(), ptrs.len());
364 prop_assert_eq!(map.is_empty(), ptrs.is_empty());
365 Ok(())
366 }
367
368 fn assert_collisions(map: &PointerMap, num_collisions: usize, num_not: usize) -> R {
369 prop_assert_eq!(map.num_collisions(), num_collisions);
370 prop_assert_eq!(map.num_non_collisions(), num_not);
371 Ok(())
372 }
373
374 fn ensure_unique<T: Eq + Hash>(xs: &[T]) -> R {
375 if !xs.iter().all_unique() {
376 return Err(TestCaseError::reject("all elements must be unique"));
377 }
378 Ok(())
379 }
380
381 proptest! {
382 #[test]
383 fn insert_same_twice_idempotence(
384 (hash, ptrs) in (
385 any::<RowHash>(),
386 vec(gen_row_pointer(), 3..10)
387 )
388 ) {
389 ensure_unique(&ptrs)?;
390
391 let mut map = PointerMap::default();
392
393 let ptr = ptrs[0];
395 prop_assert_eq!(map.insert(hash, ptr), false);
396 let old_map = map.clone(); prop_assert_eq!(map.insert(hash, ptr), true); prop_assert_eq!(&map, &old_map); assert_ptrs_and_len(&mut map, hash, &[ptr])?;
401 assert_collisions(&map, 0, 1)?;
402 prop_assert_eq!(collect_entries(&map), [entry(hash, ptr)]);
404 prop_assert!(map.colliders.is_empty());
405 prop_assert!(map.emptied_collider_slots.is_empty());
406
407 for ptr in &ptrs[1..] {
410 prop_assert_eq!(map.insert(hash, *ptr), false);
411 }
412 assert_ptrs_and_len(&mut map, hash, &ptrs)?;
413 assert_collisions(&map, ptrs.len(), 0)?;
414 let old_map = map.clone(); prop_assert_eq!(map.insert(hash, ptr), true); prop_assert_eq!(&map, &old_map); assert_ptrs_and_len(&mut map, hash, &ptrs)?;
420 assert_collisions(&map, ptrs.len(), 0)?;
421 prop_assert_eq!(collect_entries(&map), [(hash, ColliderSlotIndex::new(0).into())]);
423 prop_assert_eq!(map.colliders, [ptrs.to_owned()]);
424 prop_assert!(map.emptied_collider_slots.is_empty());
425 }
426
427 #[test]
428 fn insert_same_ptr_under_diff_hash(
429 (hashes, ptr) in (vec(any::<RowHash>(), 2..10), gen_row_pointer())
430 ) {
431 ensure_unique(&hashes)?;
432
433 let mut map = PointerMap::default();
435 for hash in &hashes {
436 prop_assert_eq!(map.insert(*hash, ptr), false);
437 }
438 for hash in &hashes {
440 assert_ptrs_are(&mut map, *hash, &[ptr])?;
441 }
442 prop_assert_eq!(map.len(), hashes.len());
443 prop_assert_eq!(map.is_empty(), false);
444 assert_collisions(&map, 0, hashes.len())?;
445 let mut entries = collect_entries(&map);
447 entries.sort();
448 prop_assert_eq!(
449 entries,
450 hashes.iter().copied().sorted().map(|hash| entry(hash, ptr)).collect::<Vec<_>>()
451 );
452 prop_assert!(map.colliders.is_empty());
453 prop_assert!(map.emptied_collider_slots.is_empty());
454 }
455
456 #[test]
457 fn insert_different_for_same_hash_handles_collision(
458 (hash, ptrs) in (any::<RowHash>(), vec(gen_row_pointer(), 3..10))
459 ) {
460 ensure_unique(&ptrs)?;
461
462 let mut map = PointerMap::default();
463
464 prop_assert_eq!(map.insert(hash, ptrs[0]), false);
466 assert_ptrs_and_len(&mut map, hash, &ptrs[..1])?;
468 assert_collisions(&map, 0, 1)?;
469 prop_assert_eq!(collect_entries(&map), [entry(hash, ptrs[0])]);
471 prop_assert!(map.colliders.is_empty());
472 prop_assert!(map.emptied_collider_slots.is_empty());
473
474 prop_assert_eq!(map.insert(hash, ptrs[1]), false);
477 assert_ptrs_and_len(&mut map, hash, &ptrs[..2])?;
479 assert_collisions(&map, 2, 0)?;
480 let first_collider_idx = ColliderSlotIndex::new(0);
482 let one_collider_entry = [(hash, first_collider_idx.into())];
483 prop_assert_eq!(collect_entries(&map), one_collider_entry);
484 prop_assert_eq!(&*map.colliders, [ptrs[..2].to_owned()]);
485 prop_assert!(map.emptied_collider_slots.is_empty());
486
487 for (ptr, i) in ptrs[2..].iter().copied().zip(2..) {
489 prop_assert_eq!(map.insert(hash, ptr), false);
491 assert_ptrs_and_len(&mut map, hash, &ptrs[..=i])?;
493 assert_collisions(&map, i + 1, 0)?;
494 prop_assert_eq!(collect_entries(&map), one_collider_entry);
496 prop_assert_eq!(&*map.colliders, [ptrs[..=i].to_owned()]);
497 prop_assert!(map.emptied_collider_slots.is_empty());
498 }
499
500 let last = ptrs.len() - 1;
502 for ptr in &ptrs[..last] {
503 prop_assert!(map.remove(hash, *ptr));
504 }
505 assert_ptrs_and_len(&mut map, hash, &ptrs[last..])?;
507 assert_collisions(&map, 0, 1)?;
508 prop_assert_eq!(collect_entries(&map), [entry(hash, ptrs[last])]);
510 prop_assert_eq!(&*map.colliders, [vec![]]);
511 prop_assert_eq!(&*map.emptied_collider_slots, [first_collider_idx]);
512
513 let penultimate = last - 1;
516 prop_assert_eq!(map.insert(hash, ptrs[penultimate]), false);
517 let pointers = ptrs[penultimate..].iter().copied().rev().collect::<Vec<_>>();
519 assert_ptrs_and_len(&mut map, hash, &pointers)?;
520 assert_collisions(&map, 2, 0)?;
521 prop_assert_eq!(collect_entries(&map), one_collider_entry);
523 prop_assert_eq!(&*map.colliders, [pointers]);
524 prop_assert!(map.emptied_collider_slots.is_empty());
525 }
526
527 #[test]
528 fn remove_non_existing_fails((hash, ptr) in (any::<RowHash>(), gen_row_pointer())) {
529 let mut map = PointerMap::default();
530 prop_assert_eq!(map.remove(hash, ptr), false);
531 }
532
533 #[test]
534 fn remove_uncollided_hash_works((hash, ptr) in (any::<RowHash>(), gen_row_pointer())) {
535 let mut map = PointerMap::default();
536
537 prop_assert_eq!(map.insert(hash, ptr), false);
539 prop_assert_eq!(map.remove(hash, ptr), true);
540 assert_ptrs_and_len(&mut map, hash, &[])?;
542 assert_collisions(&map, 0, 0)?;
543 prop_assert_eq!(collect_entries(&map), []);
545 prop_assert!(map.colliders.is_empty());
546 prop_assert!(map.emptied_collider_slots.is_empty());
547 }
548
549 #[test]
550 fn remove_same_hash_wrong_ptr_fails(
551 (hash, ptr_a, ptr_b) in (
552 any::<RowHash>(),
553 gen_row_pointer(),
554 gen_row_pointer(),
555 )
556 ) {
557 ensure_unique(&[ptr_a, ptr_b])?;
558
559 let mut map = PointerMap::default();
560
561 prop_assert_eq!(map.insert(hash, ptr_a), false);
563 prop_assert_eq!(map.remove(hash, ptr_b), false);
564 assert_ptrs_and_len(&mut map, hash, &[ptr_a])?;
566 assert_collisions(&map, 0, 1)?;
567 prop_assert_eq!(collect_entries(&map), [entry(hash, ptr_a)]);
569 prop_assert!(map.colliders.is_empty());
570 prop_assert!(map.emptied_collider_slots.is_empty());
571 }
572
573 #[test]
574 fn remove_collided_hash_wrong_ptr_fails(
575 (hash, ptrs) in (any::<RowHash>(), vec(gen_row_pointer(), 3..10))
576 ) {
577 ensure_unique(&ptrs)?;
578
579 let mut map = PointerMap::default();
580
581 let last = ptrs.len() - 1;
583 let but_last = &ptrs[0..last];
584 for ptr in but_last {
585 prop_assert_eq!(map.insert(hash, *ptr), false);
586 }
587 prop_assert_eq!(map.remove(hash, ptrs[last]), false);
588 assert_ptrs_and_len(&mut map, hash, but_last)?;
590 assert_collisions(&map, but_last.len(), 0)?;
591 prop_assert_eq!(collect_entries(&map), [(hash, ColliderSlotIndex::new(0).into())]);
593 prop_assert_eq!(&*map.colliders, [but_last.to_owned()]);
594 prop_assert!(map.emptied_collider_slots.is_empty());
595 }
596
597 #[test]
598 fn remove_collided_hash_reduction_works(
599 (hash, mut ptr_a, mut ptr_b, pick_b) in (
600 any::<RowHash>(),
601 gen_row_pointer(),
602 gen_row_pointer(),
603 any::<bool>(),
604 )
605 ) {
606 ensure_unique(&[ptr_a, ptr_b])?;
607
608 let mut map = PointerMap::default();
610 prop_assert_eq!(map.insert(hash, ptr_a), false);
611 prop_assert_eq!(map.insert(hash, ptr_b), false);
612 assert_collisions(&map, 2, 0)?;
613
614 if pick_b {
616 mem::swap(&mut ptr_a, &mut ptr_b);
617 }
618 prop_assert_eq!(map.remove(hash, ptr_b), true);
619 assert_ptrs_and_len(&mut map, hash, &[ptr_a])?;
620 assert_collisions(&map, 0, 1)?;
621 prop_assert_eq!(map.emptied_collider_slots, [ColliderSlotIndex(0)]);
622 }
623
624 #[test]
625 fn remove_collided_hash_works(
626 (hash, mut ptrs, pick_remove_idx) in (
627 any::<RowHash>(),
628 vec(gen_row_pointer(), 3..10),
629 0..10usize,
630 )
631 ) {
632 ensure_unique(&ptrs)?;
633
634 let pick_remove_idx = pick_remove_idx.min(ptrs.len() - 1);
635
636 let mut map = PointerMap::default();
638 for ptr in &ptrs {
639 prop_assert_eq!(map.insert(hash, *ptr), false);
640 }
641 assert_collisions(&map, ptrs.len(), 0)?;
642
643 let ptr_to_remove = ptrs.remove(pick_remove_idx);
645 prop_assert_eq!(map.remove(hash, ptr_to_remove), true);
646 assert_ptrs_and_len(&mut map, hash, &ptrs)?;
647 assert_collisions(&map, ptrs.len(), 0)?;
648 prop_assert_eq!(sorted(&map.colliders[0]), sorted(&ptrs));
649 prop_assert_eq!(map.emptied_collider_slots, []);
650 }
651 }
652}