1use alloc::boxed::Box;
2use alloc::vec::Vec;
3use core::any::{Any, TypeId, type_name};
4use core::fmt::{self, Debug};
5
6#[derive(Copy, Clone)]
7pub(crate) struct TypeKey(pub(crate) TypeId, &'static str);
8
9impl TypeKey {
10 pub(crate) fn of<T: 'static>() -> TypeKey {
11 TypeKey(TypeId::of::<T>(), type_name::<T>())
12 }
13}
14
15impl PartialEq for TypeKey {
16 fn eq(&self, other: &Self) -> bool {
17 self.0 == other.0
18 }
19}
20
21impl Eq for TypeKey {}
22
23impl Debug for TypeKey {
24 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25 write!(f, "{}", self.1)
26 }
27}
28
29pub(crate) trait DebugAny: Any + Debug + Send + Sync {
30 fn as_any(&self) -> &dyn Any;
31 fn as_any_mut(&mut self) -> &mut dyn Any;
32 fn into_any(self: Box<Self>) -> Box<dyn Any>;
33}
34
35impl<T: Any + Debug + Send + Sync + 'static> DebugAny for T {
36 fn as_any(&self) -> &dyn Any {
37 self
38 }
39
40 fn as_any_mut(&mut self) -> &mut dyn Any {
41 self
42 }
43
44 fn into_any(self: Box<Self>) -> Box<dyn Any> {
45 self
46 }
47}
48
49trait EventAny: DebugAny + Sync {}
52
53impl<T: DebugAny + Sync> EventAny for T {}
54
55impl dyn DebugAny {
63 #[inline(always)]
69 pub(crate) unsafe fn downcast_ref_unchecked<T: Any>(&self) -> &T {
70 debug_assert!(self.as_any().is::<T>());
71 unsafe { &*(self as *const dyn DebugAny).cast::<T>() }
73 }
74
75 #[inline(always)]
81 pub(crate) unsafe fn downcast_mut_unchecked<T: Any>(&mut self) -> &mut T {
82 debug_assert!(self.as_any_mut().is::<T>());
83 unsafe { &mut *(self as *mut dyn DebugAny).cast::<T>() }
85 }
86}
87
88impl dyn EventAny {
89 #[inline(always)]
95 unsafe fn downcast_ref_unchecked<T: Any>(&self) -> &T {
96 debug_assert!(self.as_any().is::<T>());
97 unsafe { &*(self as *const dyn EventAny).cast::<T>() }
99 }
100
101 #[inline(always)]
107 unsafe fn downcast_mut_unchecked<T: Any>(&mut self) -> &mut T {
108 debug_assert!(self.as_any_mut().is::<T>());
109 unsafe { &mut *(self as *mut dyn EventAny).cast::<T>() }
111 }
112}
113
114#[derive(Copy, Clone)]
116struct CloneFns {
117 clone: fn(&dyn DebugAny) -> Box<dyn DebugAny>,
118 clone_into: fn(&mut dyn DebugAny, &dyn DebugAny),
119}
120
121impl CloneFns {
122 fn of<T: Clone + Debug + Send + Sync + 'static>() -> CloneFns {
123 CloneFns {
124 clone: |value| Box::new(value.as_any().downcast_ref::<T>().unwrap().clone()),
125 clone_into: |target, value| {
126 target
127 .as_any_mut()
128 .downcast_mut::<T>()
129 .unwrap()
130 .clone_from(value.as_any().downcast_ref::<T>().unwrap())
131 },
132 }
133 }
134}
135
136#[derive(Copy, Clone)]
138struct EventFns {
139 clone: fn(&dyn EventAny) -> Box<dyn EventAny>,
140 clone_into: fn(&mut dyn EventAny, &dyn EventAny),
141}
142
143impl EventFns {
144 fn of<T: Clone + Debug + Send + Sync + 'static>() -> EventFns {
145 EventFns {
146 clone: |value| Box::new(value.as_any().downcast_ref::<T>().unwrap().clone()),
147 clone_into: |target, value| {
148 target
149 .as_any_mut()
150 .downcast_mut::<T>()
151 .unwrap()
152 .clone_from(value.as_any().downcast_ref::<T>().unwrap())
153 },
154 }
155 }
156}
157
158struct EventEntry {
160 key: TypeKey,
161 active: bool,
164 value: Box<dyn EventAny>,
165 fns: EventFns,
166}
167
168#[derive(Default)]
174pub(crate) struct Extensions {
175 entries: Vec<(TypeKey, Box<dyn DebugAny>)>,
177 replayable: Vec<(TypeKey, CloneFns)>,
178 events: Vec<EventEntry>,
180 has_event_data: bool,
182}
183
184impl Debug for Extensions {
185 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
186 f.debug_map()
187 .entries(self.entries.iter().map(|(key, value)| (key, value)))
188 .entries(
189 self.events
190 .iter()
191 .filter(|entry| entry.active)
192 .map(|entry| (entry.key, &entry.value)),
193 )
194 .finish()
195 }
196}
197
198impl Extensions {
199 #[inline]
200 fn position(&self, key: TypeId) -> Option<usize> {
201 self.entries.iter().position(|(k, _)| k.0 == key)
202 }
203
204 #[inline]
205 pub(crate) fn get<T: Debug + Send + Sync + 'static>(&self) -> Option<&T> {
206 let index = self.position(TypeId::of::<T>())?;
207 Some(unsafe { self.entries[index].1.downcast_ref_unchecked::<T>() })
209 }
210
211 #[inline]
212 pub(crate) fn get_mut<T: Default + Debug + Send + Sync + 'static>(&mut self) -> &mut T {
213 let index = match self.position(TypeId::of::<T>()) {
214 Some(index) => index,
215 None => self.insert_default::<T>(),
216 };
217 unsafe { self.entries[index].1.downcast_mut_unchecked::<T>() }
219 }
220
221 #[cold]
222 fn insert_default<T: Default + Debug + Send + Sync + 'static>(&mut self) -> usize {
223 self.entries
224 .push((TypeKey::of::<T>(), Box::new(T::default())));
225 self.entries.len() - 1
226 }
227
228 pub(crate) fn set_replayable<T: Clone + Debug + Send + Sync + 'static>(&mut self) {
230 let key = TypeKey::of::<T>();
231 if !self.replayable.iter().any(|(k, _)| *k == key) {
232 self.replayable.push((key, CloneFns::of::<T>()));
233 }
234 }
235
236 #[cfg(test)]
238 pub fn has_event_data(&self) -> bool {
239 self.has_event_data
240 }
241
242 #[inline]
243 fn event_position(&self, key: TypeId) -> Option<usize> {
244 self.events.iter().position(|entry| entry.key.0 == key)
245 }
246
247 #[inline(always)]
252 pub(crate) fn event<T: Debug + Send + Sync + 'static>(&self) -> Option<&T> {
253 if !self.has_event_data {
254 return None;
255 }
256 self.find_event()
257 }
258
259 #[inline]
260 fn find_event<T: Debug + Send + Sync + 'static>(&self) -> Option<&T> {
261 let entry = &self.events[self.event_position(TypeId::of::<T>())?];
262 if !entry.active {
263 return None;
264 }
265 Some(unsafe { entry.value.downcast_ref_unchecked::<T>() })
267 }
268
269 #[inline]
273 pub(crate) fn event_mut<T: Default + Clone + Debug + Send + Sync + 'static>(
274 &mut self,
275 ) -> &mut T {
276 let index = match self.event_position(TypeId::of::<T>()) {
277 Some(index) => index,
278 None => self.insert_event::<T>(),
279 };
280 let entry = &mut self.events[index];
281 let value = unsafe { entry.value.downcast_mut_unchecked::<T>() };
283 if !entry.active {
284 value.clone_from(&T::default());
287 entry.active = true;
288 self.has_event_data = true;
289 }
290 value
291 }
292
293 pub(crate) fn take_event<T: Default + Debug + Send + Sync + 'static>(&mut self) -> Option<T> {
297 if !self.has_event_data {
298 return None;
299 }
300 let index = self.event_position(TypeId::of::<T>())?;
301 let entry = &mut self.events[index];
302 if !entry.active {
303 return None;
304 }
305 let value = core::mem::take((*entry.value).as_any_mut().downcast_mut::<T>()?);
307 entry.active = false;
308 self.has_event_data = self.events.iter().any(|entry| entry.active);
309 Some(value)
310 }
311
312 #[cold]
313 fn insert_event<T: Default + Clone + Debug + Send + Sync + 'static>(&mut self) -> usize {
314 self.events.push(EventEntry {
315 key: TypeKey::of::<T>(),
316 active: false,
317 value: Box::new(T::default()),
318 fns: EventFns::of::<T>(),
319 });
320 self.events.len() - 1
321 }
322
323 #[inline(always)]
325 pub(crate) fn clear_event_data(&mut self) {
326 if self.has_event_data {
327 self.deactivate_events();
328 }
329 }
330
331 #[inline(never)]
332 fn deactivate_events(&mut self) {
333 for entry in self.events.iter_mut() {
334 entry.active = false;
335 }
336 self.has_event_data = false;
337 }
338
339 #[inline]
344 pub(crate) fn snapshot_if_any(&self) -> Option<Snapshot> {
345 if self.replayable.is_empty() && !self.has_event_data {
346 return None;
347 }
348 Some(self.snapshot()).filter(|snapshot| !snapshot.is_empty())
349 }
350
351 pub(crate) fn snapshot(&self) -> Snapshot {
353 let mut snapshot = Snapshot::default();
354 if !self.replayable.is_empty() {
355 snapshot.replayable = self
356 .replayable
357 .iter()
358 .filter_map(|&(key, fns)| {
359 self.position(key.0)
360 .map(|index| (key, (fns.clone)(&*self.entries[index].1), fns))
361 })
362 .collect();
363 }
364 snapshot.events = self.capture_event_data();
365 snapshot
366 }
367
368 pub(crate) fn restore(&mut self, snapshot: &Snapshot) {
372 for (key, value, fns) in snapshot.replayable.iter() {
374 match self.position(key.0) {
375 Some(index) => (fns.clone_into)(&mut *self.entries[index].1, &**value),
376 None => self.entries.push((*key, (fns.clone)(&**value))),
377 }
378 }
379 self.restore_event_data(&snapshot.events);
380 }
381
382 pub(crate) fn capture_event_data(&self) -> EventData {
384 if !self.has_event_data {
385 return EventData::default();
386 }
387 EventData {
388 entries: self
389 .events
390 .iter()
391 .filter(|entry| entry.active)
392 .map(|entry| EventDataEntry {
393 key: entry.key,
394 value: (entry.fns.clone)(&*entry.value),
395 fns: entry.fns,
396 })
397 .collect(),
398 }
399 }
400
401 pub(crate) fn restore_event_data(&mut self, data: &EventData) {
403 self.clear_event_data();
404 self.attach_event_data(data);
405 }
406
407 pub(crate) fn attach_event_data(&mut self, data: &EventData) {
412 for entry in data.entries.iter() {
413 match self.event_position(entry.key.0) {
414 Some(index) => {
415 let target = &mut self.events[index];
416 (entry.fns.clone_into)(&mut *target.value, &*entry.value);
417 target.active = true;
418 }
419 None => self.events.push(EventEntry {
420 key: entry.key,
421 active: true,
422 value: (entry.fns.clone)(&*entry.value),
423 fns: entry.fns,
424 }),
425 }
426 self.has_event_data = true;
427 }
428 }
429}
430
431#[derive(Default)]
433pub(crate) struct Snapshot {
434 replayable: Vec<(TypeKey, Box<dyn DebugAny>, CloneFns)>,
435 events: EventData,
436}
437
438impl Snapshot {
439 pub(crate) fn is_empty(&self) -> bool {
441 self.replayable.is_empty() && self.events.is_empty()
442 }
443
444 pub(crate) fn event_data(&self) -> &EventData {
446 &self.events
447 }
448}
449
450impl Clone for Snapshot {
451 fn clone(&self) -> Snapshot {
452 Snapshot {
453 replayable: self
454 .replayable
455 .iter()
456 .map(|(key, value, fns)| (*key, (fns.clone)(&**value), *fns))
457 .collect(),
458 events: self.events.clone(),
459 }
460 }
461}
462
463impl Debug for Snapshot {
464 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
465 f.debug_map()
466 .entries(self.replayable.iter().map(|(key, value, _)| (key, value)))
467 .entries(
468 self.events
469 .entries
470 .iter()
471 .map(|entry| (entry.key, &entry.value)),
472 )
473 .finish()
474 }
475}
476
477#[derive(Default)]
510pub struct EventData {
511 entries: Vec<EventDataEntry>,
514}
515
516struct EventDataEntry {
517 key: TypeKey,
518 value: Box<dyn EventAny>,
519 fns: EventFns,
520}
521
522impl EventData {
523 pub const fn new() -> EventData {
525 EventData {
526 entries: Vec::new(),
527 }
528 }
529
530 pub fn is_empty(&self) -> bool {
532 self.entries.is_empty()
533 }
534
535 fn position(&self, key: TypeId) -> Option<usize> {
536 self.entries.iter().position(|entry| entry.key.0 == key)
537 }
538
539 pub fn get<T: Debug + Send + 'static>(&self) -> Option<&T> {
541 let entry = &self.entries[self.position(TypeId::of::<T>())?];
542 Some(unsafe { entry.value.downcast_ref_unchecked::<T>() })
544 }
545
546 pub fn get_mut<T: Default + Clone + Debug + Send + Sync + 'static>(&mut self) -> &mut T {
550 let index = match self.position(TypeId::of::<T>()) {
551 Some(index) => index,
552 None => {
553 self.entries.push(EventDataEntry {
554 key: TypeKey::of::<T>(),
555 value: Box::new(T::default()),
556 fns: EventFns::of::<T>(),
557 });
558 self.entries.len() - 1
559 }
560 };
561 let entry = &mut self.entries[index];
562 unsafe { entry.value.downcast_mut_unchecked::<T>() }
564 }
565
566 pub fn insert<T: Clone + Debug + Send + Sync + 'static>(&mut self, value: T) {
568 let entry = EventDataEntry {
569 key: TypeKey::of::<T>(),
570 value: Box::new(value),
571 fns: EventFns::of::<T>(),
572 };
573 match self.position(TypeId::of::<T>()) {
574 Some(index) => self.entries[index] = entry,
575 None => self.entries.push(entry),
576 }
577 }
578
579 pub fn remove<T: Debug + Send + 'static>(&mut self) -> Option<T> {
581 let entry = self.entries.remove(self.position(TypeId::of::<T>())?);
582 entry
583 .value
584 .into_any()
585 .downcast::<T>()
586 .ok()
587 .map(|value| *value)
588 }
589}
590
591impl Clone for EventData {
592 fn clone(&self) -> EventData {
593 EventData {
594 entries: self
595 .entries
596 .iter()
597 .map(|entry| EventDataEntry {
598 key: entry.key,
599 value: (entry.fns.clone)(&*entry.value),
600 fns: entry.fns,
601 })
602 .collect(),
603 }
604 }
605}
606
607impl Debug for EventData {
608 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
609 f.debug_map()
610 .entries(self.entries.iter().map(|entry| (entry.key, &entry.value)))
611 .finish()
612 }
613}
614
615#[test]
616fn test_extensions() {
617 #[derive(Debug, Default, Clone, PartialEq)]
618 struct A(u32);
619 #[derive(Debug, Default, Clone, PartialEq)]
620 struct B(String);
621 #[derive(Debug, Default, Clone, PartialEq)]
622 struct C(u8);
623
624 let mut ext = Extensions::default();
625 assert_eq!(ext.get::<A>(), None);
626 ext.get_mut::<A>().0 = 42;
627 ext.get_mut::<B>().0.push_str("hello");
628 ext.get_mut::<C>().0 = 1;
629 assert_eq!(ext.get::<A>(), Some(&A(42)));
630 assert_eq!(ext.get::<B>(), Some(&B("hello".into())));
631 assert_eq!(ext.get::<C>(), Some(&C(1)));
632
633 ext.set_replayable::<A>();
634 ext.set_replayable::<B>();
635 ext.set_replayable::<A>();
636 let snapshot = ext.snapshot();
637 ext.get_mut::<A>().0 = 1;
638 ext.get_mut::<B>().0.clear();
639 ext.get_mut::<C>().0 = 2;
640 ext.restore(&snapshot);
641 assert_eq!(ext.get::<A>(), Some(&A(42)));
642 assert_eq!(ext.get::<B>(), Some(&B("hello".into())));
643 assert_eq!(ext.get::<C>(), Some(&C(2)));
644
645 let mut other = Extensions::default();
647 other.restore(&snapshot.clone());
648 assert_eq!(other.get::<A>(), Some(&A(42)));
649 assert_eq!(other.get::<C>(), None);
650 assert_eq!(
651 format!("{:?}", other),
652 format!(
653 "{{{}: A(42), {}: B(\"hello\")}}",
654 core::any::type_name::<A>(),
655 core::any::type_name::<B>()
656 )
657 );
658}
659
660#[test]
661fn test_event_data() {
662 #[derive(Debug, Default, PartialEq)]
663 struct Tags(Vec<u64>);
664
665 impl Clone for Tags {
667 fn clone(&self) -> Tags {
668 Tags(self.0.clone())
669 }
670
671 fn clone_from(&mut self, source: &Tags) {
672 self.0.clone_from(&source.0);
673 }
674 }
675 #[derive(Debug, Default, Clone, PartialEq)]
676 struct Span(usize, usize);
677
678 let mut ext = Extensions::default();
679 assert!(!ext.has_event_data());
680 assert_eq!(ext.event::<Tags>(), None);
681
682 ext.event_mut::<Tags>().0.extend([1, 2]);
683 assert!(ext.has_event_data());
684 assert_eq!(ext.event::<Tags>(), Some(&Tags(vec![1, 2])));
685 assert_eq!(ext.event::<Span>(), None);
686
687 ext.clear_event_data();
689 assert!(!ext.has_event_data());
690 assert_eq!(ext.event::<Tags>(), None);
691 let tags = ext.event_mut::<Tags>();
692 assert_eq!(tags, &Tags(vec![]));
693 assert!(tags.0.capacity() >= 2);
694 tags.0.push(3);
695 *ext.event_mut::<Span>() = Span(1, 2);
696
697 let snapshot = ext.snapshot();
699 ext.clear_event_data();
700 *ext.event_mut::<Span>() = Span(3, 4);
701 ext.restore(&snapshot);
702 assert_eq!(ext.event::<Tags>(), Some(&Tags(vec![3])));
703 assert_eq!(ext.event::<Span>(), Some(&Span(1, 2)));
704
705 ext.clear_event_data();
706 ext.event_mut::<Span>();
707 let empty = ext.snapshot();
708 *ext.event_mut::<Tags>() = Tags(vec![4]);
709 ext.restore(&empty);
710 assert_eq!(ext.event::<Tags>(), None);
711 assert_eq!(ext.event::<Span>(), Some(&Span(0, 0)));
712
713 *ext.event_mut::<Tags>() = Tags(vec![5]);
715 assert_eq!(ext.take_event::<Tags>(), Some(Tags(vec![5])));
716 assert_eq!(ext.take_event::<Tags>(), None);
717 assert_eq!(ext.event::<Tags>(), None);
718 assert_eq!(ext.event::<Span>(), Some(&Span(0, 0)));
719 assert!(ext.capture_event_data().entries.len() == 1);
720 assert_eq!(ext.take_event::<Span>(), Some(Span(0, 0)));
721 assert!(!ext.has_event_data());
722 assert!(ext.capture_event_data().is_empty());
723}
724
725#[test]
726fn test_downcast_through_box() {
727 use alloc::sync::Arc;
728
729 let boxed: Box<dyn DebugAny> = Box::new(42u32);
732 assert_eq!(unsafe { *boxed.downcast_ref_unchecked::<u32>() }, 42);
734 let arc: Arc<dyn DebugAny> = Arc::new(42u32);
735 assert_eq!(unsafe { *arc.downcast_ref_unchecked::<u32>() }, 42);
737 let mut boxed: Box<dyn EventAny> = Box::new(1u32);
738 unsafe { *boxed.downcast_mut_unchecked::<u32>() += 1 };
740 assert_eq!(unsafe { *boxed.downcast_ref_unchecked::<u32>() }, 2);
742}