1use crate::types::ScrObject;
2use serde::{de::DeserializeOwned, Deserialize, Serialize};
3use std::{
4 cell::{Cell, RefCell},
5 collections::HashMap,
6 rc::Rc,
7};
8
9mod ffi {
10 #[link(wasm_import_module = "host")]
11 extern "C" {
12 pub fn canonicalize_ref(ref_idx: u32, buffer: *mut i8, buffer_size: usize) -> i32;
13 }
14}
15
16thread_local! {
17 static BUFFER: RefCell<Vec<u8>> = RefCell::new(vec![0; 1 << 15]);
18 static RETVAL: RefCell<ScrObject> = RefCell::new(ScrObject { data: 0, length: 0 });
19 static HANDLERS: RefCell<HashMap<u32, Rc<InnerRefFunction>>> = RefCell::new(HashMap::new());
20 static REF_IDX: RefCell<u32> = RefCell::new(0);
21}
22
23#[no_mangle]
24pub unsafe extern "C" fn __cfx_call_ref(
25 ref_idx: u32,
26 args: *const u8,
27 args_len: usize,
28) -> *const ScrObject {
29 let args = std::slice::from_raw_parts(args, args_len);
30
31 HANDLERS.with(|handlers| {
32 let handlers = handlers.borrow();
33
34 if let Some(handler) = handlers.get(&ref_idx) {
35 BUFFER.with(|buf| {
36 handler.handle(args, buf);
37
38 RETVAL.with(|retval| {
39 let mut retval = retval.borrow_mut();
40 let buf = buf.borrow();
41
42 retval.data = buf.as_ptr() as _;
43 retval.length = buf.len() as _;
44 });
45 });
46 }
47 });
48
49 RETVAL.with(|scr| scr.as_ptr())
50}
51
52#[no_mangle]
53pub unsafe extern "C" fn __cfx_duplicate_ref(ref_idx: u32) -> u32 {
54 HANDLERS.with(|handlers| {
55 let handlers = handlers.borrow();
56 if let Some(handler) = handlers.get(&ref_idx) {
57 let refs = handler.refs.get();
58 handler.refs.set(refs + 1);
59 }
60 });
61
62 ref_idx
63}
64
65#[no_mangle]
66pub unsafe extern "C" fn __cfx_remove_ref(ref_idx: u32) {
67 HANDLERS.with(|handlers| {
68 let remove = {
69 let handlers = handlers.borrow();
70
71 if let Some(handler) = handlers.get(&ref_idx) {
72 let refs = handler.refs.get();
73 handler.refs.set(refs - 1);
74
75 refs <= 1
76 } else {
77 false
78 }
79 };
80
81 if remove {
82 handlers.borrow_mut().remove(&ref_idx);
83 }
84 });
85}
86
87fn canonicalize_ref(ref_idx: u32) -> String {
88 thread_local! {
89 static CANON_REF: RefCell<Vec<u8>> = RefCell::new(vec![0; 1024]);
90 }
91
92 CANON_REF.with(|vec| {
93 let mut resized = false;
94
95 loop {
96 let (buffer, buffer_size) = {
97 let mut vec = vec.borrow_mut();
98
99 (vec.as_mut_ptr() as *mut _, vec.capacity())
100 };
101
102 let result = unsafe { ffi::canonicalize_ref(ref_idx, buffer, buffer_size) };
103
104 if result == 0 {
105 return String::new();
107 }
108
109 if result < 0 {
110 if resized {
111 return String::new();
112 }
113
114 vec.borrow_mut().resize(result.abs() as _, 0);
115 resized = true;
116 } else {
117 unsafe {
118 let slice = std::slice::from_raw_parts(buffer as *mut u8, (result - 1) as _);
119
120 return std::str::from_utf8_unchecked(slice).to_owned();
121 }
122 }
123 }
124 })
125}
126
127struct InnerRefFunction {
128 _idx: u32,
129 func: Box<dyn Fn(&[u8], &RefCell<Vec<u8>>)>,
130 refs: Cell<i32>,
131}
132
133impl InnerRefFunction {
134 fn handle(&self, input: &[u8], output: &RefCell<Vec<u8>>) {
135 (self.func)(input, output);
136 }
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
140#[serde(rename = "_ExtStruct")]
141pub struct ExternRefFunction((i8, serde_bytes::ByteBuf));
142
143impl ExternRefFunction {
144 pub(crate) fn new(name: &str) -> ExternRefFunction {
145 let bytes = serde_bytes::ByteBuf::from(name.bytes().collect::<Vec<u8>>());
146 ExternRefFunction((10, bytes))
147 }
148
149 pub(crate) fn name(&self) -> &str {
150 unsafe { std::str::from_utf8_unchecked(&self.0 .1) }
151 }
152
153 pub fn invoke<Out, In>(&self, args: In) -> Option<Out>
154 where
155 In: Serialize,
156 Out: DeserializeOwned,
157 {
158 crate::invoker::invoke_ref_func(&self, args)
159 }
160}
161
162#[derive(Clone)]
163pub struct RefFunction {
164 name: String,
165 inner: Rc<InnerRefFunction>,
166}
167
168impl RefFunction {
169 pub fn new<Handler, Input, Output>(handler: Handler) -> RefFunction
170 where
171 Handler: Fn(Input) -> Output + 'static,
172 Input: DeserializeOwned,
173 Output: Serialize,
174 {
175 let idx = REF_IDX.with(|idx| {
176 let mut idx = idx.borrow_mut();
177 *idx += 1;
178 *idx
179 });
180
181 let name = canonicalize_ref(idx);
182
183 let func = move |input: &[u8], out_buf: &RefCell<Vec<u8>>| {
184 let input = rmp_serde::decode::from_read(input).unwrap();
185 let out = handler(input);
186 {
187 let mut out_buf = out_buf.borrow_mut();
188
189 unsafe {
190 out_buf.set_len(0);
191 }
192
193 let _ = rmp_serde::encode::write_named(&mut *out_buf, &out);
194 }
195 };
196
197 let inner = InnerRefFunction {
198 _idx: idx,
199 func: Box::new(func),
200 refs: Cell::new(0),
201 };
202
203 let inner = Rc::new(inner);
204
205 HANDLERS.with(|handlers| {
206 let mut handlers = handlers.borrow_mut();
207 handlers.insert(idx, inner.clone());
208 });
209
210 RefFunction { name, inner }
211 }
212
213 pub fn new_raw<Handler>(handler: Handler) -> RefFunction
214 where
215 Handler: Fn(&[u8]) -> Vec<u8> + 'static,
216 {
217 let idx = REF_IDX.with(|idx| {
218 let mut idx = idx.borrow_mut();
219 *idx += 1;
220 *idx
221 });
222
223 let name = canonicalize_ref(idx);
224
225 let func = move |input: &[u8], out_buf: &RefCell<Vec<u8>>| {
226 let out = handler(input);
227 {
228 let mut out_buf = out_buf.borrow_mut();
229
230 unsafe {
231 out_buf.set_len(0);
232 }
233
234 out_buf.extend(out.iter());
235 }
236 };
237
238 let inner = InnerRefFunction {
239 _idx: idx,
240 func: Box::new(func),
241 refs: Cell::new(0),
242 };
243
244 let inner = Rc::new(inner);
245
246 HANDLERS.with(|handlers| {
247 let mut handlers = handlers.borrow_mut();
248 handlers.insert(idx, inner.clone());
249 });
250
251 RefFunction { name, inner }
252 }
253
254 pub(crate) fn name(&self) -> &str {
255 &self.name
256 }
257
258 pub fn as_extern_ref_func(&self) -> ExternRefFunction {
259 ExternRefFunction::new(self.name())
260 }
261}