Skip to main content

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 gpu_layout;
19mod gpu_module;
20mod http_module;
21mod llm_module;
22mod root_module;
23pub use gpu_layout::{GpuFieldLayout, GpuStructLayout};
24
25use std::sync::{Mutex, OnceLock, Weak};
26static PTR_TYPE: OnceLock<types::Type> = OnceLock::new();
27pub fn ptr_type() -> types::Type {
28    PTR_TYPE.get().cloned().unwrap()
29}
30
31pub fn get_type(ty: &Type) -> Result<types::Type> {
32    if ty.is_f64() {
33        Ok(types::F64)
34    } else if ty.is_f32() {
35        Ok(types::F32)
36    } else if ty.is_int() | ty.is_uint() {
37        match ty.width() {
38            1 => Ok(types::I8),
39            2 => Ok(types::I16),
40            4 => Ok(types::I32),
41            8 => Ok(types::I64),
42            _ => Err(anyhow!("非法类型 {:?}", ty)),
43        }
44    } else if let Type::Bool = ty {
45        Ok(types::I8)
46    } else {
47        Ok(ptr_type())
48    }
49}
50
51use compiler::Symbol;
52use cranelift::prelude::*;
53
54pub fn init_jit(mut jit: JITRunTime) -> Result<JITRunTime> {
55    jit.add_all()?;
56    Ok(jit)
57}
58
59use std::sync::Arc;
60unsafe impl Send for JITRunTime {}
61unsafe impl Sync for JITRunTime {}
62
63pub(crate) fn with_vm_context<T>(context: *const Weak<Mutex<JITRunTime>>, f: impl FnOnce(&Vm) -> Result<T>) -> Result<T> {
64    if context.is_null() {
65        return Err(anyhow!("VM context is null"));
66    }
67    let jit = unsafe { &*context }.upgrade().ok_or_else(|| anyhow!("VM context has expired"))?;
68    let vm = Vm { jit };
69    f(&vm)
70}
71
72fn add_method_field(jit: &mut JITRunTime, def: &str, method: &str, id: u32) -> Result<()> {
73    let def_id = jit.get_id(def)?;
74    if let Some((_, define)) = jit.compiler.symbols.get_symbol_mut(def_id) {
75        if let Symbol::Struct(Type::Struct { params, fields }, _) = define {
76            fields.push((method.into(), Type::Symbol { id, params: params.clone() }));
77        }
78    }
79    Ok(())
80}
81
82fn add_native_module_fns(jit: &mut JITRunTime, module: &str, fns: &[(&str, &[Type], Type, *const u8)]) -> Result<()> {
83    jit.add_module(module);
84    for (name, arg_tys, ret_ty, fn_ptr) in fns {
85        let full_name = format!("{}::{}", module, name);
86        jit.add_native_ptr(&full_name, name, arg_tys, ret_ty.clone(), *fn_ptr)?;
87    }
88    jit.pop_module();
89    Ok(())
90}
91
92impl JITRunTime {
93    pub fn add_module(&mut self, name: &str) {
94        self.compiler.symbols.add_module(name.into());
95    }
96
97    pub fn pop_module(&mut self) {
98        self.compiler.symbols.pop_module();
99    }
100
101    pub fn add_type(&mut self, name: &str, ty: Type, is_pub: bool) -> u32 {
102        self.compiler.add_symbol(name, Symbol::Struct(ty, is_pub))
103    }
104
105    pub fn add_empty_type(&mut self, name: &str) -> Result<u32> {
106        match self.get_id(name) {
107            Ok(id) => Ok(id),
108            Err(_) => Ok(self.add_type(name, Type::Struct { params: Vec::new(), fields: Vec::new() }, true)),
109        }
110    }
111
112    pub fn add_native_module_ptr(&mut self, module: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
113        self.add_module(module);
114        let full_name = format!("{}::{}", module, name);
115        let result = self.add_native_ptr(&full_name, name, arg_tys, ret_ty, fn_ptr);
116        self.pop_module();
117        result
118    }
119
120    pub(crate) fn add_native_module_context_ptr(&mut self, module: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
121        self.add_module(module);
122        let full_name = format!("{}::{}", module, name);
123        let result = self.add_context_native_ptr(&full_name, name, arg_tys, ret_ty, fn_ptr);
124        self.pop_module();
125        result
126    }
127
128    pub fn add_native_method_ptr(&mut self, def: &str, method: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
129        self.add_empty_type(def)?;
130        let full_name = format!("{}::{}", def, method);
131        let id = self.add_native_ptr(&full_name, &full_name, arg_tys, ret_ty, fn_ptr)?;
132        add_method_field(self, def, method, id)?;
133        Ok(id)
134    }
135
136    pub fn add_std(&mut self) -> Result<()> {
137        self.add_module("std");
138        for (name, arg_tys, ret_ty, fn_ptr) in STD {
139            self.add_native_ptr(name, name, arg_tys, ret_ty, fn_ptr)?;
140        }
141        self.add_context_native_ptr("import", "import", &[Type::Any, Type::Any], Type::Bool, native::import_with_vm as *const u8)?;
142        Ok(())
143    }
144
145    pub fn add_any(&mut self) -> Result<()> {
146        for (name, arg_tys, ret_ty, fn_ptr) in ANY {
147            let (_, method) = name.split_once("::").ok_or_else(|| anyhow!("非法 Any 方法名 {}", name))?;
148            self.add_native_method_ptr("Any", method, arg_tys, ret_ty, fn_ptr)?;
149        }
150        Ok(())
151    }
152
153    pub fn add_vec(&mut self) -> Result<()> {
154        self.add_empty_type("Vec")?;
155        let vec_def = Type::Symbol { id: self.get_id("Vec")?, params: Vec::new() };
156        self.add_inline("Vec::swap", vec![vec_def.clone(), Type::I64, Type::I64], Type::Void, |ctx: Option<&mut BuildContext>, args: Vec<Value>| {
157            if let Some(ctx) = ctx {
158                let width = ctx.builder.ins().iconst(types::I64, 4);
159                let offset_val = ctx.builder.ins().imul(args[1], width); // i * 4 i32大小四字节
160                let final_addr = ctx.builder.ins().iadd(args[0], offset_val); // base + (i*4)
161                let dest = ctx.builder.ins().imul(args[2], width);
162                let dest_addr = ctx.builder.ins().iadd(args[0], dest); // base + (i*4)
163                let dest_val = ctx.builder.ins().load(types::I32, MemFlags::trusted(), dest_addr, 0);
164                let v = ctx.builder.ins().load(types::I32, MemFlags::trusted(), final_addr, 0);
165                ctx.builder.ins().store(MemFlags::trusted(), v, dest_addr, 0);
166                ctx.builder.ins().store(MemFlags::trusted(), dest_val, final_addr, 0);
167            }
168            Err(anyhow!("无返回值"))
169        })?;
170
171        self.add_inline("Vec::get_idx", vec![vec_def.clone(), Type::I64], Type::I32, |ctx: Option<&mut BuildContext>, args: Vec<Value>| {
172            if let Some(ctx) = ctx {
173                let width = ctx.builder.ins().iconst(types::I64, 4);
174                let offset_val = ctx.builder.ins().imul(args[1], width); // i * 4 i32大小四字节
175                let final_addr = ctx.builder.ins().iadd(args[0], offset_val);
176                Ok((Some(ctx.builder.ins().load(types::I32, MemFlags::trusted(), final_addr, 0)), Type::I32))
177            } else {
178                Ok((None, Type::I32))
179            }
180        })?;
181        Ok(())
182    }
183
184    pub fn add_llm(&mut self) -> Result<()> {
185        add_native_module_fns(self, "llm", &llm_module::LLM_NATIVE)
186    }
187
188    pub fn add_root(&mut self) -> Result<()> {
189        add_native_module_fns(self, "root", &root_module::ROOT_NATIVE)?;
190        self.add_native_module_context_ptr("root", "add_fn", &[Type::Any, Type::Any], Type::Bool, root_module::root_add_fn_with_vm as *const u8)?;
191        Ok(())
192    }
193
194    pub fn add_http(&mut self) -> Result<()> {
195        add_native_module_fns(self, "http", &http_module::HTTP_NATIVE)
196    }
197
198    pub fn add_db(&mut self) -> Result<()> {
199        add_native_module_fns(self, "db", &db_module::DB_NATIVE)
200    }
201
202    pub fn add_gpu(&mut self) -> Result<()> {
203        add_native_module_fns(self, "gpu", &gpu_module::GPU_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        self.add_gpu()?;
215        Ok(())
216    }
217}
218
219#[derive(Clone)]
220pub struct Vm {
221    jit: Arc<Mutex<JITRunTime>>,
222}
223
224#[derive(Clone)]
225pub struct CompiledFn {
226    ptr: usize,
227    ret: Type,
228    owner: Vm,
229}
230
231impl CompiledFn {
232    pub fn ptr(&self) -> *const u8 {
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        let jit = Arc::new(Mutex::new(JITRunTime::new(|_| {})));
248        jit.lock().unwrap().set_owner(Arc::downgrade(&jit));
249        Self { jit }
250    }
251
252    pub fn with_all() -> Result<Self> {
253        let vm = Self::new();
254        vm.add_all()?;
255        Ok(vm)
256    }
257
258    pub fn add_module(&self, name: &str) {
259        self.jit.lock().unwrap().add_module(name)
260    }
261
262    pub fn pop_module(&self) {
263        self.jit.lock().unwrap().pop_module()
264    }
265
266    pub fn add_type(&self, name: &str, ty: Type, is_pub: bool) -> u32 {
267        self.jit.lock().unwrap().add_type(name, ty, is_pub)
268    }
269
270    pub fn add_empty_type(&self, name: &str) -> Result<u32> {
271        self.jit.lock().unwrap().add_empty_type(name)
272    }
273
274    pub fn add_std(&self) -> Result<()> {
275        self.jit.lock().unwrap().add_std()
276    }
277
278    pub fn add_any(&self) -> Result<()> {
279        self.jit.lock().unwrap().add_any()
280    }
281
282    pub fn add_vec(&self) -> Result<()> {
283        self.jit.lock().unwrap().add_vec()
284    }
285
286    pub fn add_llm(&self) -> Result<()> {
287        self.jit.lock().unwrap().add_llm()
288    }
289
290    pub fn add_root(&self) -> Result<()> {
291        self.jit.lock().unwrap().add_root()
292    }
293
294    pub fn add_http(&self) -> Result<()> {
295        self.jit.lock().unwrap().add_http()
296    }
297
298    pub fn add_db(&self) -> Result<()> {
299        self.jit.lock().unwrap().add_db()
300    }
301
302    pub fn add_gpu(&self) -> Result<()> {
303        self.jit.lock().unwrap().add_gpu()
304    }
305
306    pub fn add_all(&self) -> Result<()> {
307        self.jit.lock().unwrap().add_all()
308    }
309
310    pub fn add_native_ptr(&self, full_name: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
311        self.jit.lock().unwrap().add_native_ptr(full_name, name, arg_tys, ret_ty, fn_ptr)
312    }
313
314    pub fn add_native_module_ptr(&self, module: &str, name: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
315        self.jit.lock().unwrap().add_native_module_ptr(module, name, arg_tys, ret_ty, fn_ptr)
316    }
317
318    pub fn add_native_method_ptr(&self, def: &str, method: &str, arg_tys: &[Type], ret_ty: Type, fn_ptr: *const u8) -> Result<u32> {
319        self.jit.lock().unwrap().add_native_method_ptr(def, method, arg_tys, ret_ty, fn_ptr)
320    }
321
322    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> {
323        self.jit.lock().unwrap().add_inline(name, args, ret, f)
324    }
325
326    pub fn import_code(&self, name: &str, code: Vec<u8>) -> Result<()> {
327        self.jit.lock().unwrap().import_code(name, code)
328    }
329
330    pub fn import_file(&self, name: &str, path: &str) -> Result<()> {
331        self.jit.lock().unwrap().compiler.import_file(name, path)?;
332        Ok(())
333    }
334
335    pub fn import(&self, name: &str, path: &str) -> Result<()> {
336        if root::contains(path) {
337            let code = root::get(path).unwrap();
338            if code.is_str() {
339                self.import_code(name, code.as_str().as_bytes().to_vec())
340            } else {
341                self.import_code(name, code.get_dynamic("code").ok_or(anyhow!("{:?} 没有 code 成员", code))?.as_str().as_bytes().to_vec())
342            }
343        } else {
344            self.import_file(name, path)
345        }
346    }
347
348    pub fn infer(&self, name: &str, arg_tys: &[Type]) -> Result<Type> {
349        self.jit.lock().unwrap().get_type(name, arg_tys)
350    }
351
352    pub fn get_fn_ptr(&self, name: &str, arg_tys: &[Type]) -> Result<(*const u8, Type)> {
353        self.jit.lock().unwrap().get_fn_ptr(name, arg_tys)
354    }
355
356    pub fn get_fn(&self, name: &str, arg_tys: &[Type]) -> Result<CompiledFn> {
357        let (ptr, ret) = self.get_fn_ptr(name, arg_tys)?;
358        Ok(CompiledFn { ptr: ptr as usize, ret, owner: self.clone() })
359    }
360
361    pub fn load(&self, code: Vec<u8>, arg_name: SmolStr) -> Result<(i64, Type)> {
362        self.jit.lock().unwrap().load(code, arg_name)
363    }
364
365    pub fn get_symbol(&self, name: &str, params: Vec<Type>) -> Result<Type> {
366        Ok(Type::Symbol { id: self.jit.lock().unwrap().get_id(name)?, params })
367    }
368
369    pub fn gpu_struct_layout(&self, name: &str, params: &[Type]) -> Result<GpuStructLayout> {
370        let jit = self.jit.lock().unwrap();
371        GpuStructLayout::from_symbol_table(&jit.compiler.symbols, name, params)
372    }
373
374    pub fn disassemble(&self, name: &str) -> Result<String> {
375        self.jit.lock().unwrap().compiler.symbols.disassemble(name)
376    }
377
378    #[cfg(feature = "ir-disassembly")]
379    pub fn disassemble_ir(&self, name: &str) -> Result<String> {
380        self.jit.lock().unwrap().disassemble_ir(name)
381    }
382}
383
384impl Default for Vm {
385    fn default() -> Self {
386        Self::new()
387    }
388}
389
390#[cfg(test)]
391mod tests {
392    use super::Vm;
393    use dynamic::{Dynamic, ToJson, Type};
394
395    extern "C" fn math_double(value: i64) -> i64 {
396        value * 2
397    }
398
399    #[test]
400    fn vm_can_add_native_after_jit_creation() -> anyhow::Result<()> {
401        let vm = Vm::new();
402        vm.add_native_module_ptr("math", "double", &[Type::I64], Type::I64, math_double as *const u8)?;
403        vm.import_code(
404            "vm_dynamic_native",
405            br#"
406            pub fn run(value: i64) {
407                math::double(value)
408            }
409            "#
410            .to_vec(),
411        )?;
412
413        let compiled = vm.get_fn("vm_dynamic_native::run", &[Type::I64])?;
414        assert_eq!(compiled.ret_ty(), &Type::I64);
415        let run: extern "C" fn(i64) -> i64 = unsafe { std::mem::transmute(compiled.ptr()) };
416        assert_eq!(run(21), 42);
417        Ok(())
418    }
419
420    #[test]
421    fn compares_any_with_string_literal_as_string() -> anyhow::Result<()> {
422        let vm = Vm::with_all()?;
423        vm.import_code(
424            "vm_string_compare_any",
425            br#"
426            pub fn any_ne_empty(chat_path) {
427                chat_path != ""
428            }
429            "#
430            .to_vec(),
431        )?;
432
433        let compiled = vm.get_fn("vm_string_compare_any::any_ne_empty", &[Type::Any])?;
434        assert_eq!(compiled.ret_ty(), &Type::Bool);
435
436        let any_ne_empty: extern "C" fn(*const Dynamic) -> bool = unsafe { std::mem::transmute(compiled.ptr()) };
437        let empty = Dynamic::from("");
438        let non_empty = Dynamic::from("chat");
439
440        assert!(!any_ne_empty(&empty));
441        assert!(any_ne_empty(&non_empty));
442        Ok(())
443    }
444
445    #[test]
446    fn compares_concrete_value_with_string_literal_as_string() -> anyhow::Result<()> {
447        let vm = Vm::with_all()?;
448        vm.import_code(
449            "vm_string_compare_imm",
450            br#"
451            pub fn int_eq_str(value: i64) {
452                value == "42"
453            }
454
455            pub fn int_to_str(value: i64) {
456                value + ""
457            }
458            "#
459            .to_vec(),
460        )?;
461
462        let compiled = vm.get_fn("vm_string_compare_imm::int_eq_str", &[Type::I64])?;
463        assert_eq!(compiled.ret_ty(), &Type::Bool);
464
465        let int_eq_str: extern "C" fn(i64) -> bool = unsafe { std::mem::transmute(compiled.ptr()) };
466
467        let compiled = vm.get_fn("vm_string_compare_imm::int_to_str", &[Type::I64])?;
468        assert_eq!(compiled.ret_ty(), &Type::Any);
469        let int_to_str: extern "C" fn(i64) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
470        let text = int_to_str(42);
471        assert_eq!(unsafe { &*text }.as_str(), "42");
472
473        assert!(int_eq_str(42));
474        assert!(!int_eq_str(7));
475        Ok(())
476    }
477
478    #[test]
479    fn dynamic_field_value_participates_in_or_expression() -> anyhow::Result<()> {
480        let vm = Vm::with_all()?;
481        vm.import_code(
482            "vm_dynamic_field_or",
483            r#"
484            pub fn next_or_start() {
485                let choice = {
486                    label: "颜色",
487                    next: "color"
488                };
489                choice.next || "start"
490            }
491
492            pub fn direct_next() {
493                let choice = {
494                    label: "颜色",
495                    next: "color"
496                };
497                choice.next
498            }
499
500            pub fn bracket_next() {
501                let choice = {
502                    label: "颜色",
503                    next: "color"
504                };
505                choice["next"]
506            }
507
508            pub fn assigned_preview() {
509                let choice = {
510                    next: "tax_free"
511                };
512                choice.preview = choice.next || "start";
513                choice
514            }
515            "#
516            .as_bytes()
517            .to_vec(),
518        )?;
519
520        let compiled = vm.get_fn("vm_dynamic_field_or::direct_next", &[])?;
521        assert_eq!(compiled.ret_ty(), &Type::Any);
522        let direct_next: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
523        assert_eq!(unsafe { &*direct_next() }.as_str(), "color");
524
525        let compiled = vm.get_fn("vm_dynamic_field_or::bracket_next", &[])?;
526        assert_eq!(compiled.ret_ty(), &Type::Any);
527        let bracket_next: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
528        assert_eq!(unsafe { &*bracket_next() }.as_str(), "color");
529
530        let compiled = vm.get_fn("vm_dynamic_field_or::next_or_start", &[])?;
531        assert_eq!(compiled.ret_ty(), &Type::Any);
532        let next_or_start: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
533        assert_eq!(unsafe { &*next_or_start() }.as_str(), "color");
534
535        let compiled = vm.get_fn("vm_dynamic_field_or::assigned_preview", &[])?;
536        assert_eq!(compiled.ret_ty(), &Type::Any);
537        let assigned_preview: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
538        let choice = unsafe { &*assigned_preview() };
539        assert_eq!(choice.get_dynamic("preview").unwrap().as_str(), "tax_free");
540        Ok(())
541    }
542
543    #[test]
544    fn empty_object_literal_in_if_branch_stays_dynamic() -> anyhow::Result<()> {
545        let vm = Vm::with_all()?;
546        vm.import_code(
547            "vm_if_empty_object_branch",
548            r#"
549            pub fn first_note(steps) {
550                let first = if steps.len() > 0 { steps[0] } else { {} };
551                let first_note = first.note || "fallback";
552                first_note
553            }
554
555            pub fn first_ja(steps) {
556                let first = if steps.len() > 0 { steps[0] } else { {} };
557                first.ja || "すみません"
558            }
559
560            pub fn assign_first_note(steps) {
561                let first = {};
562                first = if steps.len() > 0 { steps[0] } else { {} };
563                first.note || "fallback"
564            }
565            "#
566            .as_bytes()
567            .to_vec(),
568        )?;
569
570        let compiled = vm.get_fn("vm_if_empty_object_branch::first_note", &[Type::Any])?;
571        assert_eq!(compiled.ret_ty(), &Type::Any);
572        let first_note: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
573
574        let empty_steps = Dynamic::list(Vec::new());
575        assert_eq!(unsafe { &*first_note(&empty_steps) }.as_str(), "fallback");
576
577        let mut step = std::collections::BTreeMap::new();
578        step.insert("note".into(), "hello".into());
579        let steps = Dynamic::list(vec![Dynamic::map(step)]);
580        assert_eq!(unsafe { &*first_note(&steps) }.as_str(), "hello");
581
582        let compiled = vm.get_fn("vm_if_empty_object_branch::first_ja", &[Type::Any])?;
583        assert_eq!(compiled.ret_ty(), &Type::Any);
584        let first_ja: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
585        assert_eq!(unsafe { &*first_ja(&empty_steps) }.as_str(), "すみません");
586
587        let compiled = vm.get_fn("vm_if_empty_object_branch::assign_first_note", &[Type::Any])?;
588        assert_eq!(compiled.ret_ty(), &Type::Any);
589        let assign_first_note: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
590        assert_eq!(unsafe { &*assign_first_note(&empty_steps) }.as_str(), "fallback");
591        assert_eq!(unsafe { &*assign_first_note(&steps) }.as_str(), "hello");
592        Ok(())
593    }
594
595    #[test]
596    fn list_literal_can_be_function_tail_expression() -> anyhow::Result<()> {
597        let vm = Vm::with_all()?;
598        vm.import_code(
599            "vm_tail_list_literal",
600            r#"
601            pub fn numbers() {
602                [1, 2, 3]
603            }
604
605            pub fn maps() {
606                [
607                    {note: "first"},
608                    {note: "second"}
609                ]
610            }
611
612            pub fn object_with_maps() {
613                {
614                    steps: [
615                        {note: "first"},
616                        {note: "second"}
617                    ]
618                }
619            }
620
621            pub fn return_maps() {
622                return [
623                    {note: "first"},
624                    {note: "second"}
625                ];
626            }
627
628            pub fn return_maps_without_semicolon() {
629                return [
630                    {note: "first"},
631                    {note: "second"}
632                ]
633            }
634
635            pub fn tail_bare_variable() {
636                let value = [
637                    {note: "first"},
638                    {note: "second"}
639                ];
640                value
641            }
642
643            pub fn return_bare_variable_without_semicolon() {
644                let value = [
645                    {note: "first"},
646                    {note: "second"}
647                ];
648                return value
649            }
650
651            pub fn tail_object_variable() {
652                let result = {
653                    steps: [
654                        {note: "first"},
655                        {note: "second"}
656                    ]
657                };
658                result
659            }
660            "#
661            .as_bytes()
662            .to_vec(),
663        )?;
664
665        let compiled = vm.get_fn("vm_tail_list_literal::numbers", &[])?;
666        assert_eq!(compiled.ret_ty(), &Type::Any);
667        let numbers: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
668        let result = unsafe { &*numbers() };
669        assert_eq!(result.len(), 3);
670        assert_eq!(result.get_idx(1).and_then(|value| value.as_int()), Some(2));
671
672        let compiled = vm.get_fn("vm_tail_list_literal::maps", &[])?;
673        assert_eq!(compiled.ret_ty(), &Type::Any);
674        let maps: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
675        let result = unsafe { &*maps() };
676        assert_eq!(result.len(), 2);
677        assert_eq!(result.get_idx(1).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("second".to_string()));
678
679        let compiled = vm.get_fn("vm_tail_list_literal::object_with_maps", &[])?;
680        assert_eq!(compiled.ret_ty(), &Type::Any);
681        let object_with_maps: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
682        let result = unsafe { &*object_with_maps() };
683        let steps = result.get_dynamic("steps").expect("steps");
684        assert_eq!(steps.len(), 2);
685        assert_eq!(steps.get_idx(0).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("first".to_string()));
686
687        let compiled = vm.get_fn("vm_tail_list_literal::return_maps", &[])?;
688        assert_eq!(compiled.ret_ty(), &Type::Any);
689        let return_maps: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
690        let result = unsafe { &*return_maps() };
691        assert_eq!(result.len(), 2);
692        assert_eq!(result.get_idx(1).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("second".to_string()));
693
694        let compiled = vm.get_fn("vm_tail_list_literal::return_maps_without_semicolon", &[])?;
695        assert_eq!(compiled.ret_ty(), &Type::Any);
696        let return_maps_without_semicolon: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
697        let result = unsafe { &*return_maps_without_semicolon() };
698        assert_eq!(result.len(), 2);
699        assert_eq!(result.get_idx(0).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("first".to_string()));
700
701        let compiled = vm.get_fn("vm_tail_list_literal::tail_bare_variable", &[])?;
702        assert_eq!(compiled.ret_ty(), &Type::Any);
703        let tail_bare_variable: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
704        let result = unsafe { &*tail_bare_variable() };
705        assert_eq!(result.len(), 2);
706        assert_eq!(result.get_idx(1).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("second".to_string()));
707
708        let compiled = vm.get_fn("vm_tail_list_literal::return_bare_variable_without_semicolon", &[])?;
709        assert_eq!(compiled.ret_ty(), &Type::Any);
710        let return_bare_variable_without_semicolon: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
711        let result = unsafe { &*return_bare_variable_without_semicolon() };
712        assert_eq!(result.len(), 2);
713        assert_eq!(result.get_idx(0).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("first".to_string()));
714
715        let compiled = vm.get_fn("vm_tail_list_literal::tail_object_variable", &[])?;
716        assert_eq!(compiled.ret_ty(), &Type::Any);
717        let tail_object_variable: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
718        let result = unsafe { &*tail_object_variable() };
719        let steps = result.get_dynamic("steps").expect("steps");
720        assert_eq!(steps.len(), 2);
721        assert_eq!(steps.get_idx(1).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("second".to_string()));
722        Ok(())
723    }
724
725    #[test]
726    fn repeated_deep_step_literals_import_successfully() -> anyhow::Result<()> {
727        fn extra_page_literal(depth: usize) -> String {
728            let mut value = "{leaf: \"done\"}".to_string();
729            for idx in 0..depth {
730                value = format!("{{kind: \"page\", idx: {idx}, children: [{value}], meta: {{title: \"extra\", visible: true}}}}");
731            }
732            value
733        }
734
735        let extra = extra_page_literal(48);
736        let code = format!(
737            r#"
738            pub fn script() {{
739                return [
740                    {{ja: "一つ目", note: "first", extra: {extra}}},
741                    {{ja: "二つ目", note: "second", extra: {extra}}},
742                    {{ja: "三つ目", note: "third", extra: {extra}}}
743                ]
744            }}
745            "#
746        );
747
748        let vm = Vm::with_all()?;
749        vm.import_code("vm_repeated_deep_step_literals", code.into_bytes())?;
750        let compiled = vm.get_fn("vm_repeated_deep_step_literals::script", &[])?;
751        assert_eq!(compiled.ret_ty(), &Type::Any);
752        let script: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
753        let result = unsafe { &*script() };
754        assert_eq!(result.len(), 3);
755        assert_eq!(result.get_idx(2).and_then(|value| value.get_dynamic("note")).map(|value| value.as_str().to_string()), Some("third".to_string()));
756        Ok(())
757    }
758
759    #[test]
760    fn native_import_uses_owning_vm() -> anyhow::Result<()> {
761        let module_path = std::env::temp_dir().join(format!("zust_vm_import_owner_{}.zs", std::process::id()));
762        std::fs::write(&module_path, "pub fn value() { 41 }")?;
763        let module_path = module_path.to_string_lossy().replace('\\', "\\\\").replace('"', "\\\"");
764
765        let vm1 = Vm::with_all()?;
766        vm1.import_code(
767            "vm_import_owner",
768            format!(
769                r#"
770                pub fn run() {{
771                    import("vm_imported_owner", "{module_path}");
772                }}
773                "#
774            )
775            .into_bytes(),
776        )?;
777        let compiled = vm1.get_fn("vm_import_owner::run", &[])?;
778
779        let vm2 = Vm::with_all()?;
780        vm2.import_code("vm_import_other", b"pub fn run() { 0 }".to_vec())?;
781        let _ = vm2.get_fn("vm_import_other::run", &[])?;
782
783        let run: extern "C" fn() = unsafe { std::mem::transmute(compiled.ptr()) };
784        run();
785
786        assert!(vm1.get_fn("vm_imported_owner::value", &[]).is_ok());
787        assert!(vm2.get_fn("vm_imported_owner::value", &[]).is_err());
788        Ok(())
789    }
790
791    #[test]
792    fn object_last_field_call_does_not_need_trailing_comma() -> anyhow::Result<()> {
793        let vm = Vm::with_all()?;
794        vm.import_code(
795            "vm_object_last_call_field",
796            r#"
797            pub fn extra_page() {
798                {
799                    title: "extra",
800                    pages: [
801                        {note: "nested"}
802                    ]
803                }
804            }
805
806            pub fn data() {
807                return [
808                    {
809                        note: "first",
810                        choices: ["a", "b"],
811                        extras: extra_page()
812                    },
813                    {
814                        note: "second",
815                        choices: ["c"],
816                        extras: extra_page()
817                    }
818                ]
819            }
820            "#
821            .as_bytes()
822            .to_vec(),
823        )?;
824
825        let compiled = vm.get_fn("vm_object_last_call_field::data", &[])?;
826        assert_eq!(compiled.ret_ty(), &Type::Any);
827        let data: extern "C" fn() -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
828        let result = unsafe { &*data() };
829        assert_eq!(result.len(), 2);
830        let first = result.get_idx(0).expect("first step");
831        assert_eq!(first.get_dynamic("extras").and_then(|extras| extras.get_dynamic("title")).map(|title| title.as_str().to_string()), Some("extra".to_string()));
832        Ok(())
833    }
834
835    #[test]
836    fn gpu_struct_layout_packs_and_unpacks_dynamic_maps() -> anyhow::Result<()> {
837        let vm = Vm::with_all()?;
838        vm.import_code(
839            "vm_gpu_layout",
840            br#"
841            pub struct Params {
842                a: u32,
843                b: u32,
844                c: u32,
845            }
846            "#
847            .to_vec(),
848        )?;
849
850        let layout = vm.gpu_struct_layout("vm_gpu_layout::Params", &[])?;
851        assert_eq!(layout.size, 16);
852        assert_eq!(layout.fields.iter().map(|field| (field.name.as_str(), field.offset)).collect::<Vec<_>>(), vec![("a", 0), ("b", 4), ("c", 8)]);
853
854        let value = dynamic::map!("a"=> 1u32, "b"=> 2u32, "c"=> 3u32);
855        let bytes = layout.pack_map(&value)?;
856        assert_eq!(bytes.len(), 16);
857        assert_eq!(&bytes[0..4], &1u32.to_ne_bytes());
858        assert_eq!(&bytes[4..8], &2u32.to_ne_bytes());
859        assert_eq!(&bytes[8..12], &3u32.to_ne_bytes());
860
861        let read = layout.unpack_map(&bytes)?;
862        assert_eq!(read.get_dynamic("a").and_then(|value| value.as_uint()), Some(1));
863        assert_eq!(read.get_dynamic("b").and_then(|value| value.as_uint()), Some(2));
864        assert_eq!(read.get_dynamic("c").and_then(|value| value.as_uint()), Some(3));
865        Ok(())
866    }
867
868    #[test]
869    fn root_native_calls_do_not_take_ownership_of_dynamic_args() -> anyhow::Result<()> {
870        let vm = Vm::with_all()?;
871        vm.import_code(
872            "vm_root_clone_bridge",
873            br#"
874            pub fn add_then_reuse(arg) {
875                let user = {
876                    address: "test-wallet",
877                    points: 20
878                };
879                root::add("local/root-clone-bridge-user", user);
880                user.points = user.points - 7;
881                root::add("local/root-clone-bridge-user", user);
882                {
883                    user: user,
884                    points: user.points
885                }
886            }
887            "#
888            .to_vec(),
889        )?;
890
891        let compiled = vm.get_fn("vm_root_clone_bridge::add_then_reuse", &[Type::Any])?;
892        assert_eq!(compiled.ret_ty(), &Type::Any);
893        let add_then_reuse: extern "C" fn(*const Dynamic) -> *const Dynamic = unsafe { std::mem::transmute(compiled.ptr()) };
894        let arg = Dynamic::Null;
895        let result = add_then_reuse(&arg);
896        let result = unsafe { &*result };
897
898        assert_eq!(result.get_dynamic("points").and_then(|value| value.as_int()), Some(13));
899        let mut json = String::new();
900        result.to_json(&mut json);
901        assert!(json.contains("\"points\": 13"));
902        Ok(())
903    }
904}