clique_wasm_runtime/
lib.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, Mutex},
4};
5
6pub use clique_wasm_types::{val::Val, val_type::ValType};
7pub use wasmtime::Config;
8use wasmtime::{Caller, Engine, Func, Linker, Memory, Module, Store};
9
10const MEMORY_NAME: &'static str = "memory";
11const MAIN_METHOD_NAME: &'static str = "evaluate";
12const CONSOLE_LOG_METHOD_NAME: &'static str = "console_log_internal";
13const GET_PARAM_LEN_METHOD_NAME: &'static str = "get_param_len_internal";
14const GET_PARAM_VALUE_METHOD_NAME: &'static str = "get_param_value_internal";
15const SET_OUTPUT_METHOD_NAME: &'static str = "set_output_internal";
16
17pub struct Runtime {
18    function: Func,
19    store: Store<()>,
20    params: Arc<Mutex<HashMap<String, Val>>>,
21    outputs: Arc<Mutex<HashMap<String, Val>>>,
22}
23
24impl Runtime {
25    fn get_memory(caller: &mut Caller<'_, ()>) -> anyhow::Result<Memory> {
26        match caller
27            .get_export(&MEMORY_NAME)
28            .and_then(|e| e.into_memory())
29        {
30            Some(v) => Ok(v),
31            None => anyhow::bail!("Can't find `{}` export", MEMORY_NAME),
32        }
33    }
34
35    fn get_key(
36        caller: &mut Caller<'_, ()>,
37        key_ptr: i32,
38        key_len: i32,
39    ) -> anyhow::Result<Option<String>> {
40        let memory = Self::get_memory(caller)?;
41        if let Some(data) = memory
42            .data(&caller)
43            .get(key_ptr as usize..(key_ptr + key_len) as usize)
44        {
45            Ok(Some(String::from_utf8(data.to_vec())?))
46        } else {
47            Ok(None)
48        }
49    }
50
51    fn get_param_len(
52        caller: &mut Caller<'_, ()>,
53        key_ptr: i32,
54        key_len: i32,
55        params: &HashMap<String, Val>,
56    ) -> anyhow::Result<usize> {
57        match Self::get_key(caller, key_ptr, key_len)? {
58            Some(key) => match params.get(&key) {
59                Some(v) => Ok(v.to_bytes().len()),
60                None => Ok(0),
61            },
62            None => Ok(0),
63        }
64    }
65
66    fn get_param_value(
67        caller: &mut Caller<'_, ()>,
68        key_ptr: i32,
69        key_len: i32,
70        value_ptr: i32,
71        params: &HashMap<String, Val>,
72    ) -> anyhow::Result<bool> {
73        match Self::get_key(caller, key_ptr, key_len)? {
74            Some(key) => match params.get(&key) {
75                Some(v) => {
76                    let memory = Self::get_memory(caller)?;
77                    memory.write(caller, value_ptr as usize, &v.to_bytes())?;
78                    Ok(true)
79                }
80                None => Ok(false),
81            },
82            None => Ok(false),
83        }
84    }
85
86    fn set_output(
87        caller: &mut Caller<'_, ()>,
88        key_ptr: i32,
89        key_len: i32,
90        value_ptr: i32,
91        value_len: i32,
92        outputs: &mut HashMap<String, Val>,
93    ) -> anyhow::Result<bool> {
94        match Self::get_key(caller, key_ptr, key_len)? {
95            Some(key) => {
96                let memory = Self::get_memory(caller)?;
97                match memory
98                    .data(&caller)
99                    .get(value_ptr as usize..(value_ptr + value_len) as usize)
100                {
101                    Some(bytes) => {
102                        let value = Val::from_bytes(bytes).map_err(|e| anyhow::anyhow!(e))?;
103                        outputs.insert(key, value);
104                        Ok(true)
105                    }
106                    None => Ok(false),
107                }
108            }
109            None => Ok(false),
110        }
111    }
112
113    pub fn new(wasm: &[u8], config: &Config) -> anyhow::Result<Self> {
114        let params: Arc<Mutex<HashMap<String, Val>>> = Arc::new(Mutex::new(HashMap::new()));
115        let outputs: Arc<Mutex<HashMap<String, Val>>> = Arc::new(Mutex::new(HashMap::new()));
116        let engine = Engine::new(config)?;
117        let mut store = Store::new(&engine, ());
118        let module = Module::new(&engine, wasm)?;
119        let mut linker = Linker::new(&engine);
120        linker.func_wrap(
121            "env",
122            CONSOLE_LOG_METHOD_NAME,
123            |mut caller: Caller<'_, ()>, key_ptr: i32, key_len: i32| match Self::get_key(
124                &mut caller,
125                key_ptr,
126                key_len,
127            ) {
128                Ok(v) => match v {
129                    Some(v) => println!("{}:{}", CONSOLE_LOG_METHOD_NAME, v),
130                    None => {}
131                },
132                Err(err) => println!("{}:{}", CONSOLE_LOG_METHOD_NAME, err),
133            },
134        )?;
135        let get_param_len_params = params.clone();
136        linker.func_wrap(
137            "env",
138            GET_PARAM_LEN_METHOD_NAME,
139            move |mut caller: Caller<'_, ()>, key_ptr: i32, key_len: i32| {
140                let params = get_param_len_params.lock().unwrap();
141                match Self::get_param_len(&mut caller, key_ptr, key_len, &params) {
142                    Ok(v) => v as i32,
143                    Err(err) => {
144                        println!("{}:{}", GET_PARAM_LEN_METHOD_NAME, err);
145                        0
146                    }
147                }
148            },
149        )?;
150        let get_param_value_params = params.clone();
151        linker.func_wrap(
152            "env",
153            &GET_PARAM_VALUE_METHOD_NAME,
154            move |mut caller: Caller<'_, ()>, key_ptr: i32, key_len: i32, value_ptr: i32| {
155                let params = get_param_value_params.lock().unwrap();
156                match Self::get_param_value(&mut caller, key_ptr, key_len, value_ptr, &params) {
157                    Ok(v) => {
158                        if v {
159                            1
160                        } else {
161                            0
162                        }
163                    }
164                    Err(err) => {
165                        println!("{}:{}", GET_PARAM_LEN_METHOD_NAME, err);
166                        0
167                    }
168                }
169            },
170        )?;
171        let set_output_outputs = outputs.clone();
172        linker.func_wrap(
173            "env",
174            &SET_OUTPUT_METHOD_NAME,
175            move |mut caller: Caller<'_, ()>,
176                  key_ptr: i32,
177                  key_len: i32,
178                  value_ptr: i32,
179                  value_len: i32| {
180                let mut outputs = set_output_outputs.lock().unwrap();
181                match Self::set_output(
182                    &mut caller,
183                    key_ptr,
184                    key_len,
185                    value_ptr,
186                    value_len,
187                    &mut outputs,
188                ) {
189                    Ok(v) => {
190                        if v {
191                            1
192                        } else {
193                            0
194                        }
195                    }
196                    Err(err) => {
197                        println!("{}:{}", SET_OUTPUT_METHOD_NAME, err);
198                        0
199                    }
200                }
201            },
202        )?;
203        let instance = linker.instantiate(&mut store, &module)?;
204        let function = match instance.get_func(&mut store, MAIN_METHOD_NAME) {
205            Some(v) => v,
206            None => anyhow::bail!(format!("Can't find `{}` method", MAIN_METHOD_NAME)),
207        };
208
209        Ok(Self {
210            store,
211            function,
212            params,
213            outputs,
214        })
215    }
216
217    pub fn set_param(&mut self, key: &str, value: &Val) {
218        let mut params = self.params.lock().unwrap();
219        params.insert(key.to_owned(), value.clone());
220    }
221
222    pub fn reset_params(&mut self) {
223        let mut params = self.params.lock().unwrap();
224        params.clear();
225    }
226
227    pub fn evaluate(&mut self) -> anyhow::Result<HashMap<String, Val>> {
228        self.outputs.lock().unwrap().clear();
229        self.function.call(&mut self.store, &vec![], &mut [])?;
230        Ok(self.outputs.lock().unwrap().clone())
231    }
232}
clique_wasm_runtime/lib.rs

clique_wasm_runtime/
lib.rs
