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         for (_, instance) in store.store_data_mut().components.instances.iter_mut() {
41             let Some(instance) = instance.as_mut() else {
42                 continue;
43             };
44 
45             instance
46                 .get_mut()
47                 .concurrent_state_mut()
48                 .take_fibers_and_futures(&mut fibers, &mut futures);
49         }
50 
51         for mut fiber in fibers {
52             fiber.dispose(store);
53         }
54 
55         crate::component::concurrent::tls::set(store, move || drop(futures));
56     }
57 }
58 
59 impl StoreData {
60     pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
61         self.components.instances[id].as_ref().unwrap().get()
62     }
63 
64     pub(crate) fn component_instance_mut(
65         &mut self,
66         id: ComponentInstanceId,
67     ) -> Pin<&mut ComponentInstance> {
68         self.components.instances[id].as_mut().unwrap().get_mut()
69     }
70 }
71 
72 impl StoreOpaque {
73     pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
74         self.store_data().component_instance(id)
75     }
76 }
77 
78 /// A type used to represent an allocated `ComponentInstance` located within a
79 /// store.
80 ///
81 /// This type is held in various locations as a "safe index" into a store. This
82 /// encapsulates a `StoreId` which owns the instance as well as the index within
83 /// the store's list of which instance it's pointing to.
84 ///
85 /// This type can notably be used to index into a `StoreOpaque` to project out
86 /// the `ComponentInstance` that is associated with this id.
87 #[repr(C)] // used by reference in the C API
88 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
89 pub struct StoreComponentInstanceId {
90     store_id: StoreId,
91     instance: ComponentInstanceId,
92 }
93 
94 impl StoreComponentInstanceId {
95     pub(crate) fn new(
96         store_id: StoreId,
97         instance: ComponentInstanceId,
98     ) -> StoreComponentInstanceId {
99         StoreComponentInstanceId { store_id, instance }
100     }
101 
102     #[inline]
103     pub fn assert_belongs_to(&self, store: StoreId) {
104         self.store_id.assert_belongs_to(store)
105     }
106 
107     #[inline]
108     pub(crate) fn store_id(&self) -> StoreId {
109         self.store_id
110     }
111 
112     #[inline]
113     pub(crate) fn instance(&self) -> ComponentInstanceId {
114         self.instance
115     }
116 
117     /// Looks up the `vm::ComponentInstance` within `store` that this id points
118     /// to.
119     ///
120     /// # Panics
121     ///
122     /// Panics if `self` does not belong to `store`.
123     pub(crate) fn get<'a>(&self, store: &'a StoreOpaque) -> &'a ComponentInstance {
124         self.assert_belongs_to(store.id());
125         store.component_instance(self.instance)
126     }
127 
128     /// Mutable version of `get` above.
129     ///
130     /// # Panics
131     ///
132     /// Panics if `self` does not belong to `store`.
133     pub(crate) fn get_mut<'a>(&self, store: &'a mut StoreOpaque) -> Pin<&'a mut ComponentInstance> {
134         self.from_data_get_mut(store.store_data_mut())
135     }
136 
137     /// Same as `get_mut`, but borrows less of a store.
138     pub(crate) fn from_data_get_mut<'a>(
139         &self,
140         store: &'a mut StoreData,
141     ) -> Pin<&'a mut ComponentInstance> {
142         self.assert_belongs_to(store.id());
143         store.component_instance_mut(self.instance)
144     }
145 }
146