codex_extension_api/
state.rs1use std::any::Any;
2use std::any::TypeId;
3use std::collections::HashMap;
4use std::sync::Arc;
5use std::sync::Mutex;
6use std::sync::PoisonError;
7
8type ErasedData = Arc<dyn Any + Send + Sync>;
9
10#[derive(Clone, Debug, Default)]
17pub struct ExtensionDataInit {
18 entries: HashMap<TypeId, ErasedData>,
19}
20
21impl ExtensionDataInit {
22 pub fn new() -> Self {
24 Self::default()
25 }
26
27 pub fn insert<T>(&mut self, value: T) -> Option<Arc<T>>
29 where
30 T: Any + Send + Sync,
31 {
32 self.entries
33 .insert(TypeId::of::<T>(), Arc::new(value))
34 .map(downcast_data)
35 }
36
37 pub fn get<T>(&self) -> Option<Arc<T>>
39 where
40 T: Any + Send + Sync,
41 {
42 let value = self.entries.get(&TypeId::of::<T>())?.clone();
43 Some(downcast_data(value))
44 }
45}
46
47#[derive(Debug)]
49pub struct ExtensionData {
50 level_id: String,
51 entries: Mutex<HashMap<TypeId, ErasedData>>,
52}
53
54impl ExtensionData {
55 pub fn new(level_id: impl Into<String>) -> Self {
57 Self::new_with_init(level_id, ExtensionDataInit::default())
58 }
59
60 pub fn new_with_init(level_id: impl Into<String>, init: ExtensionDataInit) -> Self {
62 Self {
63 level_id: level_id.into(),
64 entries: Mutex::new(init.entries),
65 }
66 }
67
68 pub fn level_id(&self) -> &str {
70 &self.level_id
71 }
72
73 pub fn get<T>(&self) -> Option<Arc<T>>
75 where
76 T: Any + Send + Sync,
77 {
78 let value = self.entries().get(&TypeId::of::<T>())?.clone();
79 Some(downcast_data(value))
80 }
81
82 pub fn get_or_init<T>(&self, init: impl FnOnce() -> T) -> Arc<T>
87 where
88 T: Any + Send + Sync,
89 {
90 let mut entries = self.entries();
91 let value = entries
92 .entry(TypeId::of::<T>())
93 .or_insert_with(|| Arc::new(init()));
94 downcast_data(Arc::clone(value))
95 }
96
97 pub fn insert<T>(&self, value: T) -> Option<Arc<T>>
99 where
100 T: Any + Send + Sync,
101 {
102 self.entries()
103 .insert(TypeId::of::<T>(), Arc::new(value))
104 .map(downcast_data)
105 }
106
107 pub fn remove<T>(&self) -> Option<Arc<T>>
109 where
110 T: Any + Send + Sync,
111 {
112 self.entries().remove(&TypeId::of::<T>()).map(downcast_data)
113 }
114
115 fn entries(&self) -> std::sync::MutexGuard<'_, HashMap<TypeId, ErasedData>> {
116 self.entries.lock().unwrap_or_else(PoisonError::into_inner)
117 }
118}
119
120fn downcast_data<T>(value: ErasedData) -> Arc<T>
121where
122 T: Any + Send + Sync,
123{
124 let Ok(value) = value.downcast::<T>() else {
125 unreachable!("typed extension data stored an incompatible value");
126 };
127 value
128}