intuicio_core/
host.rs

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}