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 {
64 #[inline(always)]
66 pub(crate) fn downcast_ref<T: Any>(&self) -> Option<&T> {
67 (self as &dyn Any).downcast_ref()
68 }
69
70 #[inline(always)]
72 pub(crate) fn downcast_mut<T: Any>(&mut self) -> Option<&mut T> {
73 (self as &mut dyn Any).downcast_mut()
74 }
75}
76
77impl dyn EventAny {
78 #[inline(always)]
80 fn downcast_ref<T: Any>(&self) -> Option<&T> {
81 (self as &dyn Any).downcast_ref()
82 }
83
84 #[inline(always)]
86 fn downcast_mut<T: Any>(&mut self) -> Option<&mut T> {
87 (self as &mut dyn Any).downcast_mut()
88 }
89}
90
91#[cold]
93#[inline(never)]
94fn wrong_type() -> ! {
95 panic!("value stored with the key of another type")
96}
97
98#[derive(Copy, Clone)]
100struct CloneFns {
101 clone: fn(&dyn DebugAny) -> Box<dyn DebugAny>,
102 clone_into: fn(&mut dyn DebugAny, &dyn DebugAny),
103}
104
105impl CloneFns {
106 fn of<T: Clone + Debug + Send + Sync + 'static>() -> CloneFns {
107 CloneFns {
108 clone: |value| Box::new(value.as_any().downcast_ref::<T>().unwrap().clone()),
109 clone_into: |target, value| {
110 target
111 .as_any_mut()
112 .downcast_mut::<T>()
113 .unwrap()
114 .clone_from(value.as_any().downcast_ref::<T>().unwrap())
115 },
116 }
117 }
118}
119
120#[derive(Copy, Clone)]
122struct EventFns {
123 clone: fn(&dyn EventAny) -> Box<dyn EventAny>,
124 clone_into: fn(&mut dyn EventAny, &dyn EventAny),
125}
126
127impl EventFns {
128 fn of<T: Clone + Debug + Send + Sync + 'static>() -> EventFns {
129 EventFns {
130 clone: |value| Box::new(value.as_any().downcast_ref::<T>().unwrap().clone()),
131 clone_into: |target, value| {
132 target
133 .as_any_mut()
134 .downcast_mut::<T>()
135 .unwrap()
136 .clone_from(value.as_any().downcast_ref::<T>().unwrap())
137 },
138 }
139 }
140}
141
142struct EventEntry {
144 key: TypeKey,
145 active: bool,
148 value: Box<dyn EventAny>,
149 fns: EventFns,
150}
151
152#[derive(Default)]
158pub(crate) struct Extensions {
159 entries: Vec<(TypeKey, Box<dyn DebugAny>)>,
161 replayable: Vec<(TypeKey, CloneFns)>,
162 events: Vec<EventEntry>,
164 has_event_data: bool,
166}
167
168impl Debug for Extensions {
169 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
170 f.debug_map()
171 .entries(self.entries.iter().map(|(key, value)| (key, value)))
172 .entries(
173 self.events
174 .iter()
175 .filter(|entry| entry.active)
176 .map(|entry| (entry.key, &entry.value)),
177 )
178 .finish()
179 }
180}
181
182impl Extensions {
183 #[inline]
184 fn position(&self, key: TypeId) -> Option<usize> {
185 self.entries.iter().position(|(k, _)| k.0 == key)
186 }
187
188 #[inline]
189 pub(crate) fn get<T: Debug + Send + Sync + 'static>(&self) -> Option<&T> {
190 let index = self.position(TypeId::of::<T>())?;
191 self.entries[index].1.downcast_ref::<T>()
192 }
193
194 #[inline]
195 pub(crate) fn get_mut<T: Default + Debug + Send + Sync + 'static>(&mut self) -> &mut T {
196 let index = match self.position(TypeId::of::<T>()) {
197 Some(index) => index,
198 None => self.insert_default::<T>(),
199 };
200 match self.entries[index].1.downcast_mut::<T>() {
201 Some(value) => value,
202 None => wrong_type(),
203 }
204 }
205
206 #[cold]
207 fn insert_default<T: Default + Debug + Send + Sync + 'static>(&mut self) -> usize {
208 self.entries
209 .push((TypeKey::of::<T>(), Box::new(T::default())));
210 self.entries.len() - 1
211 }
212
213 pub(crate) fn set_replayable<T: Clone + Debug + Send + Sync + 'static>(&mut self) {
215 let key = TypeKey::of::<T>();
216 if !self.replayable.iter().any(|(k, _)| *k == key) {
217 self.replayable.push((key, CloneFns::of::<T>()));
218 }
219 }
220
221 #[cfg(test)]
223 pub fn has_event_data(&self) -> bool {
224 self.has_event_data
225 }
226
227 #[inline]
228 fn event_position(&self, key: TypeId) -> Option<usize> {
229 self.events.iter().position(|entry| entry.key.0 == key)
230 }
231
232 #[inline(always)]
237 pub(crate) fn event<T: Debug + Send + Sync + 'static>(&self) -> Option<&T> {
238 if !self.has_event_data {
239 return None;
240 }
241 self.find_event()
242 }
243
244 #[inline]
245 fn find_event<T: Debug + Send + Sync + 'static>(&self) -> Option<&T> {
246 let entry = &self.events[self.event_position(TypeId::of::<T>())?];
247 if !entry.active {
248 return None;
249 }
250 entry.value.downcast_ref::<T>()
251 }
252
253 #[inline]
257 pub(crate) fn event_mut<T: Default + Clone + Debug + Send + Sync + 'static>(
258 &mut self,
259 ) -> &mut T {
260 let index = match self.event_position(TypeId::of::<T>()) {
261 Some(index) => index,
262 None => self.insert_event::<T>(),
263 };
264 let entry = &mut self.events[index];
265 let Some(value) = entry.value.downcast_mut::<T>() else {
266 wrong_type()
267 };
268 if !entry.active {
269 value.clone_from(&T::default());
272 entry.active = true;
273 self.has_event_data = true;
274 }
275 value
276 }
277
278 pub(crate) fn take_event<T: Default + Debug + Send + Sync + 'static>(&mut self) -> Option<T> {
282 if !self.has_event_data {
283 return None;
284 }
285 let index = self.event_position(TypeId::of::<T>())?;
286 let entry = &mut self.events[index];
287 if !entry.active {
288 return None;
289 }
290 let value = core::mem::take((*entry.value).as_any_mut().downcast_mut::<T>()?);
292 entry.active = false;
293 self.has_event_data = self.events.iter().any(|entry| entry.active);
294 Some(value)
295 }
296
297 #[cold]
298 fn insert_event<T: Default + Clone + Debug + Send + Sync + 'static>(&mut self) -> usize {
299 self.events.push(EventEntry {
300 key: TypeKey::of::<T>(),
301 active: false,
302 value: Box::new(T::default()),
303 fns: EventFns::of::<T>(),
304 });
305 self.events.len() - 1
306 }
307
308 #[inline(always)]
310 pub(crate) fn clear_event_data(&mut self) {
311 if self.has_event_data {
312 self.deactivate_events();
313 }
314 }
315
316 #[inline(never)]
317 fn deactivate_events(&mut self) {
318 for entry in self.events.iter_mut() {
319 entry.active = false;
320 }
321 self.has_event_data = false;
322 }
323
324 #[inline]
329 pub(crate) fn snapshot_if_any(&self) -> Option<Snapshot> {
330 if self.replayable.is_empty() && !self.has_event_data {
331 return None;
332 }
333 Some(self.snapshot()).filter(|snapshot| !snapshot.is_empty())
334 }
335
336 pub(crate) fn snapshot(&self) -> Snapshot {
338 let mut snapshot = Snapshot::default();
339 if !self.replayable.is_empty() {
340 snapshot.replayable = self
341 .replayable
342 .iter()
343 .filter_map(|&(key, fns)| {
344 self.position(key.0)
345 .map(|index| (key, (fns.clone)(&*self.entries[index].1), fns))
346 })
347 .collect();
348 }
349 snapshot.events = self.capture_event_data();
350 snapshot
351 }
352
353 pub(crate) fn restore(&mut self, snapshot: &Snapshot) {
357 for (key, value, fns) in snapshot.replayable.iter() {
359 match self.position(key.0) {
360 Some(index) => (fns.clone_into)(&mut *self.entries[index].1, &**value),
361 None => self.entries.push((*key, (fns.clone)(&**value))),
362 }
363 }
364 self.restore_event_data(&snapshot.events);
365 }
366
367 pub(crate) fn capture_event_data(&self) -> EventData {
369 if !self.has_event_data {
370 return EventData::default();
371 }
372 EventData {
373 entries: self
374 .events
375 .iter()
376 .filter(|entry| entry.active)
377 .map(|entry| EventDataEntry {
378 key: entry.key,
379 value: (entry.fns.clone)(&*entry.value),
380 fns: entry.fns,
381 })
382 .collect(),
383 }
384 }
385
386 pub(crate) fn restore_event_data(&mut self, data: &EventData) {
388 self.clear_event_data();
389 self.attach_event_data(data);
390 }
391
392 pub(crate) fn attach_event_data(&mut self, data: &EventData) {
397 for entry in data.entries.iter() {
398 match self.event_position(entry.key.0) {
399 Some(index) => {
400 let target = &mut self.events[index];
401 (entry.fns.clone_into)(&mut *target.value, &*entry.value);
402 target.active = true;
403 }
404 None => self.events.push(EventEntry {
405 key: entry.key,
406 active: true,
407 value: (entry.fns.clone)(&*entry.value),
408 fns: entry.fns,
409 }),
410 }
411 self.has_event_data = true;
412 }
413 }
414}
415
416#[derive(Default)]
418pub(crate) struct Snapshot {
419 replayable: Vec<(TypeKey, Box<dyn DebugAny>, CloneFns)>,
420 events: EventData,
421}
422
423impl Snapshot {
424 pub(crate) fn is_empty(&self) -> bool {
426 self.replayable.is_empty() && self.events.is_empty()
427 }
428
429 pub(crate) fn event_data(&self) -> &EventData {
431 &self.events
432 }
433}
434
435impl Clone for Snapshot {
436 fn clone(&self) -> Snapshot {
437 Snapshot {
438 replayable: self
439 .replayable
440 .iter()
441 .map(|(key, value, fns)| (*key, (fns.clone)(&**value), *fns))
442 .collect(),
443 events: self.events.clone(),
444 }
445 }
446}
447
448impl Debug for Snapshot {
449 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
450 f.debug_map()
451 .entries(self.replayable.iter().map(|(key, value, _)| (key, value)))
452 .entries(
453 self.events
454 .entries
455 .iter()
456 .map(|entry| (entry.key, &entry.value)),
457 )
458 .finish()
459 }
460}
461
462#[derive(Default)]
495pub struct EventData {
496 entries: Vec<EventDataEntry>,
499}
500
501struct EventDataEntry {
502 key: TypeKey,
503 value: Box<dyn EventAny>,
504 fns: EventFns,
505}
506
507impl EventData {
508 pub const fn new() -> EventData {
510 EventData {
511 entries: Vec::new(),
512 }
513 }
514
515 pub fn is_empty(&self) -> bool {
517 self.entries.is_empty()
518 }
519
520 fn position(&self, key: TypeId) -> Option<usize> {
521 self.entries.iter().position(|entry| entry.key.0 == key)
522 }
523
524 pub fn get<T: Debug + Send + 'static>(&self) -> Option<&T> {
526 let entry = &self.entries[self.position(TypeId::of::<T>())?];
527 entry.value.downcast_ref::<T>()
528 }
529
530 pub fn get_mut<T: Default + Clone + Debug + Send + Sync + 'static>(&mut self) -> &mut T {
534 let index = match self.position(TypeId::of::<T>()) {
535 Some(index) => index,
536 None => {
537 self.entries.push(EventDataEntry {
538 key: TypeKey::of::<T>(),
539 value: Box::new(T::default()),
540 fns: EventFns::of::<T>(),
541 });
542 self.entries.len() - 1
543 }
544 };
545 let entry = &mut self.entries[index];
546 match entry.value.downcast_mut::<T>() {
547 Some(value) => value,
548 None => wrong_type(),
549 }
550 }
551
552 pub fn insert<T: Clone + Debug + Send + Sync + 'static>(&mut self, value: T) {
554 let entry = EventDataEntry {
555 key: TypeKey::of::<T>(),
556 value: Box::new(value),
557 fns: EventFns::of::<T>(),
558 };
559 match self.position(TypeId::of::<T>()) {
560 Some(index) => self.entries[index] = entry,
561 None => self.entries.push(entry),
562 }
563 }
564
565 pub fn remove<T: Debug + Send + 'static>(&mut self) -> Option<T> {
567 let entry = self.entries.remove(self.position(TypeId::of::<T>())?);
568 entry
569 .value
570 .into_any()
571 .downcast::<T>()
572 .ok()
573 .map(|value| *value)
574 }
575}
576
577impl Clone for EventData {
578 fn clone(&self) -> EventData {
579 EventData {
580 entries: self
581 .entries
582 .iter()
583 .map(|entry| EventDataEntry {
584 key: entry.key,
585 value: (entry.fns.clone)(&*entry.value),
586 fns: entry.fns,
587 })
588 .collect(),
589 }
590 }
591}
592
593impl Debug for EventData {
594 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
595 f.debug_map()
596 .entries(self.entries.iter().map(|entry| (entry.key, &entry.value)))
597 .finish()
598 }
599}
600
601#[test]
602fn test_extensions() {
603 #[derive(Debug, Default, Clone, PartialEq)]
604 struct A(u32);
605 #[derive(Debug, Default, Clone, PartialEq)]
606 struct B(String);
607 #[derive(Debug, Default, Clone, PartialEq)]
608 struct C(u8);
609
610 let mut ext = Extensions::default();
611 assert_eq!(ext.get::<A>(), None);
612 ext.get_mut::<A>().0 = 42;
613 ext.get_mut::<B>().0.push_str("hello");
614 ext.get_mut::<C>().0 = 1;
615 assert_eq!(ext.get::<A>(), Some(&A(42)));
616 assert_eq!(ext.get::<B>(), Some(&B("hello".into())));
617 assert_eq!(ext.get::<C>(), Some(&C(1)));
618
619 ext.set_replayable::<A>();
620 ext.set_replayable::<B>();
621 ext.set_replayable::<A>();
622 let snapshot = ext.snapshot();
623 ext.get_mut::<A>().0 = 1;
624 ext.get_mut::<B>().0.clear();
625 ext.get_mut::<C>().0 = 2;
626 ext.restore(&snapshot);
627 assert_eq!(ext.get::<A>(), Some(&A(42)));
628 assert_eq!(ext.get::<B>(), Some(&B("hello".into())));
629 assert_eq!(ext.get::<C>(), Some(&C(2)));
630
631 let mut other = Extensions::default();
633 other.restore(&snapshot.clone());
634 assert_eq!(other.get::<A>(), Some(&A(42)));
635 assert_eq!(other.get::<C>(), None);
636 assert_eq!(
637 format!("{:?}", other),
638 format!(
639 "{{{}: A(42), {}: B(\"hello\")}}",
640 core::any::type_name::<A>(),
641 core::any::type_name::<B>()
642 )
643 );
644}
645
646#[test]
647fn test_event_data() {
648 #[derive(Debug, Default, PartialEq)]
649 struct Tags(Vec<u64>);
650
651 impl Clone for Tags {
653 fn clone(&self) -> Tags {
654 Tags(self.0.clone())
655 }
656
657 fn clone_from(&mut self, source: &Tags) {
658 self.0.clone_from(&source.0);
659 }
660 }
661 #[derive(Debug, Default, Clone, PartialEq)]
662 struct Span(usize, usize);
663
664 let mut ext = Extensions::default();
665 assert!(!ext.has_event_data());
666 assert_eq!(ext.event::<Tags>(), None);
667
668 ext.event_mut::<Tags>().0.extend([1, 2]);
669 assert!(ext.has_event_data());
670 assert_eq!(ext.event::<Tags>(), Some(&Tags(vec![1, 2])));
671 assert_eq!(ext.event::<Span>(), None);
672
673 ext.clear_event_data();
675 assert!(!ext.has_event_data());
676 assert_eq!(ext.event::<Tags>(), None);
677 let tags = ext.event_mut::<Tags>();
678 assert_eq!(tags, &Tags(vec![]));
679 assert!(tags.0.capacity() >= 2);
680 tags.0.push(3);
681 *ext.event_mut::<Span>() = Span(1, 2);
682
683 let snapshot = ext.snapshot();
685 ext.clear_event_data();
686 *ext.event_mut::<Span>() = Span(3, 4);
687 ext.restore(&snapshot);
688 assert_eq!(ext.event::<Tags>(), Some(&Tags(vec![3])));
689 assert_eq!(ext.event::<Span>(), Some(&Span(1, 2)));
690
691 ext.clear_event_data();
692 ext.event_mut::<Span>();
693 let empty = ext.snapshot();
694 *ext.event_mut::<Tags>() = Tags(vec![4]);
695 ext.restore(&empty);
696 assert_eq!(ext.event::<Tags>(), None);
697 assert_eq!(ext.event::<Span>(), Some(&Span(0, 0)));
698
699 *ext.event_mut::<Tags>() = Tags(vec![5]);
701 assert_eq!(ext.take_event::<Tags>(), Some(Tags(vec![5])));
702 assert_eq!(ext.take_event::<Tags>(), None);
703 assert_eq!(ext.event::<Tags>(), None);
704 assert_eq!(ext.event::<Span>(), Some(&Span(0, 0)));
705 assert!(ext.capture_event_data().entries.len() == 1);
706 assert_eq!(ext.take_event::<Span>(), Some(Span(0, 0)));
707 assert!(!ext.has_event_data());
708 assert!(ext.capture_event_data().is_empty());
709}
710
711#[test]
712fn test_downcast_through_box() {
713 use alloc::sync::Arc;
714
715 let boxed: Box<dyn DebugAny> = Box::new(42u32);
718 assert_eq!(boxed.downcast_ref::<u32>(), Some(&42));
719 let arc: Arc<dyn DebugAny> = Arc::new(42u32);
720 assert_eq!(arc.downcast_ref::<u32>(), Some(&42));
721 let mut boxed: Box<dyn EventAny> = Box::new(1u32);
722 *boxed.downcast_mut::<u32>().unwrap() += 1;
723 assert_eq!(boxed.downcast_ref::<u32>(), Some(&2));
724}