Skip to main content

codex_extension_api/
state.rs

1use 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/// Typed values supplied before an [`ExtensionData`] scope is created.
11///
12/// Hosts may retain a clone when later operations must use the same initial
13/// inputs. Cloning freezes the attachment map and shares each value by `Arc`;
14/// values with interior mutability remain shared. This type does not install
15/// extensions or provide persistence.
16#[derive(Clone, Debug, Default)]
17pub struct ExtensionDataInit {
18    entries: HashMap<TypeId, ErasedData>,
19}
20
21impl ExtensionDataInit {
22    /// Creates an empty extension data initializer.
23    pub fn new() -> Self {
24        Self::default()
25    }
26
27    /// Stores `value` as the initial attachment of type `T`.
28    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    /// Returns a host-supplied initial attachment without creating a mutable scope.
38    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/// Typed extension-owned data attached to one host object.
48#[derive(Debug)]
49pub struct ExtensionData {
50    level_id: String,
51    entries: Mutex<HashMap<TypeId, ErasedData>>,
52}
53
54impl ExtensionData {
55    /// Creates an empty attachment map for one host-owned scope.
56    pub fn new(level_id: impl Into<String>) -> Self {
57        Self::new_with_init(level_id, ExtensionDataInit::default())
58    }
59
60    /// Creates an attachment map seeded with host-supplied initial data.
61    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    /// Returns the host identity for the scope this data is attached to.
69    pub fn level_id(&self) -> &str {
70        &self.level_id
71    }
72
73    /// Returns the attached value of type `T`, if one exists.
74    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    /// Returns the attached value of type `T`, inserting one from `init` when absent.
83    ///
84    /// The initializer runs while this map is locked, so it should stay cheap;
85    /// heavyweight lazy work belongs inside the attached value itself.
86    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    /// Stores `value` as the attachment of type `T`, returning any previous value.
98    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    /// Removes and returns the attached value of type `T`, if one exists.
108    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}