1#![deny(warnings, missing_docs, missing_debug_implementations)]
42extern crate alloc;
58
59use crate::slab::Slab;
60use alloc::alloc::Layout;
61use alloc::rc::Rc;
62use core::any::{Any, TypeId};
63use core::fmt;
64use core::marker::PhantomData;
65use core::mem::ManuallyDrop;
66use hashbrown::raw::{self as hashbrown_raw, RawTable};
67use std::hash::Hash;
68use thiserror::Error;
69
70fn type_id_to_u64(v: TypeId) -> u64 {
71 use std::mem::{size_of, transmute_copy};
72 match size_of::<TypeId>() {
73 8 => unsafe { transmute_copy::<_, u64>(&v) },
74 16 => unsafe { transmute_copy::<_, u128>(&v) as u64 },
75 _ => unreachable!(),
76 }
77}
78
79pub mod prealloc_tx;
80mod slab;
81
82pub struct Repo {
88 entity_catalog: Slab<EntityRecord>,
89 entity_table: RawTable<EntityStorage>,
90 marker: PhantomData<Rc<()>>,
91}
92
93impl Default for Repo {
94 fn default() -> Self {
95 Repo::new()
96 }
97}
98
99impl fmt::Debug for Repo {
100 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
101 write!(f, "{{Repo}}")?;
102 Ok(())
103 }
104}
105
106impl Repo {
107 pub fn new() -> Self {
109 Repo {
110 entity_catalog: Slab::new(),
111 entity_table: RawTable::new(),
112 marker: PhantomData,
113 }
114 }
115
116 fn find_type_storage_bucket<T: Any>(&self) -> Option<hashbrown_raw::Bucket<EntityStorage>> {
117 let type_id = TypeId::of::<T>();
118 let bucket = self.entity_table.find(type_id_to_u64(type_id), |storage| {
119 if storage.type_id != type_id {
120 return false;
121 }
122 debug_assert_eq!(storage.type_layout, Layout::new::<T>());
123 true
124 })?;
125 Some(bucket)
126 }
127
128 fn ensure_type_storage_bucket<T: Any>(&mut self) -> hashbrown_raw::Bucket<EntityStorage> {
129 if let Some(bucket) = self.find_type_storage_bucket::<T>() {
130 bucket
131 } else {
132 let type_id = TypeId::of::<T>();
133 self.entity_table.insert(
134 type_id_to_u64(type_id),
135 EntityStorage::new::<T>(),
136 |storage| type_id_to_u64(storage.type_id),
137 )
138 }
139 }
140
141 fn create_entity<T: Any>(&mut self, v: T) -> EntityId {
142 let bucket = self.ensure_type_storage_bucket::<T>();
143 let bucket_id = unsafe { self.entity_table.bucket_index(&bucket) };
144 let storage_mut = unsafe { bucket.as_mut() };
145 let (_, entity_id) = storage_mut
146 .view_mut::<T>(&mut self.entity_catalog, bucket_id)
147 .unwrap()
148 .insert(v);
149 EntityId(entity_id, PhantomData)
150 }
151
152 fn preallocate_entity<T: Any>(&mut self) -> EntityPtr<T> {
153 let bucket = self.ensure_type_storage_bucket::<T>();
154 let bucket_id = unsafe { self.entity_table.bucket_index(&bucket) };
155 let storage_mut = unsafe { bucket.as_mut() };
156 let storage_idx = storage_mut
157 .view_mut::<T>(&mut self.entity_catalog, bucket_id)
158 .unwrap()
159 .preallocate();
160 let record = EntityRecord {
161 bucket_id,
162 storage_idx,
163 };
164 EntityPtr {
165 record,
166 phantom: PhantomData,
167 }
168 }
169
170 fn init_preallocate_entity<T: Any>(
171 &mut self,
172 record: EntityRecord,
173 value: T,
174 ) -> Result<(), (Error, T)> {
175 let bucket = self.ensure_type_storage_bucket::<T>();
176 let bucket_id = unsafe { self.entity_table.bucket_index(&bucket) };
177 let storage_mut = unsafe { bucket.as_mut() };
178 let mut storage_view_mut = storage_mut
179 .view_mut::<T>(&mut self.entity_catalog, bucket_id)
180 .unwrap();
181 if !storage_view_mut.is_preallocated(record.storage_idx) {
182 return Err((Error::InvalidPtr, value));
183 }
184 let _ = storage_view_mut.init_preallocated(record.storage_idx, value);
185 Ok(())
186 }
187
188 fn cancel_preallocate_entity(&mut self, record: EntityRecord) -> Result<bool, Error> {
189 let bucket_id = record.bucket_id;
190 if bucket_id >= self.entity_table.buckets() {
191 return Err(Error::InvalidPtr);
192 }
193 let bucket = unsafe { self.entity_table.bucket(bucket_id) };
194 let storage_mut = unsafe { bucket.as_mut() };
195 match storage_mut.untyped_reattach_vacant(record.storage_idx) {
196 Ok(()) => Ok(true),
197 Err(slab::EntryError::EntryIsOccupied) => Ok(false),
198 Err(slab::EntryError::InvalidIdx) => Err(Error::InvalidPtr),
199 Err(slab::EntryError::EntryIsVacant) => Err(Error::InvalidPtr),
200 Err(slab::EntryError::EntryIsDetachedVacant) => unreachable!(),
201 }
202 }
203
204 fn validate_record_type<T: Any>(&self, record: EntityRecord) -> bool {
205 let bucket = match self.find_type_storage_bucket::<T>() {
206 Some(bucket) => bucket,
207 None => return false,
208 };
209 let bucket_id = unsafe { self.entity_table.bucket_index(&bucket) };
210 bucket_id == record.bucket_id
211 }
212
213 pub fn insert<T: Any>(&mut self, v: T) -> EntityPtr<T> {
216 let entity_id = self.create_entity(v);
217 let entity_ptr = entity_id
218 .cast_ptr(self)
219 .ok_or(Error::InternalError001)
220 .unwrap();
221 entity_ptr
222 }
223
224 pub fn insert_for_id<T: Any>(&mut self, v: T) -> EntityId {
227 self.create_entity(v)
228 }
229
230 pub fn remove<T: Any>(&mut self, entity_ptr: EntityPtr<T>) -> Option<T> {
232 let record = entity_ptr.record;
233 let bucket = unsafe { self.entity_table.bucket(record.bucket_id) };
234 let mut storage_view_mut = unsafe {
235 bucket
236 .as_mut()
237 .view_mut::<T>(&mut self.entity_catalog, record.bucket_id)
238 .ok()?
239 };
240 storage_view_mut.remove(record.storage_idx)
241 }
242}
243
244#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
245struct EntityRecord {
246 bucket_id: usize,
247 storage_idx: usize,
248}
249
250struct EntityStorage {
251 type_id: TypeId,
252 type_layout: Layout,
253 type_drop: fn(slab::ErasedSlab),
254 type_reattach_vacant: fn(&mut slab::ErasedSlab, usize) -> Result<(), slab::EntryError>,
255 type_data: slab::ErasedSlab,
256}
257
258struct EntityStorageEntry<T> {
259 data: T,
260 entity_id: usize,
261}
262
263impl Drop for EntityStorage {
264 fn drop(&mut self) {
265 (self.type_drop)(self.type_data);
266 }
267}
268
269impl EntityStorage {
270 fn new<T: Any>() -> Self {
271 EntityStorage {
272 type_id: TypeId::of::<T>(),
273 type_layout: Layout::new::<T>(),
274 type_drop: |erased_slab| {
275 let _ = unsafe { Slab::<EntityStorageEntry<T>>::from_erased_slab(erased_slab) };
276 },
277 type_reattach_vacant: |erased_slab, index| {
278 let mut slab = unsafe {
279 ManuallyDrop::new(Slab::<EntityStorageEntry<T>>::from_erased_slab(
280 *erased_slab,
281 ))
282 };
283 let result = slab.reattach_vacant(index);
284 *erased_slab = ManuallyDrop::into_inner(slab).into_erased_slab();
285 result
286 },
287 type_data: Slab::<EntityStorageEntry<T>>::new().into_erased_slab(),
288 }
289 }
290
291 fn view<T: Any>(&self) -> Result<EntityStorageView<'_, T>, Error> {
292 let type_id = TypeId::of::<T>();
293 if self.type_id != type_id {
294 return Err(Error::InvalidPtr);
295 }
296 debug_assert_eq!(self.type_layout, Layout::new::<T>());
297 Ok(EntityStorageView {
298 storage: self,
299 phantom: PhantomData,
300 })
301 }
302
303 fn view_mut<'a, T: Any>(
304 &'a mut self,
305 catalog: &'a mut Slab<EntityRecord>,
306 bucket_id: usize,
307 ) -> Result<EntityStorageViewMut<'a, T>, Error> {
308 let type_id = TypeId::of::<T>();
309 if self.type_id != type_id {
310 return Err(Error::InvalidPtr);
311 }
312 debug_assert_eq!(self.type_layout, Layout::new::<T>());
313 Ok(EntityStorageViewMut {
314 storage: self,
315 catalog,
316 bucket_id,
317 phantom: PhantomData,
318 })
319 }
320
321 unsafe fn raw_slab<T: Any>(&self) -> Result<ManuallyDrop<Slab<EntityStorageEntry<T>>>, Error> {
322 let type_id = TypeId::of::<T>();
323 if self.type_id != type_id {
324 return Err(Error::InvalidPtr);
325 }
326 debug_assert_eq!(self.type_layout, Layout::new::<T>());
327 Ok(ManuallyDrop::new(Slab::from_erased_slab(self.type_data)))
328 }
329 unsafe fn finish_raw_slab<T: Any>(
330 &mut self,
331 slab: ManuallyDrop<Slab<EntityStorageEntry<T>>>,
332 ) -> Result<(), Error> {
333 let type_id = TypeId::of::<T>();
334 if self.type_id != type_id {
335 return Err(Error::InvalidPtr);
336 }
337 debug_assert_eq!(self.type_layout, Layout::new::<T>());
338 self.type_data = ManuallyDrop::into_inner(slab).into_erased_slab();
339 Ok(())
340 }
341
342 fn untyped_reattach_vacant(&mut self, storage_idx: usize) -> Result<(), slab::EntryError> {
343 (self.type_reattach_vacant)(&mut self.type_data, storage_idx)
344 }
345}
346
347#[derive(Copy, Clone)]
348struct EntityStorageView<'a, T: Any> {
349 storage: &'a EntityStorage,
350 phantom: PhantomData<&'a [T]>,
351}
352
353impl<'a, T: Any> EntityStorageView<'a, T> {
354 fn is_valid_idx(&self, v: usize) -> bool {
355 let slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
356 slab.is_occupied(v)
357 }
358
359 fn entity_id(self, index: usize) -> EntityId {
360 let entity_id = unsafe {
361 self.storage
362 .type_data
363 .index::<EntityStorageEntry<T>>(index)
364 .entity_id
365 };
366 EntityId(entity_id, PhantomData)
367 }
368
369 fn index(self, index: usize) -> &'a T {
370 unsafe {
371 &self
372 .storage
373 .type_data
374 .index::<EntityStorageEntry<T>>(index)
375 .data
376 }
377 }
378}
379
380struct EntityStorageViewMut<'a, T: Any> {
381 storage: &'a mut EntityStorage,
382 catalog: &'a mut Slab<EntityRecord>,
383 bucket_id: usize,
384 phantom: PhantomData<&'a mut [T]>,
385}
386
387impl<'a, T: Any> EntityStorageViewMut<'a, T> {
388 fn is_valid_idx(&self, v: usize) -> bool {
389 let slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
390 slab.is_occupied(v)
391 }
392
393 fn preallocate(&mut self) -> usize {
394 let mut slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
395 let storage_idx = slab.detach_vacant();
396 unsafe { self.storage.finish_raw_slab(slab).unwrap() };
397 storage_idx
398 }
399
400 fn insert(&mut self, data: T) -> (usize, usize) {
401 let mut slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
402 let entity_id = self.catalog.detach_vacant();
403 let entry = EntityStorageEntry { data, entity_id };
404 let storage_idx = slab.push(entry);
405 self.catalog
406 .occupy_detached_vacant(
407 entity_id,
408 EntityRecord {
409 bucket_id: self.bucket_id,
410 storage_idx,
411 },
412 )
413 .unwrap();
414 unsafe { self.storage.finish_raw_slab(slab).unwrap() };
415 (storage_idx, entity_id)
416 }
417
418 fn is_preallocated(&self, storage_idx: usize) -> bool {
419 let slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
420 slab.is_detached_vacant(storage_idx)
421 }
422
423 fn init_preallocated(&mut self, storage_idx: usize, data: T) -> usize {
424 let mut slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
425 let entity_id = self.catalog.detach_vacant();
426 let entry = EntityStorageEntry { data, entity_id };
427 slab.occupy_detached_vacant(storage_idx, entry).unwrap();
428 self.catalog
429 .occupy_detached_vacant(
430 entity_id,
431 EntityRecord {
432 bucket_id: self.bucket_id,
433 storage_idx,
434 },
435 )
436 .unwrap();
437 unsafe { self.storage.finish_raw_slab(slab).unwrap() };
438 entity_id
439 }
440
441 fn remove(&mut self, storage_idx: usize) -> Option<T> {
442 let mut slab = unsafe { self.storage.raw_slab::<T>().unwrap() };
443 let entry = slab.remove(storage_idx)?;
444 self.catalog.remove(entry.entity_id);
445 Some(entry.data)
446 }
447
448 #[allow(dead_code)]
449 fn index(self, index: usize) -> &'a T {
450 unsafe {
451 &self
452 .storage
453 .type_data
454 .index::<EntityStorageEntry<T>>(index)
455 .data
456 }
457 }
458
459 fn index_mut(self, index: usize) -> &'a mut T {
460 unsafe {
461 &mut self
462 .storage
463 .type_data
464 .index_mut::<EntityStorageEntry<T>>(index)
465 .data
466 }
467 }
468}
469pub trait AsRepoRef {
471 fn as_repo_ref(&self) -> &Repo;
473}
474
475pub trait AsRepoMut: AsRepoRef {
477 fn as_repo_mut(&mut self) -> &mut Repo;
479}
480
481impl AsRepoRef for Repo {
482 fn as_repo_ref(&self) -> &Repo {
483 self
484 }
485}
486
487impl AsRepoMut for Repo {
488 fn as_repo_mut(&mut self) -> &mut Repo {
489 self
490 }
491}
492
493pub struct EntityId<R = Repo>(usize, PhantomData<R>);
495
496impl<R> Copy for EntityId<R> {}
497
498impl<R> Clone for EntityId<R> {
499 fn clone(&self) -> Self {
500 Self(self.0, self.1)
501 }
502}
503
504impl<R> PartialEq for EntityId<R> {
505 fn eq(&self, other: &Self) -> bool {
506 PartialEq::eq(&self.0, &other.0)
507 }
508}
509
510impl<R> Eq for EntityId<R> {}
511
512impl<R> PartialOrd for EntityId<R> {
513 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
514 PartialOrd::partial_cmp(&self.0, &other.0)
515 }
516}
517
518impl<R> Ord for EntityId<R> {
519 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
520 Ord::cmp(&self.0, &other.0)
521 }
522}
523
524impl<R> Hash for EntityId<R> {
525 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
526 Hash::hash(&self.0, state)
527 }
528}
529
530impl<R> fmt::Debug for EntityId<R> {
531 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
532 write!(f, "EntityId({})", self.0)
533 }
534}
535
536impl<R> EntityId<R> {
537 pub fn cast_repo<R2>(self) -> EntityId<R2> {
539 EntityId(self.0, PhantomData)
540 }
541}
542
543impl<R: AsRepoRef> EntityId<R> {
544 pub fn cast_ptr<T: Any>(self, repo: &R) -> Option<EntityPtr<T>> {
546 let repo = repo.as_repo_ref();
547 let &record = repo.entity_catalog.get(self.0)?;
548 let bucket = unsafe { repo.entity_table.bucket(record.bucket_id) };
549 let storage_view = unsafe { bucket.as_ref().view::<T>().ok()? };
550 debug_assert!(storage_view.is_valid_idx(record.storage_idx));
551 Some(EntityPtr {
552 record,
553 phantom: PhantomData,
554 })
555 }
556
557 pub fn cast_ref<T: Any>(self, repo: &R) -> Option<&'_ T> {
559 let repo = repo.as_repo_ref();
560 let &record = repo.entity_catalog.get(self.0)?;
561 let bucket = unsafe { repo.entity_table.bucket(record.bucket_id) };
562 let storage_view = unsafe { bucket.as_ref().view::<T>().ok()? };
563 debug_assert!(storage_view.is_valid_idx(record.storage_idx));
564 Some(storage_view.index(record.storage_idx))
565 }
566}
567
568impl<R: AsRepoMut> EntityId<R> {
569 pub fn cast_mut<T: Any>(self, repo: &mut R) -> Option<&'_ mut T> {
571 let repo = repo.as_repo_mut();
572 let &record = repo.entity_catalog.get(self.0)?;
573 let bucket = unsafe { repo.entity_table.bucket(record.bucket_id) };
574 let storage_view_mut = unsafe {
575 bucket
576 .as_mut()
577 .view_mut::<T>(&mut repo.entity_catalog, record.bucket_id)
578 .ok()?
579 };
580 debug_assert!(storage_view_mut.is_valid_idx(record.storage_idx));
581 Some(storage_view_mut.index_mut(record.storage_idx))
582 }
583}
584
585pub struct EntityPtr<T: Any, R = Repo> {
587 record: EntityRecord,
588 phantom: PhantomData<(*mut T, *mut R)>,
589}
590
591impl<T: Any, R> Copy for EntityPtr<T, R> {}
592
593impl<T: Any, R> Clone for EntityPtr<T, R> {
594 fn clone(&self) -> Self {
595 *self
596 }
597}
598
599impl<T: Any, R> PartialEq for EntityPtr<T, R> {
600 fn eq(&self, other: &Self) -> bool {
601 self.record.eq(&other.record)
602 }
603}
604
605impl<T: Any, R> PartialOrd for EntityPtr<T, R> {
606 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
607 self.record.partial_cmp(&other.record)
608 }
609}
610
611impl<T: Any, R> Eq for EntityPtr<T, R> {}
612
613impl<T: Any, R> Ord for EntityPtr<T, R> {
614 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
615 self.record.cmp(&other.record)
616 }
617}
618
619impl<T: Any, R> Hash for EntityPtr<T, R> {
620 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
621 self.record.hash(state)
622 }
623}
624
625unsafe impl<T: Any, R> Send for EntityPtr<T, R> {}
626unsafe impl<T: Any, R> Sync for EntityPtr<T, R> {}
627
628impl<T: Any, R> fmt::Debug for EntityPtr<T, R> {
629 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
630 write!(
631 f,
632 "EntityPtr(#{}, #{})",
633 self.record.bucket_id, self.record.storage_idx
634 )
635 }
636}
637
638impl<T: Any, R> EntityPtr<T, R> {
639 pub fn cast_repo<R2>(self) -> EntityPtr<T, R2> {
641 EntityPtr {
642 record: self.record,
643 phantom: PhantomData,
644 }
645 }
646}
647
648impl<T: Any, R: AsRepoRef> EntityPtr<T, R> {
649 pub fn entity_id(self, repo: &R) -> Result<EntityId<R>, Error> {
651 let repo = repo.as_repo_ref();
652 let record = self.record;
653 let bucket = unsafe { repo.entity_table.bucket(record.bucket_id) };
654 let storage_view = unsafe { bucket.as_ref().view::<T>()? };
655 debug_assert!(storage_view.is_valid_idx(record.storage_idx));
656 Ok(storage_view.entity_id(record.storage_idx).cast_repo())
657 }
658
659 pub fn get_ref(self, repo: &R) -> Result<&'_ T, Error> {
661 let repo = repo.as_repo_ref();
662 let record = self.record;
663 let bucket = unsafe { repo.entity_table.bucket(record.bucket_id) };
664 let storage_view = unsafe { bucket.as_ref().view::<T>()? };
665 debug_assert!(storage_view.is_valid_idx(record.storage_idx));
666 Ok(storage_view.index(record.storage_idx))
667 }
668}
669
670impl<T: Any, R: AsRepoMut> EntityPtr<T, R> {
671 pub fn get_mut(self, repo: &mut R) -> Result<&'_ mut T, Error> {
673 let repo = repo.as_repo_mut();
674 let record = self.record;
675 let bucket = unsafe { repo.entity_table.bucket(record.bucket_id) };
676 let storage_view_mut = unsafe {
677 bucket
678 .as_mut()
679 .view_mut::<T>(&mut repo.entity_catalog, record.bucket_id)?
680 };
681
682 debug_assert!(storage_view_mut.is_valid_idx(record.storage_idx));
683 Ok(storage_view_mut.index_mut(record.storage_idx))
684 }
685}
686
687#[derive(Error, Debug)]
689#[non_exhaustive]
690pub enum Error {
691 #[error("Invalid pointer is specified.")]
692 InvalidPtr,
694 #[error("Internal bucket inconsistency.")]
695 InternalError001,
697}
698
699#[test]
700fn test_repo_basic() {
701 fn s(v: &str) -> String {
702 v.to_string()
703 }
704
705 let mut repo = Repo::new();
706 let a = repo.insert_for_id(42i32);
707 let a_ptr = a.cast_ptr::<i32>(&repo).unwrap();
708 assert_eq!(42i32, *a_ptr.get_ref(&repo).unwrap());
709 let a_wrong_ptr = a.cast_ptr::<u32>(&repo);
710 assert_eq!(None, a_wrong_ptr);
711
712 let b = repo.insert_for_id("hello");
713 assert_eq!(42i32, *a_ptr.get_ref(&repo).unwrap());
714 let b_ref = b.cast_ref::<&'static str>(&repo).unwrap();
715 assert_eq!("hello", *b_ref);
716
717 let c = repo.insert_for_id(s("world"));
718 assert_eq!(s("world"), *c.cast_mut::<String>(&mut repo).unwrap());
719 *c.cast_mut::<String>(&mut repo).unwrap() = s("World");
720 assert_eq!(s("World"), *c.cast_mut::<String>(&mut repo).unwrap());
721}