cfx_core/
ref_funcs.rs

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                // some error?
106                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}