1 #[cfg(feature = "component-model-async")]
2 use crate::runtime::vm::VMStore;
3 use crate::runtime::vm::component::{ComponentInstance, OwnedComponentInstance};
4 use crate::store::{StoreData, StoreId, StoreOpaque};
5 #[cfg(feature = "component-model-async")]
6 use alloc::vec::Vec;
7 use core::pin::Pin;
8 use wasmtime_environ::PrimaryMap;
9 
10 #[derive(Default)]
11 pub struct ComponentStoreData {
12     instances: PrimaryMap<ComponentInstanceId, Option<OwnedComponentInstance>>,
13 }
14 
15 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
16 pub struct ComponentInstanceId(u32);
17 wasmtime_environ::entity_impl!(ComponentInstanceId);
18 
19 impl StoreData {
20     pub(crate) fn push_component_instance(
21         &mut self,
22         data: OwnedComponentInstance,
23     ) -> ComponentInstanceId {
24         let expected = data.get().id();
25         let ret = self.components.instances.push(Some(data));
26         assert_eq!(expected, ret);
27         ret
28     }
29 }
30 
31 impl ComponentStoreData {
32     pub fn next_component_instance_id(&self) -> ComponentInstanceId {
33         self.instances.next_key()
34     }
35 
36     #[cfg(feature = "component-model-async")]
37     pub(crate) fn drop_fibers_and_futures(store: &mut dyn VMStore) {
38         let mut fibers = Vec::new();
39         let mut futures = Vec::new();
40         store
41             .concurrent_state_mut()
42             .take_fibers_and_futures(&mut fibers, &mut futures);
43 
44         for mut fiber in fibers {
45             fiber.dispose(store);
46         }
47 
48         crate::component::concurrent::tls::set(store, move || drop(futures));
49     }
50 
51     #[cfg(feature = "component-model-async")]
52     pub(crate) fn assert_instance_states_empty(&mut self) {
53         for (_, instance) in self.instances.iter_mut() {
54             let Some(instance) = instance.as_mut() else {
55                 continue;
56             };
57 
58             assert!(
59                 instance
60                     .get_mut()
61                     .instance_states()
62                     .0
63                     .iter_mut()
64                     .all(|(_, state)| state.handle_table().is_empty()
65                         && state.concurrent_state().pending_is_empty())
66             );
67         }
68     }
69 }
70 
71 impl StoreData {
72     pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
73         self.components.instances[id].as_ref().unwrap().get()
74     }
75 
76     pub(crate) fn component_instance_mut(
77         &mut self,
78         id: ComponentInstanceId,
79     ) -> Pin<&mut ComponentInstance> {
80         self.components.instances[id].as_mut().unwrap().get_mut()
81     }
82 }
83 
84 impl StoreOpaque {
85     pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
86         self.store_data().component_instance(id)
87     }
88 }
89 
90 /// A type used to represent an allocated `ComponentInstance` located within a
91 /// store.
92 ///
93 /// This type is held in various locations as a "safe index" into a store. This
94 /// encapsulates a `StoreId` which owns the instance as well as the index within
95 /// the store's list of which instance it's pointing to.
96 ///
97 /// This type can notably be used to index into a `StoreOpaque` to project out
98 /// the `ComponentInstance` that is associated with this id.
99 #[repr(C)] // used by reference in the C API
100 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
101 pub struct StoreComponentInstanceId {
102     store_id: StoreId,
103     instance: ComponentInstanceId,
104 }
105 
106 impl StoreComponentInstanceId {
107     pub(crate) fn new(
108         store_id: StoreId,
109         instance: ComponentInstanceId,
110     ) -> StoreComponentInstanceId {
111         StoreComponentInstanceId { store_id, instance }
112     }
113 
114     #[inline]
115     pub fn assert_belongs_to(&self, store: StoreId) {
116         self.store_id.assert_belongs_to(store)
117     }
118 
119     #[inline]
120     pub(crate) fn store_id(&self) -> StoreId {
121         self.store_id
122     }
123 
124     #[inline]
125     pub(crate) fn instance(&self) -> ComponentInstanceId {
126         self.instance
127     }
128 
129     /// Looks up the `vm::ComponentInstance` within `store` that this id points
130     /// to.
131     ///
132     /// # Panics
133     ///
134     /// Panics if `self` does not belong to `store`.
135     pub(crate) fn get<'a>(&self, store: &'a StoreOpaque) -> &'a ComponentInstance {
136         self.assert_belongs_to(store.id());
137         store.component_instance(self.instance)
138     }
139 
140     /// Mutable version of `get` above.
141     ///
142     /// # Panics
143     ///
144     /// Panics if `self` does not belong to `store`.
145     pub(crate) fn get_mut<'a>(&self, store: &'a mut StoreOpaque) -> Pin<&'a mut ComponentInstance> {
146         self.from_data_get_mut(store.store_data_mut())
147     }
148 
149     /// Return a mutable `ComponentInstance` and a `ModuleRegistry`
150     /// from the store.
151     ///
152     /// # Panics
153     ///
154     /// Panics if `self` does not belong to `store`.
155     #[cfg(feature = "component-model-async")]
156     pub(crate) fn get_mut_and_registry<'a>(
157         &self,
158         store: &'a mut StoreOpaque,
159     ) -> (
160         Pin<&'a mut ComponentInstance>,
161         &'a crate::module::ModuleRegistry,
162     ) {
163         let (store_data, registry) = store.store_data_mut_and_registry();
164         let instance = self.from_data_get_mut(store_data);
165         (instance, registry)
166     }
167 
168     /// Same as `get_mut`, but borrows less of a store.
169     pub(crate) fn from_data_get_mut<'a>(
170         &self,
171         store: &'a mut StoreData,
172     ) -> Pin<&'a mut ComponentInstance> {
173         self.assert_belongs_to(store.id());
174         store.component_instance_mut(self.instance)
175     }
176 }
177