vm/
lib.rs

1//使用 cranelift 作为后端 直接 jit 解释脚本
2mod binary;
3mod native;
4pub use native::{ANY, STD};
5
6mod fns;
7use anyhow::{Result, anyhow};
8pub use fns::{FnInfo, FnVariant};
9mod context;
10pub use context::BuildContext;
11
12mod rt;
13use cranelift::prelude::types;
14use dynamic::Type;
15pub use rt::JITRunTime;
16use smol_str::SmolStr;
17mod db_module;
18mod http_module;
19mod llm_module;
20mod root_module;
21
22use std::cell::RefCell;
23use std::sync::{Mutex, OnceLock, Weak};
24static PTR_TYPE: OnceLock<types::Type> = OnceLock::new();
25pub fn ptr_type() -> types::Type {
26    PTR_TYPE.get().cloned().unwrap()
27}
28
29pub fn get_type(ty: &Type) -> Result<types::Type> {
30    if ty.is_f64() {
31        Ok(types::F64)
32    } else if ty.is_f32() {
33        Ok(types::F32)
34    } else if ty.is_int() | ty.is_uint() {
35        match ty.width() {
36            1 => Ok(types::I8),
37            2 => Ok(types::I16),
38            4 => Ok(types::I32),
39            8 => Ok(types::I64),
40            _ => Err(anyhow!("非法类型 {:?}", ty)),
41        }
42    } else if let Type::Bool = ty {
43        Ok(types::I8)
44    } else {
45        Ok(ptr_type())
46    }
47}
48
49use compiler::Symbol;
50use cranelift::prelude::*;
51
52pub fn init_jit(mut jit: JITRunTime) -> Result<JITRunTime> {
53    jit.add_all()?;
54    Ok(jit)
55}
56
57use std::sync::Arc;
58unsafe impl Send for JITRunTime {}
59unsafe impl Sync for JITRunTime {}
60
61thread_local! {
62    static CURRENT_VM: RefCell<Option<Weak<Mutex<JITRunTime>>>> = const { RefCell::new(None) };
63}
64
65fn set_current_vm(vm: &Vm) {
66    CURRENT_VM.with(|current| {
67        *current.borrow_mut() = Some(Arc::downgrade(&vm.jit));
68    });
69}
70
71fn with_current_vm<T>(f: impl FnOnce(&Vm) -> Result<T>) -> Result<T> {
72    CURRENT_VM.with(|current| {
73        let jit = current.borrow().as_ref().and_then(Weak::upgrade).ok_or_else(|| anyhow!("当前线程没有 VM"))?;
74        let vm = Vm { jit };
75        f(&vm)
76    })
77}
78
79pub(crate) fn import_current(name: &str, path: &str) -> Result<()> {
80    with_current_vm(|vm| vm.import(name, path))
81}
82
83pub(crate) fn get_current_fn_ptr(name: &str, arg_tys: &[Type]) -> Result<(*const u8, Type)> {
84    with_current_vm(|vm| vm.get_fn_ptr(name, arg_tys))
85}
86
87fn add_method_field(jit: &mut JITRunTime, def: &str, method: &str, id: u32) -> Result<()> {
88    let def_id = jit.get_id(def)?;
89    if let Some((_, define)) = jit.compiler.symbols.get_symbol_mut(def_id) {
90        if let Symbol::Struct(Type::Struct { params, fields }, _) = define {
91            fields.push((method.into(), Type::Symbol { id, params: params.clone() }));
92        }
93    }
94    Ok(())
95}
96
97fn add_native_module_fns(jit: &mut JITRunTime, module: &str, fns: &[(&str, &[Type], Type, *const u8)]) -> Result<()> {
98    jit.add_module(module);
99    for (name, arg_tys, ret_ty, fn_ptr) in fns {
100        let full_name = format!("{}::{}", module, name);
101        jit.add_native_ptr(&full_name, name, arg_tys, ret_ty.clone(), *fn_ptr)?;
102    }
103    jit.pop_module();
104    Ok(())
105}
106
107impl JITRunTime {
108    pub fn add_module(&mut self, name: &str) {
109        self.compiler.symbols.add_module(name.into());
110    }
111
112    pub fn pop_module(&mut self) {
113        self.compiler.symbols.pop_module();
114    }
115
116    pub fn add_type(&mut self, name: &str, ty: Type, is_pub: bool) -> u32 {
117        self.compiler.add_symbol(name, Symbol::Struct(ty, is_pub))
118    }
119
120    pub fn add_empty_type(&mut self, name: &str) -> Result<u32> {
121        match self.get_id(name) {
122            Ok(id) => Ok(id),
123            Err(_) => Ok(self.add_type(name, Type::Struct { params: Vec::new(), fields: Vec::new() }, true)),
124        }
125    }
126
127    pub fn add_native_module_ptr(&mut self, module: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
128        self.add_module(module);
129        let full_name = format!("{}::{}", module, name);
130        let result = self.add_native_ptr(&full_name, name, arg_tys, ret_ty, fn_ptr);
131        self.pop_module();
132        result
133    }
134
135    pub fn add_native_method_ptr(&mut self, def: &str, method: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
136        self.add_empty_type(def)?;
137        let full_name = format!("{}::{}", def, method);
138        let id = self.add_native_ptr(&full_name, &full_name, arg_tys, ret_ty, fn_ptr)?;
139        add_method_field(self, def, method, id)?;
140        Ok(id)
141    }
142
143    pub fn add_std(&mut self) -> Result<()> {
144        self.add_module("std");
145        for (name, arg_tys, ret_ty, fn_ptr) in STD {
146            self.add_native_ptr(name, name, arg_tys, ret_ty, fn_ptr)?;
147        }
148        Ok(())
149    }
150
151    pub fn add_any(&mut self) -> Result<()> {
152        for (name, arg_tys, ret_ty, fn_ptr) in ANY {
153            let (_, method) = name.split_once("::").ok_or_else(|| anyhow!("非法 Any 方法名 {}", name))?;
154            self.add_native_method_ptr("Any", method, arg_tys, ret_ty, fn_ptr)?;
155        }
156        Ok(())
157    }
158
159    pub fn add_vec(&mut self) -> Result<()> {
160        self.add_empty_type("Vec")?;
161        let vec_def = Type::Symbol { id: self.get_id("Vec")?, params: Vec::new() };
162        self.add_inline("Vec::swap", vec![vec_def.clone(), Type::I64, Type::I64], Type::Void, |ctx: Option<&mut BuildContext>, args: Vec<Value>| {
163            if let Some(ctx) = ctx {
164                let width = ctx.builder.ins().iconst(types::I64, 4);
165                let offset_val = ctx.builder.ins().imul(args[1], width); // i * 4 i32大小四字节
166                let final_addr = ctx.builder.ins().iadd(args[0], offset_val); // base + (i*4)
167                let dest = ctx.builder.ins().imul(args[2], width);
168                let dest_addr = ctx.builder.ins().iadd(args[0], dest); // base + (i*4)
169                let dest_val = ctx.builder.ins().load(types::I32, MemFlags::trusted(), dest_addr, 0);
170                let v = ctx.builder.ins().load(types::I32, MemFlags::trusted(), final_addr, 0);
171                ctx.builder.ins().store(MemFlags::trusted(), v, dest_addr, 0);
172                ctx.builder.ins().store(MemFlags::trusted(), dest_val, final_addr, 0);
173            }
174            Err(anyhow!("无返回值"))
175        })?;
176
177        self.add_inline("Vec::get_idx", vec![vec_def.clone(), Type::I64], Type::I32, |ctx: Option<&mut BuildContext>, args: Vec<Value>| {
178            if let Some(ctx) = ctx {
179                let width = ctx.builder.ins().iconst(types::I64, 4);
180                let offset_val = ctx.builder.ins().imul(args[1], width); // i * 4 i32大小四字节
181                let final_addr = ctx.builder.ins().iadd(args[0], offset_val);
182                Ok((Some(ctx.builder.ins().load(types::I32, MemFlags::trusted(), final_addr, 0)), Type::I32))
183            } else {
184                Ok((None, Type::I32))
185            }
186        })?;
187        Ok(())
188    }
189
190    pub fn add_llm(&mut self) -> Result<()> {
191        add_native_module_fns(self, "llm", &llm_module::LLM_NATIVE)
192    }
193
194    pub fn add_root(&mut self) -> Result<()> {
195        add_native_module_fns(self, "root", &root_module::ROOT_NATIVE)
196    }
197
198    pub fn add_http(&mut self) -> Result<()> {
199        add_native_module_fns(self, "http", &http_module::HTTP_NATIVE)
200    }
201
202    pub fn add_db(&mut self) -> Result<()> {
203        add_native_module_fns(self, "db", &db_module::DB_NATIVE)
204    }
205
206    pub fn add_all(&mut self) -> Result<()> {
207        self.add_std()?;
208        self.add_any()?;
209        self.add_vec()?;
210        self.add_llm()?;
211        self.add_root()?;
212        self.add_http()?;
213        self.add_db()?;
214        Ok(())
215    }
216}
217
218#[derive(Clone)]
219pub struct Vm {
220    jit: Arc<Mutex<JITRunTime>>,
221}
222
223#[derive(Clone)]
224pub struct CompiledFn {
225    ptr: usize,
226    ret: Type,
227    owner: Vm,
228}
229
230impl CompiledFn {
231    pub fn ptr(&self) -> *const u8 {
232        set_current_vm(&self.owner);
233        self.ptr as *const u8
234    }
235
236    pub fn ret_ty(&self) -> &Type {
237        &self.ret
238    }
239
240    pub fn owner(&self) -> &Vm {
241        &self.owner
242    }
243}
244
245impl Vm {
246    pub fn new() -> Self {
247        Self { jit: Arc::new(Mutex::new(JITRunTime::new(|_| {}))) }
248    }
249
250    pub fn with_all() -> Result<Self> {
251        let vm = Self::new();
252        vm.add_all()?;
253        Ok(vm)
254    }
255
256    pub fn add_module(&self, name: &str) {
257        self.jit.lock().unwrap().add_module(name)
258    }
259
260    pub fn pop_module(&self) {
261        self.jit.lock().unwrap().pop_module()
262    }
263
264    pub fn add_type(&self, name: &str, ty: Type, is_pub: bool) -> u32 {
265        self.jit.lock().unwrap().add_type(name, ty, is_pub)
266    }
267
268    pub fn add_empty_type(&self, name: &str) -> Result<u32> {
269        self.jit.lock().unwrap().add_empty_type(name)
270    }
271
272    pub fn add_std(&self) -> Result<()> {
273        self.jit.lock().unwrap().add_std()
274    }
275
276    pub fn add_any(&self) -> Result<()> {
277        self.jit.lock().unwrap().add_any()
278    }
279
280    pub fn add_vec(&self) -> Result<()> {
281        self.jit.lock().unwrap().add_vec()
282    }
283
284    pub fn add_llm(&self) -> Result<()> {
285        self.jit.lock().unwrap().add_llm()
286    }
287
288    pub fn add_root(&self) -> Result<()> {
289        self.jit.lock().unwrap().add_root()
290    }
291
292    pub fn add_http(&self) -> Result<()> {
293        self.jit.lock().unwrap().add_http()
294    }
295
296    pub fn add_db(&self) -> Result<()> {
297        self.jit.lock().unwrap().add_db()
298    }
299
300    pub fn add_all(&self) -> Result<()> {
301        self.jit.lock().unwrap().add_all()
302    }
303
304    pub fn add_native_ptr(&self, full_name: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
305        self.jit.lock().unwrap().add_native_ptr(full_name, name, arg_tys, ret_ty, fn_ptr)
306    }
307
308    pub fn add_native_module_ptr(&self, module: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
309        self.jit.lock().unwrap().add_native_module_ptr(module, name, arg_tys, ret_ty, fn_ptr)
310    }
311
312    pub fn add_native_method_ptr(&self, def: &str, method: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
313        self.jit.lock().unwrap().add_native_method_ptr(def, method, arg_tys, ret_ty, fn_ptr)
314    }
315
316    pub fn add_inline(&self, name: &str, args: Vec<Type>, ret: Type, f: fn(Option<&mut BuildContext>, Vec<Value>) -> Result<(Option<Value>, Type)>) -> Result<u32> {
317        self.jit.lock().unwrap().add_inline(name, args, ret, f)
318    }
319
320    pub fn import_code(&self, name: &str, code: Vec<u8>) -> Result<()> {
321        self.jit.lock().unwrap().import_code(name, code)
322    }
323
324    pub fn import_file(&self, name: &str, path: &str) -> Result<()> {
325        self.jit.lock().unwrap().compiler.import_file(name, path)?;
326        Ok(())
327    }
328
329    pub fn import(&self, name: &str, path: &str) -> Result<()> {
330        if root::contains(path) {
331            let code = root::get(path).unwrap();
332            if code.is_str() {
333                self.import_code(name, code.as_str().as_bytes().to_vec())
334            } else {
335                self.import_code(name, code.get_dynamic("code").ok_or(anyhow!("{:?} 没有 code 成员", code))?.as_str().as_bytes().to_vec())
336            }
337        } else {
338            self.import_file(name, path)
339        }
340    }
341
342    pub fn infer(&self, name: &str, arg_tys: &[Type]) -> Result<Type> {
343        self.jit.lock().unwrap().get_type(name, arg_tys)
344    }
345
346    pub fn get_fn_ptr(&self, name: &str, arg_tys: &[Type]) -> Result<(*const u8, Type)> {
347        self.jit.lock().unwrap().get_fn_ptr(name, arg_tys)
348    }
349
350    pub fn get_fn(&self, name: &str, arg_tys: &[Type]) -> Result<CompiledFn> {
351        set_current_vm(self);
352        let (ptr, ret) = self.get_fn_ptr(name, arg_tys)?;
353        Ok(CompiledFn { ptr: ptr as usize, ret, owner: self.clone() })
354    }
355
356    pub fn load(&self, code: Vec<u8>, arg_name: SmolStr) -> Result<(i64, Type)> {
357        self.jit.lock().unwrap().load(code, arg_name)
358    }
359
360    pub fn get_symbol(&self, name: &str, params: Vec<Type>) -> Result<Type> {
361        Ok(Type::Symbol { id: self.jit.lock().unwrap().get_id(name)?, params })
362    }
363
364    pub fn disassemble(&self, name: &str) -> Result<String> {
365        self.jit.lock().unwrap().compiler.symbols.disassemble(name)
366    }
367
368    #[cfg(feature = "ir-disassembly")]
369    pub fn disassemble_ir(&self, name: &str) -> Result<String> {
370        self.jit.lock().unwrap().disassemble_ir(name)
371    }
372}
373
374impl Default for Vm {
375    fn default() -> Self {
376        Self::new()
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::Vm;
383    use dynamic::{Dynamic, ToJson, Type};
384
385    extern "C" fn math_double(value: i64) -> i64 {
386        value * 2
387    }
388
389    #[test]
390    fn vm_can_add_native_after_jit_creation() -> anyhow::Result<()> {
391        let vm = Vm::new();
392        vm.add_native_module_ptr("math", "double", &[Type::I64], Type::I64, math_double as *const u8)?;
393        vm.import_code(
394            "vm_dynamic_native",
395            br#"
396            pub fn run(value: i64) {
397                math::double(value)
398            }
399            "#
400            .to_vec(),
401        )?;
402
403        let compiled = vm.get_fn("vm_dynamic_native::run", &[Type::I64])?;
404        assert_eq!(compiled.ret_ty(), &Type::I64);
405        let run: extern "C" fn(i64) -> i64 = unsafe { std::mem::transmute(compiled.ptr()) };
406        assert_eq!(run(21), 42);
407        Ok(())
408    }
409
410    #[test]
411    fn compares_any_with_string_literal_as_string() -> anyhow::Result<()> {
412        let vm = Vm::with_all()?;
413        vm.import_code(
414            "vm_string_compare_any",
415            br#"
416            pub fn any_ne_empty(chat_path) {
417                chat_path != ""
418            }
419            "#
420            .to_vec(),
421        )?;
422
423        let compiled = vm.get_fn("vm_string_compare_any::any_ne_empty", &[Type::Any])?;
424        assert_eq!(compiled.ret_ty(), &Type::Bool);
425
426        let any_ne_empty: extern "C" fn(*const Dynamic) -> bool = unsafe { std::mem::transmute(compiled.ptr()) };
427        let empty = Dynamic::from("");
428        let non_empty = Dynamic::from("chat");
429
430        assert!(!any_ne_empty(&empty));
431        assert!(any_ne_empty(&non_empty));
432        Ok(())
433    }
434
435    #[test]
436    fn compares_concrete_value_with_string_literal_as_string() -> anyhow::Result<()> {
437        let vm = Vm::with_all()?;
438        vm.import_code(
439            "vm_string_compare_imm",
440            br#"
441            pub fn int_eq_str(value: i64) {
442                value == "42"
443            }
444
445            pub fn int_to_str(value: i64) {
446                value + ""
447            }
448            "#
449            .to_vec(),
450        )?;
451
452        let compiled = vm.get_fn("vm_string_compare_imm::int_eq_str", &[Type::I64])?;
453        assert_eq!(compiled.ret_ty(), &Type::Bool);
454
455        let int_eq_str: extern "C" fn(i64) -> bool = unsafe { std::mem::transmute(compiled.ptr()) };
456
457        let compiled = vm.get_fn("vm_string_compare_imm::int_to_str", &[Type::I64])?;
458        assert_eq!(compiled.ret_ty(), &Type::Any);
459        let int_to_str: extern "C" fn(i64) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
460        let text = int_to_str(42);
461        assert_eq!(unsafe { &*text }.as_str(), "42");
462
463        assert!(int_eq_str(42));
464        assert!(!int_eq_str(7));
465        Ok(())
466    }
467
468    #[test]
469    fn dynamic_field_value_participates_in_or_expression() -> anyhow::Result<()> {
470        let vm = Vm::with_all()?;
471        vm.import_code(
472            "vm_dynamic_field_or",
473            r#"
474            pub fn next_or_start() {
475                let choice = {
476                    label: "颜色",
477                    next: "color"
478                };
479                choice.next || "start"
480            }
481
482            pub fn direct_next() {
483                let choice = {
484                    label: "颜色",
485                    next: "color"
486                };
487                choice.next
488            }
489
490            pub fn bracket_next() {
491                let choice = {
492                    label: "颜色",
493                    next: "color"
494                };
495                choice["next"]
496            }
497
498            pub fn assigned_preview() {
499                let choice = {
500                    next: "tax_free"
501                };
502                choice.preview = choice.next || "start";
503                choice
504            }
505            "#
506            .as_bytes()
507            .to_vec(),
508        )?;
509
510        let compiled = vm.get_fn("vm_dynamic_field_or::direct_next", &[])?;
511        assert_eq!(compiled.ret_ty(), &Type::Any);
512        let direct_next: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
513        assert_eq!(unsafe { &*direct_next() }.as_str(), "color");
514
515        let compiled = vm.get_fn("vm_dynamic_field_or::bracket_next", &[])?;
516        assert_eq!(compiled.ret_ty(), &Type::Any);
517        let bracket_next: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
518        assert_eq!(unsafe { &*bracket_next() }.as_str(), "color");
519
520        let compiled = vm.get_fn("vm_dynamic_field_or::next_or_start", &[])?;
521        assert_eq!(compiled.ret_ty(), &Type::Any);
522        let next_or_start: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
523        assert_eq!(unsafe { &*next_or_start() }.as_str(), "color");
524
525        let compiled = vm.get_fn("vm_dynamic_field_or::assigned_preview", &[])?;
526        assert_eq!(compiled.ret_ty(), &Type::Any);
527        let assigned_preview: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
528        let choice = unsafe { &*assigned_preview() };
529        assert_eq!(choice.get_dynamic("preview").unwrap().as_str(), "tax_free");
530        Ok(())
531    }
532
533    #[test]
534    fn empty_object_literal_in_if_branch_stays_dynamic() -> anyhow::Result<()> {
535        let vm = Vm::with_all()?;
536        vm.import_code(
537            "vm_if_empty_object_branch",
538            r#"
539            pub fn first_note(steps) {
540                let first = if steps.len() > 0 { steps[0] } else { {} };
541                let first_note = first.note || "fallback";
542                first_note
543            }
544
545            pub fn first_ja(steps) {
546                let first = if steps.len() > 0 { steps[0] } else { {} };
547                first.ja || "すみません"
548            }
549            "#
550            .as_bytes()
551            .to_vec(),
552        )?;
553
554        let compiled = vm.get_fn("vm_if_empty_object_branch::first_note", &[Type::Any])?;
555        assert_eq!(compiled.ret_ty(), &Type::Any);
556        let first_note: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
557
558        let empty_steps = Dynamic::list(Vec::new());
559        assert_eq!(unsafe { &*first_note(&empty_steps) }.as_str(), "fallback");
560
561        let mut step = std::collections::BTreeMap::new();
562        step.insert("note".into(), "hello".into());
563        let steps = Dynamic::list(vec![Dynamic::map(step)]);
564        assert_eq!(unsafe { &*first_note(&steps) }.as_str(), "hello");
565
566        let compiled = vm.get_fn("vm_if_empty_object_branch::first_ja", &[Type::Any])?;
567        assert_eq!(compiled.ret_ty(), &Type::Any);
568        let first_ja: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
569        assert_eq!(unsafe { &*first_ja(&empty_steps) }.as_str(), "すみません");
570        Ok(())
571    }
572
573    #[test]
574    fn root_native_calls_do_not_take_ownership_of_dynamic_args() -> anyhow::Result<()> {
575        let vm = Vm::with_all()?;
576        vm.import_code(
577            "vm_root_clone_bridge",
578            br#"
579            pub fn add_then_reuse(arg) {
580                let user = {
581                    address: "test-wallet",
582                    points: 20
583                };
584                root::add("local/root-clone-bridge-user", user);
585                user.points = user.points - 7;
586                root::add("local/root-clone-bridge-user", user);
587                {
588                    user: user,
589                    points: user.points
590                }
591            }
592            "#
593            .to_vec(),
594        )?;
595
596        let compiled = vm.get_fn("vm_root_clone_bridge::add_then_reuse", &[Type::Any])?;
597        assert_eq!(compiled.ret_ty(), &Type::Any);
598        let add_then_reuse: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
599        let arg = Dynamic::Null;
600        let result = add_then_reuse(&arg);
601        let result = unsafe { &*result };
602
603        assert_eq!(result.get_dynamic("points").and_then(|value| value.as_int()), Some(13));
604        let mut json = String::new();
605        result.to_json(&mut json);
606        assert!(json.contains("\"points\": 13"));
607        Ok(())
608    }
609}
vm/lib.rs

vm/
lib.rs
