1use crate::{
2 context::Context,
3 function::{FunctionHandle, FunctionQuery, FunctionQueryParameter},
4 registry::{Registry, RegistryHandle},
5 types::TypeQuery,
6};
7use intuicio_data::data_stack::DataStackPack;
8use std::{cell::RefCell, marker::PhantomData, sync::Arc};
9use typid::ID;
10
11thread_local! {
12 static GLOBAL_HOST_STACK: RefCell<Vec<(HostId, Host)>> = const{ RefCell::new(vec![]) };
13}
14
15pub type HostId = ID<Host>;
16
17#[derive(Clone)]
18pub struct HostProducer {
19 producer: Arc<Box<dyn Fn() -> Host + Send + Sync>>,
20}
21
22impl HostProducer {
23 pub fn new(f: impl Fn() -> Host + Send + Sync + 'static) -> Self {
24 Self {
25 producer: Arc::new(Box::new(f)),
26 }
27 }
28
29 pub fn produce(&self) -> Host {
30 (self.producer)()
31 }
32}
33
34pub struct Host {
35 context: Context,
36 registry: RegistryHandle,
37}
38
39impl Host {
40 pub fn new(context: Context, registry: RegistryHandle) -> Self {
41 Self { context, registry }
42 }
43
44 pub fn fork(&self) -> Self {
45 Self {
46 context: self.context.fork(),
47 registry: self.registry.clone(),
48 }
49 }
50
51 pub fn push_global(self) -> Result<HostId, Self> {
52 GLOBAL_HOST_STACK.with(|host| match host.try_borrow_mut() {
53 Ok(mut stack) => {
54 let id = HostId::new();
55 stack.push((id, self));
56 Ok(id)
57 }
58 Err(_) => Err(self),
59 })
60 }
61
62 pub fn pop_global() -> Option<Self> {
63 GLOBAL_HOST_STACK.with(move |stack| Some(stack.try_borrow_mut().ok()?.pop()?.1))
64 }
65
66 pub fn remove_global(id: HostId) -> Option<Self> {
67 GLOBAL_HOST_STACK.with(move |stack| {
68 let mut stack = stack.try_borrow_mut().ok()?;
69 let index = stack.iter().position(|(host_id, _)| host_id == &id)?;
70 Some(stack.remove(index).1)
71 })
72 }
73
74 pub fn with_global<T>(f: impl FnOnce(&mut Self) -> T) -> Option<T> {
75 GLOBAL_HOST_STACK.with(move |stack| {
76 let mut stack = stack.try_borrow_mut().ok()?;
77 let host = &mut stack.last_mut()?.1;
78 Some(f(host))
79 })
80 }
81
82 pub fn context(&mut self) -> &mut Context {
83 &mut self.context
84 }
85
86 pub fn registry(&self) -> &Registry {
87 &self.registry
88 }
89
90 pub fn context_and_registry(&mut self) -> (&mut Context, &Registry) {
91 (&mut self.context, &self.registry)
92 }
93
94 pub fn find_function(
95 &self,
96 name: &str,
97 module_name: &str,
98 type_name: Option<&str>,
99 ) -> Option<FunctionHandle> {
100 self.registry.find_function(FunctionQuery {
101 name: Some(name.into()),
102 module_name: Some(module_name.into()),
103 type_query: type_name.map(|type_name| TypeQuery {
104 name: Some(type_name.into()),
105 ..Default::default()
106 }),
107 ..Default::default()
108 })
109 }
110
111 pub fn call_function<O: DataStackPack, I: DataStackPack>(
112 &mut self,
113 name: &str,
114 module_name: &str,
115 type_name: Option<&str>,
116 ) -> Option<HostFunctionCall<I, O>> {
117 let inputs_query = I::pack_types()
118 .into_iter()
119 .map(|type_hash| FunctionQueryParameter {
120 type_query: Some(TypeQuery {
121 type_hash: Some(type_hash),
122 ..Default::default()
123 }),
124 ..Default::default()
125 })
126 .collect::<Vec<_>>();
127 let outputs_query = O::pack_types()
128 .into_iter()
129 .map(|type_hash| FunctionQueryParameter {
130 type_query: Some(TypeQuery {
131 type_hash: Some(type_hash),
132 ..Default::default()
133 }),
134 ..Default::default()
135 })
136 .collect::<Vec<_>>();
137 let handle = self.registry.find_function(FunctionQuery {
138 name: Some(name.into()),
139 module_name: Some(module_name.into()),
140 type_query: type_name.map(|type_name| TypeQuery {
141 name: Some(type_name.into()),
142 ..Default::default()
143 }),
144 inputs: inputs_query.into(),
145 outputs: outputs_query.into(),
146 ..Default::default()
147 })?;
148 Some(HostFunctionCall {
149 context: &mut self.context,
150 registry: &self.registry,
151 handle,
152 _phantom: Default::default(),
153 })
154 }
155}
156
157pub struct HostFunctionCall<'a, I: DataStackPack, O: DataStackPack> {
158 context: &'a mut Context,
159 registry: &'a Registry,
160 handle: FunctionHandle,
161 _phantom: PhantomData<(I, O)>,
162}
163
164impl<I: DataStackPack, O: DataStackPack> HostFunctionCall<'_, I, O> {
165 pub fn run(self, inputs: I) -> O {
166 self.handle.call(self.context, self.registry, inputs, false)
167 }
168}