1use super::indexes::{PageIndex, PageOffset, RowHash, RowPointer, SquashedOffset};
17use crate::static_assert_size;
18use core::{hint, slice};
19use spacetimedb_data_structures::map::{
20 Entry,
21 IntMap, };
23use spacetimedb_sats::memory_usage::MemoryUsage;
24
25#[derive(Clone, Copy, PartialEq, Eq, Debug)]
27struct ColliderSlotIndex(u32);
28
29impl MemoryUsage for ColliderSlotIndex {}
30
31impl ColliderSlotIndex {
32 fn new(idx: usize) -> Self {
34 Self(idx as u32)
35 }
36
37 fn idx(self) -> usize {
39 self.0 as usize
40 }
41}
42
43#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
47struct PtrOrCollider(RowPointer);
48
49impl MemoryUsage for PtrOrCollider {}
50
51enum MapSlotRef<'map> {
53 Pointer(&'map RowPointer),
55 Collider(ColliderSlotIndex),
58}
59
60enum MapSlotMut<'map> {
62 Pointer(&'map mut RowPointer),
64 Collider(ColliderSlotIndex),
67}
68
69#[inline]
73const fn ensure_ptr(rp: RowPointer) -> RowPointer {
74 rp.with_reserved_bit(false)
75}
76
77impl PtrOrCollider {
78 const fn ptr(rp: RowPointer) -> Self {
80 Self(ensure_ptr(rp))
81 }
82
83 const fn collider(c: ColliderSlotIndex) -> Self {
85 let pi = PageIndex(c.0 as u64);
87 Self(RowPointer::new(
88 true,
89 pi,
90 PageOffset::VAR_LEN_NULL,
91 SquashedOffset::COMMITTED_STATE,
92 ))
93 }
94
95 const fn is_ptr(&self) -> bool {
97 !self.0.reserved_bit()
98 }
99
100 const fn as_collider(&self) -> ColliderSlotIndex {
102 ColliderSlotIndex(self.0.page_index().0 as u32)
103 }
104
105 const fn unpack(&self) -> MapSlotRef<'_> {
107 if self.is_ptr() {
108 MapSlotRef::Pointer(&self.0)
109 } else {
110 MapSlotRef::Collider(self.as_collider())
111 }
112 }
113
114 fn unpack_mut(&mut self) -> MapSlotMut<'_> {
116 if self.is_ptr() {
117 MapSlotMut::Pointer(&mut self.0)
118 } else {
119 MapSlotMut::Collider(self.as_collider())
120 }
121 }
122}
123
124impl From<ColliderSlotIndex> for PtrOrCollider {
125 fn from(index: ColliderSlotIndex) -> Self {
126 Self::collider(index)
127 }
128}
129
130#[derive(Default, Clone, PartialEq, Eq, Debug)]
132pub struct PointerMap {
133 map: IntMap<RowHash, PtrOrCollider>,
137 colliders: Vec<Vec<RowPointer>>,
152 emptied_collider_slots: Vec<ColliderSlotIndex>,
155}
156
157impl MemoryUsage for PointerMap {
158 fn heap_usage(&self) -> usize {
159 let Self {
160 map,
161 colliders,
162 emptied_collider_slots,
163 } = self;
164 map.heap_usage() + colliders.heap_usage() + emptied_collider_slots.heap_usage()
165 }
166}
167
168static_assert_size!(PointerMap, 80);
169
170impl PointerMap {
172 pub fn num_collisions(&self) -> usize {
176 self.colliders.iter().map(|a| a.len()).sum()
177 }
178
179 pub fn num_non_collisions(&self) -> usize {
181 self.map.len() - (self.colliders.len() - self.emptied_collider_slots.len())
182 }
183
184 pub fn len(&self) -> usize {
187 self.num_collisions() + self.num_non_collisions()
188 }
189
190 pub fn is_empty(&self) -> bool {
192 self.map.is_empty()
193 }
194
195 pub fn pointers_for(&self, hash: RowHash) -> &[RowPointer] {
197 self.map.get(&hash).map_or(&[], |poc| match poc.unpack() {
198 MapSlotRef::Pointer(ro) => slice::from_ref(ro),
199 MapSlotRef::Collider(ci) => &self.colliders[ci.idx()],
200 })
201 }
202
203 pub fn pointers_for_mut(&mut self, hash: RowHash) -> &mut [RowPointer] {
208 self.map.get_mut(&hash).map_or(&mut [], |poc| match poc.unpack_mut() {
209 MapSlotMut::Pointer(ro) => slice::from_mut(ro),
210 MapSlotMut::Collider(ci) => &mut self.colliders[ci.idx()],
211 })
212 }
213
214 pub fn insert(&mut self, hash: RowHash, ptr: RowPointer) -> bool {
219 let mut was_in_map = false;
220
221 self.map
222 .entry(hash)
223 .and_modify(|v| match v.unpack() {
224 MapSlotRef::Pointer(existing) if *existing == ptr => was_in_map = true,
226 MapSlotRef::Pointer(existing) => {
228 let ptrs = [*existing, ptr].map(ensure_ptr);
229 let ci = match self.emptied_collider_slots.pop() {
230 None => {
232 let ci = ColliderSlotIndex::new(self.colliders.len());
233 self.colliders.push(ptrs.into());
234 ci
235 }
236 Some(ci) => {
238 self.colliders[ci.idx()].extend(ptrs);
239 ci
240 }
241 };
242 *v = PtrOrCollider::collider(ci);
243 }
244 MapSlotRef::Collider(ci) => {
246 let ptr = ensure_ptr(ptr);
247 let colliders = &mut self.colliders[ci.idx()];
248 if colliders.contains(&ptr) {
249 return was_in_map = true;
267 }
268 colliders.push(ptr);
269 }
270 })
271 .or_insert(PtrOrCollider::ptr(ptr));
273
274 was_in_map
275 }
276
277 pub fn remove(&mut self, hash: RowHash, ptr: RowPointer) -> bool {
281 let ret = 'fun: {
282 let Entry::Occupied(mut entry) = self.map.entry(hash) else {
283 break 'fun false;
284 };
285
286 match entry.get().unpack() {
287 MapSlotRef::Pointer(o) if *o == ptr => drop(entry.remove()),
289 MapSlotRef::Pointer(_) => break 'fun false,
290 MapSlotRef::Collider(ci) => {
291 let slot = &mut self.colliders[ci.idx()];
293 let Some(idx) = slot.iter().position(|o| *o == ptr) else {
294 break 'fun false;
295 };
296 slot.swap_remove(idx);
297
298 match slot.len() {
299 0 => unsafe { hint::unreachable_unchecked() },
301 1 => *entry.get_mut() = PtrOrCollider::ptr(slot.pop().unwrap()),
303 _ => break 'fun true,
304 }
305
306 self.emptied_collider_slots.push(ci);
308 }
309 }
310
311 true
312 };
313
314 ret
315 }
316}
317
318impl FromIterator<(RowHash, RowPointer)> for PointerMap {
319 fn from_iter<T: IntoIterator<Item = (RowHash, RowPointer)>>(iter: T) -> Self {
320 let mut map = PointerMap::default();
321 for (h, o) in iter {
322 let _ = map.insert(h, o);
323 }
324 map
325 }
326}
327
328#[cfg(test)]
329mod tests {
330 use super::*;
331 use core::hash::Hash;
332 use core::mem;
333 use itertools::Itertools;
334 use proptest::collection::vec;
335 use proptest::prelude::*;
336
337 type R = Result<(), TestCaseError>;
338
339 fn gen_row_pointer() -> impl Strategy<Value = RowPointer> {
340 (any::<PageOffset>(), any::<PageIndex>()).prop_map(|(po, pi)| RowPointer::new(false, pi, po, SquashedOffset(0)))
341 }
342
343 fn collect_entries(map: &PointerMap) -> Vec<(RowHash, PtrOrCollider)> {
344 map.map.iter().map(|(h, o)| (*h, *o)).collect::<Vec<_>>()
345 }
346
347 fn entry(hash: RowHash, ptr: RowPointer) -> (RowHash, PtrOrCollider) {
348 (hash, PtrOrCollider(ptr))
349 }
350
351 fn sorted<T: Ord + Copy>(xs: &[T]) -> Vec<T> {
352 xs.iter().copied().sorted().collect()
353 }
354
355 fn assert_ptrs_are(map: &mut PointerMap, hash: RowHash, ptrs: &[RowPointer]) -> R {
356 let ptrs = sorted(ptrs);
357 prop_assert_eq!(sorted(map.pointers_for(hash)), &*ptrs);
358 prop_assert_eq!(sorted(map.pointers_for_mut(hash)), ptrs);
359 Ok(())
360 }
361
362 fn assert_ptrs_and_len(map: &mut PointerMap, hash: RowHash, ptrs: &[RowPointer]) -> R {
363 assert_ptrs_are(map, hash, ptrs)?;
364 prop_assert_eq!(map.len(), ptrs.len());
365 prop_assert_eq!(map.is_empty(), ptrs.is_empty());
366 Ok(())
367 }
368
369 fn assert_collisions(map: &PointerMap, num_collisions: usize, num_not: usize) -> R {
370 prop_assert_eq!(map.num_collisions(), num_collisions);
371 prop_assert_eq!(map.num_non_collisions(), num_not);
372 Ok(())
373 }
374
375 fn ensure_unique<T: Eq + Hash>(xs: &[T]) -> R {
376 if !xs.iter().all_unique() {
377 return Err(TestCaseError::reject("all elements must be unique"));
378 }
379 Ok(())
380 }
381
382 proptest! {
383 #[test]
384 fn insert_same_twice_idempotence(
385 (hash, ptrs) in (
386 any::<RowHash>(),
387 vec(gen_row_pointer(), 3..10)
388 )
389 ) {
390 ensure_unique(&ptrs)?;
391
392 let mut map = PointerMap::default();
393
394 let ptr = ptrs[0];
396 prop_assert_eq!(map.insert(hash, ptr), false);
397 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])?;
402 assert_collisions(&map, 0, 1)?;
403 prop_assert_eq!(collect_entries(&map), [entry(hash, ptr)]);
405 prop_assert!(map.colliders.is_empty());
406 prop_assert!(map.emptied_collider_slots.is_empty());
407
408 for ptr in &ptrs[1..] {
411 prop_assert_eq!(map.insert(hash, *ptr), false);
412 }
413 assert_ptrs_and_len(&mut map, hash, &ptrs)?;
414 assert_collisions(&map, ptrs.len(), 0)?;
415 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)?;
421 assert_collisions(&map, ptrs.len(), 0)?;
422 prop_assert_eq!(collect_entries(&map), [(hash, ColliderSlotIndex::new(0).into())]);
424 prop_assert_eq!(map.colliders, [ptrs.to_owned()]);
425 prop_assert!(map.emptied_collider_slots.is_empty());
426 }
427
428 #[test]
429 fn insert_same_ptr_under_diff_hash(
430 (hashes, ptr) in (vec(any::<RowHash>(), 2..10), gen_row_pointer())
431 ) {
432 ensure_unique(&hashes)?;
433
434 let mut map = PointerMap::default();
436 for hash in &hashes {
437 prop_assert_eq!(map.insert(*hash, ptr), false);
438 }
439 for hash in &hashes {
441 assert_ptrs_are(&mut map, *hash, &[ptr])?;
442 }
443 prop_assert_eq!(map.len(), hashes.len());
444 prop_assert_eq!(map.is_empty(), false);
445 assert_collisions(&map, 0, hashes.len())?;
446 let mut entries = collect_entries(&map);
448 entries.sort();
449 prop_assert_eq!(
450 entries,
451 hashes.iter().copied().sorted().map(|hash| entry(hash, ptr)).collect::<Vec<_>>()
452 );
453 prop_assert!(map.colliders.is_empty());
454 prop_assert!(map.emptied_collider_slots.is_empty());
455 }
456
457 #[test]
458 fn insert_different_for_same_hash_handles_collision(
459 (hash, ptrs) in (any::<RowHash>(), vec(gen_row_pointer(), 3..10))
460 ) {
461 ensure_unique(&ptrs)?;
462
463 let mut map = PointerMap::default();
464
465 prop_assert_eq!(map.insert(hash, ptrs[0]), false);
467 assert_ptrs_and_len(&mut map, hash, &ptrs[..1])?;
469 assert_collisions(&map, 0, 1)?;
470 prop_assert_eq!(collect_entries(&map), [entry(hash, ptrs[0])]);
472 prop_assert!(map.colliders.is_empty());
473 prop_assert!(map.emptied_collider_slots.is_empty());
474
475 prop_assert_eq!(map.insert(hash, ptrs[1]), false);
478 assert_ptrs_and_len(&mut map, hash, &ptrs[..2])?;
480 assert_collisions(&map, 2, 0)?;
481 let first_collider_idx = ColliderSlotIndex::new(0);
483 let one_collider_entry = [(hash, first_collider_idx.into())];
484 prop_assert_eq!(collect_entries(&map), one_collider_entry);
485 prop_assert_eq!(&*map.colliders, [ptrs[..2].to_owned()]);
486 prop_assert!(map.emptied_collider_slots.is_empty());
487
488 for (ptr, i) in ptrs[2..].iter().copied().zip(2..) {
490 prop_assert_eq!(map.insert(hash, ptr), false);
492 assert_ptrs_and_len(&mut map, hash, &ptrs[..=i])?;
494 assert_collisions(&map, i + 1, 0)?;
495 prop_assert_eq!(collect_entries(&map), one_collider_entry);
497 prop_assert_eq!(&*map.colliders, [ptrs[..=i].to_owned()]);
498 prop_assert!(map.emptied_collider_slots.is_empty());
499 }
500
501 let last = ptrs.len() - 1;
503 for ptr in &ptrs[..last] {
504 prop_assert!(map.remove(hash, *ptr));
505 }
506 assert_ptrs_and_len(&mut map, hash, &ptrs[last..])?;
508 assert_collisions(&map, 0, 1)?;
509 prop_assert_eq!(collect_entries(&map), [entry(hash, ptrs[last])]);
511 prop_assert_eq!(&*map.colliders, [vec![]]);
512 prop_assert_eq!(&*map.emptied_collider_slots, [first_collider_idx]);
513
514 let penultimate = last - 1;
517 prop_assert_eq!(map.insert(hash, ptrs[penultimate]), false);
518 let pointers = ptrs[penultimate..].iter().copied().rev().collect::<Vec<_>>();
520 assert_ptrs_and_len(&mut map, hash, &pointers)?;
521 assert_collisions(&map, 2, 0)?;
522 prop_assert_eq!(collect_entries(&map), one_collider_entry);
524 prop_assert_eq!(&*map.colliders, [pointers]);
525 prop_assert!(map.emptied_collider_slots.is_empty());
526 }
527
528 #[test]
529 fn remove_non_existing_fails((hash, ptr) in (any::<RowHash>(), gen_row_pointer())) {
530 let mut map = PointerMap::default();
531 prop_assert_eq!(map.remove(hash, ptr), false);
532 }
533
534 #[test]
535 fn remove_uncollided_hash_works((hash, ptr) in (any::<RowHash>(), gen_row_pointer())) {
536 let mut map = PointerMap::default();
537
538 prop_assert_eq!(map.insert(hash, ptr), false);
540 prop_assert_eq!(map.remove(hash, ptr), true);
541 assert_ptrs_and_len(&mut map, hash, &[])?;
543 assert_collisions(&map, 0, 0)?;
544 prop_assert_eq!(collect_entries(&map), []);
546 prop_assert!(map.colliders.is_empty());
547 prop_assert!(map.emptied_collider_slots.is_empty());
548 }
549
550 #[test]
551 fn remove_same_hash_wrong_ptr_fails(
552 (hash, ptr_a, ptr_b) in (
553 any::<RowHash>(),
554 gen_row_pointer(),
555 gen_row_pointer(),
556 )
557 ) {
558 ensure_unique(&[ptr_a, ptr_b])?;
559
560 let mut map = PointerMap::default();
561
562 prop_assert_eq!(map.insert(hash, ptr_a), false);
564 prop_assert_eq!(map.remove(hash, ptr_b), false);
565 assert_ptrs_and_len(&mut map, hash, &[ptr_a])?;
567 assert_collisions(&map, 0, 1)?;
568 prop_assert_eq!(collect_entries(&map), [entry(hash, ptr_a)]);
570 prop_assert!(map.colliders.is_empty());
571 prop_assert!(map.emptied_collider_slots.is_empty());
572 }
573
574 #[test]
575 fn remove_collided_hash_wrong_ptr_fails(
576 (hash, ptrs) in (any::<RowHash>(), vec(gen_row_pointer(), 3..10))
577 ) {
578 ensure_unique(&ptrs)?;
579
580 let mut map = PointerMap::default();
581
582 let last = ptrs.len() - 1;
584 let but_last = &ptrs[0..last];
585 for ptr in but_last {
586 prop_assert_eq!(map.insert(hash, *ptr), false);
587 }
588 prop_assert_eq!(map.remove(hash, ptrs[last]), false);
589 assert_ptrs_and_len(&mut map, hash, but_last)?;
591 assert_collisions(&map, but_last.len(), 0)?;
592 prop_assert_eq!(collect_entries(&map), [(hash, ColliderSlotIndex::new(0).into())]);
594 prop_assert_eq!(&*map.colliders, [but_last.to_owned()]);
595 prop_assert!(map.emptied_collider_slots.is_empty());
596 }
597
598 #[test]
599 fn remove_collided_hash_reduction_works(
600 (hash, mut ptr_a, mut ptr_b, pick_b) in (
601 any::<RowHash>(),
602 gen_row_pointer(),
603 gen_row_pointer(),
604 any::<bool>(),
605 )
606 ) {
607 ensure_unique(&[ptr_a, ptr_b])?;
608
609 let mut map = PointerMap::default();
611 prop_assert_eq!(map.insert(hash, ptr_a), false);
612 prop_assert_eq!(map.insert(hash, ptr_b), false);
613 assert_collisions(&map, 2, 0)?;
614
615 if pick_b {
617 mem::swap(&mut ptr_a, &mut ptr_b);
618 }
619 prop_assert_eq!(map.remove(hash, ptr_b), true);
620 assert_ptrs_and_len(&mut map, hash, &[ptr_a])?;
621 assert_collisions(&map, 0, 1)?;
622 prop_assert_eq!(map.emptied_collider_slots, [ColliderSlotIndex(0)]);
623 }
624
625 #[test]
626 fn remove_collided_hash_works(
627 (hash, mut ptrs, pick_remove_idx) in (
628 any::<RowHash>(),
629 vec(gen_row_pointer(), 3..10),
630 0..10usize,
631 )
632 ) {
633 ensure_unique(&ptrs)?;
634
635 let pick_remove_idx = pick_remove_idx.min(ptrs.len() - 1);
636
637 let mut map = PointerMap::default();
639 for ptr in &ptrs {
640 prop_assert_eq!(map.insert(hash, *ptr), false);
641 }
642 assert_collisions(&map, ptrs.len(), 0)?;
643
644 let ptr_to_remove = ptrs.remove(pick_remove_idx);
646 prop_assert_eq!(map.remove(hash, ptr_to_remove), true);
647 assert_ptrs_and_len(&mut map, hash, &ptrs)?;
648 assert_collisions(&map, ptrs.len(), 0)?;
649 prop_assert_eq!(sorted(&map.colliders[0]), sorted(&ptrs));
650 prop_assert_eq!(map.emptied_collider_slots, []);
651 }
652 }
653}