Skip to main content

lust/vm/
api.rs

1use super::*;
2use crate::ast::Type;
3use crate::bytecode::{LustMap, ValueKey};
4use crate::config::LustConfig;
5use alloc::rc::Rc;
6use alloc::string::String;
7use alloc::vec::Vec;
8use core::result::Result as CoreResult;
9use core::{array, mem};
10use hashbrown::HashMap;
11
12#[derive(Debug, Clone)]
13pub struct NativeExportParam {
14    name: String,
15    ty: String,
16}
17
18impl NativeExportParam {
19    pub fn new(name: impl Into<String>, ty: impl Into<String>) -> Self {
20        Self {
21            name: name.into(),
22            ty: ty.into(),
23        }
24    }
25
26    pub fn name(&self) -> &str {
27        &self.name
28    }
29
30    pub fn ty(&self) -> &str {
31        &self.ty
32    }
33}
34
35#[derive(Debug, Clone)]
36pub struct NativeExport {
37    name: String,
38    params: Vec<NativeExportParam>,
39    return_type: String,
40    doc: Option<String>,
41}
42
43impl NativeExport {
44    pub fn new(
45        name: impl Into<String>,
46        params: Vec<NativeExportParam>,
47        return_type: impl Into<String>,
48    ) -> Self {
49        Self {
50            name: name.into(),
51            params,
52            return_type: return_type.into(),
53            doc: None,
54        }
55    }
56
57    pub fn with_doc(mut self, doc: impl Into<String>) -> Self {
58        self.doc = Some(doc.into());
59        self
60    }
61
62    pub fn name(&self) -> &str {
63        &self.name
64    }
65
66    pub fn params(&self) -> &[NativeExportParam] {
67        &self.params
68    }
69
70    pub fn return_type(&self) -> &str {
71        &self.return_type
72    }
73
74    pub fn doc(&self) -> Option<&str> {
75        self.doc.as_deref()
76    }
77}
78impl VM {
79    pub fn new() -> Self {
80        Self::with_config(&LustConfig::default())
81    }
82
83    pub fn with_config(config: &LustConfig) -> Self {
84        let mut vm = Self {
85            jit: JitState::new(),
86            budgets: BudgetState::default(),
87            functions: Vec::new(),
88            natives: HashMap::new(),
89            globals: HashMap::new(),
90            map_hasher: DefaultHashBuilder::default(),
91            call_stack: Vec::new(),
92            max_stack_depth: 1000,
93            pending_return_value: None,
94            pending_return_dest: None,
95            trace_recorder: None,
96            side_trace_context: None,
97            skip_next_trace_record: false,
98            trait_impls: HashMap::new(),
99            struct_tostring_cache: HashMap::new(),
100            struct_metadata: HashMap::new(),
101            call_until_depth: None,
102            task_manager: TaskManager::new(),
103            current_task: None,
104            pending_task_signal: None,
105            last_task_signal: None,
106            cycle_collector: cycle::CycleCollector::new(),
107            exported_natives: Vec::new(),
108            export_prefix_stack: Vec::new(),
109            exported_type_stubs: Vec::new(),
110        };
111        vm.jit.enabled = vm.jit.enabled && config.jit_enabled();
112        vm.trait_impls
113            .insert(("int".to_string(), "ToString".to_string()), true);
114        vm.trait_impls
115            .insert(("float".to_string(), "ToString".to_string()), true);
116        vm.trait_impls
117            .insert(("string".to_string(), "ToString".to_string()), true);
118        vm.trait_impls
119            .insert(("bool".to_string(), "ToString".to_string()), true);
120        super::corelib::install_core_builtins(&mut vm);
121        #[cfg(feature = "std")]
122        for (name, func) in super::stdlib::create_stdlib(config, &vm) {
123            vm.register_native(name, func);
124        }
125
126        vm
127    }
128
129    pub(crate) fn new_map(&self) -> LustMap {
130        HashMap::with_hasher(self.map_hasher.clone())
131    }
132
133    pub(crate) fn map_with_entries(
134        &self,
135        entries: impl IntoIterator<Item = (ValueKey, Value)>,
136    ) -> Value {
137        let mut map = self.new_map();
138        map.extend(entries);
139        Value::map(map)
140    }
141
142    pub(crate) fn new_map_value(&self) -> Value {
143        Value::map(self.new_map())
144    }
145
146    pub(super) fn observe_value(&mut self, value: &Value) {
147        self.cycle_collector.register_value(value);
148    }
149
150    pub(super) fn maybe_collect_cycles(&mut self) {
151        let mut collector = mem::take(&mut self.cycle_collector);
152        collector.maybe_collect(self);
153        self.cycle_collector = collector;
154    }
155
156    pub fn with_current<F, R>(f: F) -> CoreResult<R, String>
157    where
158        F: FnOnce(&mut VM) -> CoreResult<R, String>,
159    {
160        let ptr_opt = super::with_vm_stack(|stack| stack.last().copied());
161        if let Some(ptr) = ptr_opt {
162            let vm = unsafe { &mut *ptr };
163            f(vm)
164        } else {
165            Err("task API requires a running VM".to_string())
166        }
167    }
168
169    pub fn load_functions(&mut self, functions: Vec<Function>) {
170        self.functions = functions;
171    }
172
173    pub fn register_structs(&mut self, defs: &HashMap<String, StructDef>) {
174        for (name, def) in defs {
175            let field_names: Vec<Rc<String>> = def
176                .fields
177                .iter()
178                .map(|field| Rc::new(field.name.clone()))
179                .collect();
180            let field_storage: Vec<FieldStorage> = def
181                .fields
182                .iter()
183                .map(|field| match field.ownership {
184                    FieldOwnership::Weak => FieldStorage::Weak,
185                    FieldOwnership::Strong => FieldStorage::Strong,
186                })
187                .collect();
188            let field_types: Vec<Type> = def.fields.iter().map(|field| field.ty.clone()).collect();
189            let weak_targets: Vec<Option<Type>> = def
190                .fields
191                .iter()
192                .map(|field| field.weak_target.clone())
193                .collect();
194            let layout = Rc::new(StructLayout::new(
195                def.name.clone(),
196                field_names,
197                field_storage,
198                field_types,
199                weak_targets,
200            ));
201            self.struct_metadata.insert(
202                name.clone(),
203                RuntimeStructInfo {
204                    layout: layout.clone(),
205                },
206            );
207            if let Some(simple) = name.rsplit('.').next() {
208                self.struct_metadata.insert(
209                    simple.to_string(),
210                    RuntimeStructInfo {
211                        layout: layout.clone(),
212                    },
213                );
214            }
215        }
216    }
217
218    pub fn instantiate_struct(
219        &self,
220        struct_name: &str,
221        fields: Vec<(Rc<String>, Value)>,
222    ) -> Result<Value> {
223        let info =
224            self.struct_metadata
225                .get(struct_name)
226                .ok_or_else(|| LustError::RuntimeError {
227                    message: format!("Unknown struct '{}'", struct_name),
228                })?;
229        Self::build_struct_value(struct_name, info, fields)
230    }
231
232    fn build_struct_value(
233        struct_name: &str,
234        info: &RuntimeStructInfo,
235        mut fields: Vec<(Rc<String>, Value)>,
236    ) -> Result<Value> {
237        let layout = info.layout.clone();
238        let field_count = layout.field_names().len();
239        let mut ordered = vec![Value::Nil; field_count];
240        let mut filled = vec![false; field_count];
241        for (field_name, field_value) in fields.drain(..) {
242            let index_opt = layout
243                .index_of_rc(&field_name)
244                .or_else(|| layout.index_of_str(field_name.as_str()));
245            let index = match index_opt {
246                Some(i) => i,
247                None => {
248                    return Err(LustError::RuntimeError {
249                        message: format!("Struct '{}' has no field '{}'", struct_name, field_name),
250                    })
251                }
252            };
253            let canonical = layout
254                .canonicalize_field_value(index, field_value)
255                .map_err(|msg| LustError::RuntimeError { message: msg })?;
256            ordered[index] = canonical;
257            filled[index] = true;
258        }
259
260        if filled.iter().any(|slot| !*slot) {
261            let missing: Vec<String> = layout
262                .field_names()
263                .iter()
264                .enumerate()
265                .filter_map(|(idx, name)| (!filled[idx]).then(|| (**name).clone()))
266                .collect();
267            return Err(LustError::RuntimeError {
268                message: format!(
269                    "Struct '{}' is missing required field(s): {}",
270                    struct_name,
271                    missing.join(", ")
272                ),
273            });
274        }
275
276        Ok(Value::Struct {
277            name: struct_name.to_string(),
278            layout,
279            fields: Rc::new(RefCell::new(ordered)),
280        })
281    }
282
283    pub fn register_trait_impl(&mut self, type_name: String, trait_name: String) {
284        self.trait_impls.insert((type_name, trait_name), true);
285    }
286
287    pub fn register_native(&mut self, name: impl Into<String>, value: Value) {
288        let name = name.into();
289        match value {
290            Value::NativeFunction(_) => {
291                let cloned = value.clone();
292                self.natives.insert(name.clone(), value);
293                self.globals.insert(name, cloned);
294            }
295
296            other => {
297                self.globals.insert(name, other);
298            }
299        }
300    }
301
302    pub(crate) fn push_export_prefix(&mut self, crate_name: &str) {
303        let sanitized = crate_name.replace('-', "_");
304        self.export_prefix_stack.push(sanitized);
305    }
306
307    pub(crate) fn pop_export_prefix(&mut self) {
308        self.export_prefix_stack.pop();
309    }
310
311    fn current_export_prefix(&self) -> Option<&str> {
312        self.export_prefix_stack.last().map(|s| s.as_str())
313    }
314
315    pub fn export_prefix(&self) -> Option<String> {
316        self.current_export_prefix().map(|s| s.to_string())
317    }
318
319    fn canonicalize_export_name(&self, export: &mut NativeExport) {
320        if let Some(prefix) = self.current_export_prefix() {
321            let needs_prefix = match export.name.strip_prefix(prefix) {
322                Some(rest) => {
323                    if rest.is_empty() {
324                        false
325                    } else {
326                        !matches!(rest.chars().next(), Some('.') | Some(':'))
327                    }
328                }
329                None => true,
330            };
331            if needs_prefix {
332                export.name = if export.name.is_empty() {
333                    prefix.to_string()
334                } else {
335                    format!("{prefix}.{}", export.name)
336                };
337            }
338        }
339    }
340
341    fn push_export_metadata(&mut self, export: NativeExport) {
342        if self.exported_natives.iter().any(|existing| existing.name == export.name) {
343            return;
344        }
345        self.exported_natives.push(export);
346    }
347
348    pub fn record_exported_native(&mut self, mut export: NativeExport) {
349        self.canonicalize_export_name(&mut export);
350        self.push_export_metadata(export);
351    }
352
353    pub fn register_exported_native<F>(&mut self, export: NativeExport, func: F)
354    where
355        F: Fn(&[Value]) -> CoreResult<NativeCallResult, String> + 'static,
356    {
357        let mut export = export;
358        self.canonicalize_export_name(&mut export);
359        let name = export.name.clone();
360        self.push_export_metadata(export);
361        let native = Value::NativeFunction(Rc::new(func));
362        self.register_native(name, native);
363    }
364
365    pub fn register_type_stubs(&mut self, stubs: Vec<ModuleStub>) {
366        if stubs.is_empty() {
367            return;
368        }
369        self.exported_type_stubs.extend(stubs);
370    }
371
372    pub fn exported_type_stubs(&self) -> &[ModuleStub] {
373        &self.exported_type_stubs
374    }
375
376    pub fn take_type_stubs(&mut self) -> Vec<ModuleStub> {
377        mem::take(&mut self.exported_type_stubs)
378    }
379
380    pub fn exported_natives(&self) -> &[NativeExport] {
381        &self.exported_natives
382    }
383
384    pub fn take_exported_natives(&mut self) -> Vec<NativeExport> {
385        mem::take(&mut self.exported_natives)
386    }
387
388    pub fn clear_native_functions(&mut self) {
389        self.natives.clear();
390        self.exported_type_stubs.clear();
391    }
392
393    #[cfg(feature = "std")]
394    pub fn dump_externs_to_dir(
395        &self,
396        output_root: impl AsRef<std::path::Path>,
397    ) -> std::io::Result<Vec<std::path::PathBuf>> {
398        self.dump_externs_to_dir_with_options(
399            output_root,
400            &crate::externs::DumpExternsOptions::default(),
401        )
402    }
403
404    #[cfg(feature = "std")]
405    pub fn dump_externs_to_dir_with_options(
406        &self,
407        output_root: impl AsRef<std::path::Path>,
408        options: &crate::externs::DumpExternsOptions,
409    ) -> std::io::Result<Vec<std::path::PathBuf>> {
410        let files = crate::externs::extern_files_from_vm(self, options);
411        crate::externs::write_extern_files(output_root, &files)
412    }
413
414    pub fn get_global(&self, name: &str) -> Option<Value> {
415        if let Some(value) = self.globals.get(name) {
416            Some(value.clone())
417        } else {
418            self.natives.get(name).cloned()
419        }
420    }
421
422    pub fn global_names(&self) -> Vec<String> {
423        self.globals.keys().cloned().collect()
424    }
425
426    pub fn globals_snapshot(&self) -> Vec<(String, Value)> {
427        self.globals
428            .iter()
429            .map(|(name, value)| (name.clone(), value.clone()))
430            .collect()
431    }
432
433    pub fn set_global(&mut self, name: impl Into<String>, value: Value) {
434        let name = name.into();
435        self.observe_value(&value);
436        self.globals.insert(name.clone(), value);
437        self.natives.remove(&name);
438        self.maybe_collect_cycles();
439    }
440
441    pub fn call(&mut self, function_name: &str, args: Vec<Value>) -> Result<Value> {
442        let func_idx = self
443            .functions
444            .iter()
445            .position(|f| f.name == function_name)
446            .ok_or_else(|| LustError::RuntimeError {
447                message: format!("Function not found: {}", function_name),
448            })?;
449        let mut frame = CallFrame {
450            function_idx: func_idx,
451            ip: 0,
452            registers: array::from_fn(|_| Value::Nil),
453            base_register: 0,
454            return_dest: None,
455            upvalues: Vec::new(),
456        };
457        let func = &self.functions[func_idx];
458        if args.len() != func.param_count as usize {
459            return Err(LustError::RuntimeError {
460                message: format!(
461                    "Function {} expects {} arguments, got {}",
462                    function_name,
463                    func.param_count,
464                    args.len()
465                ),
466            });
467        }
468
469        for (i, arg) in args.into_iter().enumerate() {
470            frame.registers[i] = arg;
471        }
472
473        self.call_stack.push(frame);
474        match self.run() {
475            Ok(v) => Ok(v),
476            Err(e) => Err(self.annotate_runtime_error(e)),
477        }
478    }
479
480    pub fn function_value(&self, function_name: &str) -> Option<Value> {
481        let canonical = if function_name.contains("::") {
482            function_name.replace("::", ".")
483        } else {
484            function_name.to_string()
485        };
486        self.functions
487            .iter()
488            .position(|f| f.name == canonical)
489            .map(Value::Function)
490    }
491
492    pub fn function_name(&self, index: usize) -> Option<&str> {
493        self.functions.get(index).map(|f| f.name.as_str())
494    }
495
496    pub fn fail_task_handle(&mut self, handle: TaskHandle, error: LustError) -> Result<()> {
497        let task_id = self.task_id_from_handle(handle)?;
498        let mut task =
499            self.task_manager
500                .detach(task_id)
501                .ok_or_else(|| LustError::RuntimeError {
502                    message: format!("Invalid task handle {}", handle.id()),
503                })?;
504        task.state = TaskState::Failed;
505        task.error = Some(error.clone());
506        task.last_yield = None;
507        task.last_result = None;
508        task.yield_dest = None;
509        task.call_stack.clear();
510        task.pending_return_value = None;
511        task.pending_return_dest = None;
512        self.task_manager.attach(task);
513        Err(error)
514    }
515}