1use std::{
26 fmt::{Debug, Display},
27 hash::Hash,
28 marker::PhantomData,
29 ops::{Deref, DerefMut, Index, IndexMut},
30 slice, vec,
31};
32
33use crate::intern::Intern;
34use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeSeq};
35
36pub trait Identifier: Copy + Hash + From<usize> + Into<usize> + Eq + Ord + std::fmt::Debug {}
40
41impl Identifier for usize {}
42
43pub struct Identified<Id: Identifier, T> {
45 pub id: Id,
46 pub inner: T,
47}
48
49impl<Id: Identifier, T> Identified<Id, T> {
50 pub fn new(id: Id, data: T) -> Self {
52 Self { id, inner: data }
53 }
54}
55
56impl<Id: Identifier, T> Deref for Identified<Id, T> {
57 type Target = T;
58
59 fn deref(&self) -> &Self::Target {
60 &self.inner
61 }
62}
63
64impl<Id: Identifier, T> DerefMut for Identified<Id, T> {
65 fn deref_mut(&mut self) -> &mut Self::Target {
66 &mut self.inner
67 }
68}
69
70impl<'a, Id: Identifier, T> Identified<Id, &'a mut T> {
71 pub fn immutable(self) -> Identified<Id, &'a T> {
73 Identified::new(self.id, &*self.inner)
74 }
75}
76
77impl<Id: Identifier, T: Display> Display for Identified<Id, T> {
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 self.inner.fmt(f)
80 }
81}
82
83impl<Id: Identifier, T: Debug> Debug for Identified<Id, T> {
84 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85 self.inner.fmt(f)
86 }
87}
88
89pub struct Registry<Id: Identifier, T> {
121 chunks: Vec<Vec<T>>,
122 len: usize,
123 _marker: PhantomData<Id>,
124}
125
126#[inline]
128fn locate(n: usize) -> (usize, usize) {
129 let m = n + 1;
130 let k = (usize::BITS - 1 - m.leading_zeros()) as usize;
131 (k, m - (1 << k))
132}
133
134impl<Id: Identifier, T: Serialize> Serialize for Registry<Id, T> {
135 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
136 where
137 S: Serializer,
138 {
139 let mut seq = serializer.serialize_seq(Some(self.len))?;
144 for e in self.chunks.iter().flatten() {
145 seq.serialize_element(e)?;
146 }
147 seq.end()
148 }
149}
150
151impl<'de, Id: Identifier, T: Deserialize<'de>> Deserialize<'de> for Registry<Id, T> {
152 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
153 where
154 D: Deserializer<'de>,
155 {
156 Ok(Vec::<T>::deserialize(deserializer)?.into_iter().collect())
157 }
158}
159
160impl<Id: Identifier, T: Clone> Clone for Registry<Id, T> {
161 fn clone(&self) -> Self {
162 self.chunks.iter().flatten().cloned().collect()
166 }
167}
168
169impl<Id: Identifier, T: Debug> Debug for Registry<Id, T> {
170 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171 f.debug_list()
172 .entries(self.chunks.iter().flatten())
173 .finish()
174 }
175}
176
177impl<Id: Identifier, T: PartialEq> PartialEq for Registry<Id, T> {
178 fn eq(&self, other: &Self) -> bool {
179 self.len == other.len
180 && self
181 .chunks
182 .iter()
183 .flatten()
184 .eq(other.chunks.iter().flatten())
185 }
186}
187
188impl<Id: Identifier, T: Eq> Eq for Registry<Id, T> {}
189
190impl<Id: Identifier, T> Registry<Id, T> {
191 pub fn push(&mut self, e: T) -> Id {
193 let n = self.len;
194 let (k, offset) = locate(n);
195 if offset == 0 && self.chunks.len() == k {
196 self.chunks.push(Vec::with_capacity(1 << k));
200 }
201 self.chunks[k].push(e);
202 self.len += 1;
203 n.into()
204 }
205
206 pub fn len(&self) -> usize {
208 self.len
209 }
210
211 pub fn truncate(&mut self, len: usize) {
219 while self.len > len {
220 let (k, offset) = locate(self.len - 1);
221 self.chunks[k].truncate(offset);
222 self.len -= 1;
223 }
224 }
225
226 pub fn is_empty(&self) -> bool {
228 self.len == 0
229 }
230
231 #[inline]
232 #[track_caller]
233 fn at(&self, n: usize) -> &T {
234 debug_assert!(
235 n < self.len,
236 "registry index {n} out of bounds for length {}",
237 self.len
238 );
239 let (k, offset) = locate(n);
240 &self.chunks[k][offset]
241 }
242
243 #[inline]
244 #[track_caller]
245 fn at_mut(&mut self, n: usize) -> &mut T {
246 debug_assert!(
247 n < self.len,
248 "registry index {n} out of bounds for length {}",
249 self.len
250 );
251 let (k, offset) = locate(n);
252 &mut self.chunks[k][offset]
253 }
254
255 #[track_caller]
263 pub fn replace(&mut self, id: Id, value: T) -> T {
264 std::mem::replace(self.at_mut(id.into()), value)
265 }
266
267 #[track_caller]
272 pub fn get(&self, id: Id) -> Identified<Id, &T> {
273 Identified::new(id, self.at(id.into()))
274 }
275
276 #[track_caller]
281 pub fn get_mut(&mut self, id: Id) -> Identified<Id, &mut T> {
282 Identified::new(id, self.at_mut(id.into()))
283 }
284
285 pub fn iter(&self) -> Iter<'_, Id, T> {
306 Iter {
307 iter: self.chunks.iter().flatten(),
308 index: 0,
309 _marker: PhantomData,
310 }
311 }
312
313 pub fn iter_mut(&mut self) -> IterMut<'_, Id, T> {
335 IterMut {
336 iter: self.chunks.iter_mut().flatten(),
337 index: 0,
338 _marker: PhantomData,
339 }
340 }
341
342 pub fn select_mut(&mut self, ids: &[Id]) -> Vec<&mut T> {
358 let mut want: std::collections::HashMap<usize, usize> =
362 std::collections::HashMap::with_capacity(ids.len());
363 for (pos, id) in ids.iter().enumerate() {
364 let n: usize = (*id).into();
365 assert!(
366 n < self.len,
367 "select_mut: id index {n} out of bounds (len {})",
368 self.len
369 );
370 let prev = want.insert(n, pos);
371 assert!(prev.is_none(), "select_mut: duplicate id index {n}");
372 }
373
374 let mut slots: Vec<Option<&mut T>> = (0..ids.len()).map(|_| None).collect();
378 for item in self.iter_mut() {
379 let n: usize = item.id.into();
380 if let Some(&pos) = want.get(&n) {
381 slots[pos] = Some(item.inner);
382 }
383 }
384 slots
385 .into_iter()
386 .map(|slot| slot.expect("select_mut: requested id had no backing slot"))
387 .collect()
388 }
389}
390
391pub struct Iter<'a, Id: Identifier, T> {
392 iter: std::iter::Flatten<slice::Iter<'a, Vec<T>>>,
393 index: usize,
394 _marker: PhantomData<Id>,
395}
396
397impl<'a, Id: Identifier, T> Iterator for Iter<'a, Id, T> {
398 type Item = Identified<Id, &'a T>;
399
400 fn next(&mut self) -> Option<Self::Item> {
401 let value = self.iter.next()?;
402 let id = Id::from(self.index);
403 self.index += 1;
404
405 Some(Identified::new(id, value))
406 }
407}
408
409pub struct IterMut<'a, Id: Identifier, T> {
410 iter: std::iter::Flatten<slice::IterMut<'a, Vec<T>>>,
411 index: usize,
412 _marker: PhantomData<Id>,
413}
414
415impl<'a, Id: Identifier, T> Iterator for IterMut<'a, Id, T> {
416 type Item = Identified<Id, &'a mut T>;
417
418 fn next(&mut self) -> Option<Self::Item> {
419 let value = self.iter.next()?;
420 let id = Id::from(self.index);
421 self.index += 1;
422
423 Some(Identified::new(id, value))
424 }
425}
426
427pub struct IntoIter<Id: Identifier, T> {
428 iter: std::iter::Flatten<vec::IntoIter<Vec<T>>>,
429 index: usize,
430 _marker: PhantomData<Id>,
431}
432
433impl<Id: Identifier, T> Iterator for IntoIter<Id, T> {
434 type Item = Identified<Id, T>;
435
436 fn next(&mut self) -> Option<Self::Item> {
437 let value = self.iter.next()?;
438 let id = Id::from(self.index);
439 self.index += 1;
440
441 Some(Identified::new(id, value))
442 }
443}
444
445impl<Id: Identifier, T> Default for Registry<Id, T> {
446 fn default() -> Self {
447 Self {
448 chunks: Vec::new(),
449 len: 0,
450 _marker: PhantomData,
451 }
452 }
453}
454
455impl<'a, Id: Identifier, T> IntoIterator for &'a Registry<Id, T> {
456 type Item = Identified<Id, &'a T>;
457
458 type IntoIter = Iter<'a, Id, T>;
459
460 fn into_iter(self) -> Self::IntoIter {
461 self.iter()
462 }
463}
464
465impl<'a, Id: Identifier, T> IntoIterator for &'a mut Registry<Id, T> {
466 type Item = Identified<Id, &'a mut T>;
467
468 type IntoIter = IterMut<'a, Id, T>;
469
470 fn into_iter(self) -> Self::IntoIter {
471 self.iter_mut()
472 }
473}
474
475impl<Id: Identifier, T> IntoIterator for Registry<Id, T> {
476 type Item = Identified<Id, T>;
477
478 type IntoIter = IntoIter<Id, T>;
479
480 fn into_iter(self) -> Self::IntoIter {
481 IntoIter {
482 iter: self.chunks.into_iter().flatten(),
483 index: 0,
484 _marker: PhantomData,
485 }
486 }
487}
488
489impl<Id: Identifier, T> FromIterator<T> for Registry<Id, T> {
490 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
491 let mut reg = Self::default();
492 for e in iter {
493 reg.push(e);
494 }
495 reg
496 }
497}
498
499impl<Id: Identifier, T> Index<Id> for Registry<Id, T> {
500 type Output = T;
501
502 #[track_caller]
503 fn index(&self, index: Id) -> &Self::Output {
504 self.at(index.into())
505 }
506}
507
508impl<Id: Identifier, T> IndexMut<Id> for Registry<Id, T> {
509 #[track_caller]
510 fn index_mut(&mut self, index: Id) -> &mut Self::Output {
511 self.at_mut(index.into())
512 }
513}
514
515impl<Id: Identifier, T: Intern> Intern for Registry<Id, T> {
516 type Static = Registry<Id, T::Static>;
517
518 fn intern(self, pool: &mut super::intern::StringPool) -> Self::Static {
519 self.into_iter()
520 .map(|item| item.inner.intern(pool))
521 .collect()
522 }
523}
524
525#[cfg(test)]
526mod tests {
527 use super::*;
528 use jstd_derive::Identifier;
529
530 #[derive(Identifier)]
531 struct Id(usize);
532
533 #[test]
535 fn truncate_keeps_earlier_ids_and_chunk_capacity() {
536 let mut registry: Registry<Id, u32> = (0..7).collect();
537 let capacity: Vec<usize> = registry.chunks.iter().map(Vec::capacity).collect();
538
539 registry.truncate(2);
540 assert_eq!(registry.len(), 2);
541 assert_eq!(registry[Id(1)], 1);
542 assert_eq!(registry.chunks[0].capacity(), capacity[0]);
543 assert_eq!(registry.chunks[1].capacity(), capacity[1]);
544 assert!(registry.chunks[2].is_empty());
545
546 assert_eq!(registry.push(9), Id(2));
547 assert_eq!(registry.push(10), Id(3));
548 assert_eq!(registry.chunks[2].capacity(), capacity[2]);
549
550 registry.truncate(0);
551 assert!(registry.is_empty());
552 registry.truncate(5);
553 assert!(registry.is_empty());
554 }
555
556 #[test]
560 fn pushing_after_truncate_reuses_the_emptied_chunks() {
561 let mut registry: Registry<Id, u32> = (0..7).collect();
562 let chunks = registry.chunks.len();
563 let capacity: Vec<usize> = registry.chunks.iter().map(Vec::capacity).collect();
564 for epoch in 0..100 {
565 registry.truncate(0);
566 for i in 0..7 {
567 assert_eq!(registry.push(i + epoch), Id(i as usize));
568 }
569 assert_eq!(registry.chunks.len(), chunks, "epoch {epoch}");
570 let now: Vec<usize> = registry.chunks.iter().map(Vec::capacity).collect();
571 assert_eq!(now, capacity, "epoch {epoch}");
572 assert_eq!(registry[Id(6)], 6 + epoch);
573 }
574 for i in 7..15 {
576 registry.push(i);
577 }
578 assert_eq!(registry.chunks.len(), chunks + 1);
579 assert_eq!(registry[Id(14)], 14);
580 }
581
582 #[test]
583 fn locate_matches_doubling_layout() {
584 assert_eq!(locate(0), (0, 0)); assert_eq!(locate(1), (1, 0)); assert_eq!(locate(2), (1, 1));
587 assert_eq!(locate(3), (2, 0)); assert_eq!(locate(6), (2, 3));
589 assert_eq!(locate(7), (3, 0)); }
591
592 #[test]
594 fn push_index_iter_across_chunks() {
595 let mut reg = Registry::<Id, usize>::default();
596 let ids: Vec<Id> = (0..1000).map(|v| reg.push(v)).collect();
597 assert_eq!(reg.len(), 1000);
598 for (i, &id) in ids.iter().enumerate() {
599 assert_eq!(usize::from(id), i);
600 assert_eq!(reg[id], i);
601 }
602 let seen: Vec<usize> = reg.iter().map(|item| *item.inner).collect();
603 assert_eq!(seen, (0..1000).collect::<Vec<_>>());
604 }
605
606 #[test]
609 fn element_address_is_stable_across_pushes() {
610 let mut reg = Registry::<Id, usize>::default();
611 let first = reg.push(42);
612 let addr = ®[first] as *const usize;
613 for v in 0..10_000 {
614 reg.push(v);
615 }
616 assert_eq!(
617 ®[first] as *const usize, addr,
618 "address moved after growth"
619 );
620 assert_eq!(reg[first], 42);
621 }
622
623 #[test]
626 fn replace_swaps_in_place() {
627 let mut reg = Registry::<Id, i32>::default();
628 let a = reg.push(1);
629 let b = reg.push(2);
630 let addr_b = ®[b] as *const i32;
631 let old = reg.replace(b, 99);
632 assert_eq!(old, 2);
633 assert_eq!(reg[b], 99);
634 assert_eq!(reg[a], 1, "other ids untouched");
635 assert_eq!(®[b] as *const i32, addr_b, "slot address stable");
636 assert_eq!(reg.len(), 2, "len unchanged");
637 }
638
639 #[test]
640 #[cfg(debug_assertions)]
641 #[should_panic(expected = "registry index 0 out of bounds for length 0")]
642 fn indexing_reports_requested_index_and_length() {
643 let reg = Registry::<Id, usize>::default();
644 let _ = reg[Id::from(0)];
645 }
646
647 #[test]
651 fn select_mut_disjoint_in_input_order() {
652 let mut reg = Registry::<Id, usize>::default();
653 let ids: Vec<Id> = (0..100).map(|v| reg.push(v)).collect();
654
655 let picked = [ids[7], ids[0], ids[63], ids[64], ids[2]];
657 let refs = reg.select_mut(&picked);
658 assert_eq!(refs.len(), picked.len());
659 assert_eq!(
661 refs.iter().map(|r| **r).collect::<Vec<_>>(),
662 vec![7, 0, 63, 64, 2]
663 );
664 for r in refs {
666 *r += 1000;
667 }
668 for &id in &picked {
669 assert_eq!(reg[id], usize::from(id) + 1000);
670 }
671 assert_eq!(reg[ids[1]], 1);
673 }
674
675 #[test]
677 #[should_panic(expected = "duplicate id")]
678 fn select_mut_rejects_duplicates() {
679 let mut reg = Registry::<Id, usize>::default();
680 let a = reg.push(1);
681 reg.push(2);
682 let _ = reg.select_mut(&[a, a]);
683 }
684
685 #[test]
689 fn rebuild_from_flat_sequence() {
690 let reg: Registry<Id, i32> = (0..300).collect();
691 let flat: Vec<i32> = reg.iter().map(|item| *item.inner).collect();
692 let back: Registry<Id, i32> = flat.into_iter().collect();
693 assert_eq!(reg, back);
694 assert_eq!(back.len(), 300);
695 }
696
697 #[test]
701 fn clone_preserves_stability() {
702 let mut reg = Registry::<Id, usize>::default();
703 for v in 0..5 {
704 reg.push(v); }
706 let mut cloned = reg.clone();
707 assert_eq!(reg, cloned);
708 let last = Id::from(4);
709 let addr = &cloned[last] as *const usize;
710 cloned.push(99); assert_eq!(
712 &cloned[last] as *const usize, addr,
713 "clone's chunk reallocated"
714 );
715 }
716
717 #[test]
718 fn public_views_iterators_and_serialization_preserve_values() {
719 let mut reg = Registry::<Id, i32>::default();
720 assert!(reg.is_empty());
721 let first = reg.push(10);
722 let second = reg.push(20);
723 assert_eq!(reg.get(first).id, first);
724 assert_eq!(**reg.get(second), 20);
725
726 let mut mutable = reg.get_mut(first);
727 **mutable += 5;
728 let immutable = mutable.immutable();
729 assert_eq!(immutable.id, first);
730 assert_eq!(**immutable, 15);
731 assert_eq!(format!("{immutable}"), "15");
732 assert_eq!(format!("{immutable:?}"), "15");
733
734 let mut borrowed = (®).into_iter();
735 assert_eq!(borrowed.size_hint(), (0, None));
736 assert_eq!(borrowed.next().unwrap().id, first);
737 assert_eq!(borrowed.next().unwrap().id, second);
738 assert!(borrowed.next().is_none());
739
740 let owned: Vec<_> = reg.clone().into_iter().map(|entry| entry.inner).collect();
741 assert_eq!(owned, [15, 20]);
742 assert_eq!(format!("{reg:?}"), "[15, 20]");
743
744 let bytes = bincode::serde::encode_to_vec(®, bincode::config::standard()).unwrap();
745 let (decoded, used): (Registry<Id, i32>, _) =
746 bincode::serde::decode_from_slice(&bytes, bincode::config::standard()).unwrap();
747 assert_eq!(used, bytes.len());
748 assert_eq!(decoded, reg);
749 }
750}