lust/vm/
stdlib.rs

1use super::{
2    task::{TaskInstance, TaskState},
3    VM,
4};
5use crate::bytecode::{NativeCallResult, Value, ValueKey};
6use crate::config::LustConfig;
7use crate::LustError;
8use std::collections::HashMap;
9use std::fs;
10use std::io::{self, Read, Write};
11use std::rc::Rc;
12pub fn create_stdlib(config: &LustConfig) -> Vec<(&'static str, Value)> {
13    let mut stdlib = vec![
14        ("print", create_print_fn()),
15        ("println", create_println_fn()),
16        ("type", create_type_fn()),
17        ("tostring", create_tostring_fn()),
18        ("task", create_task_module()),
19    ];
20    if config.is_module_enabled("io") {
21        stdlib.push(("io", create_io_module()));
22    }
23
24    if config.is_module_enabled("os") {
25        stdlib.push(("os", create_os_module()));
26    }
27
28    stdlib
29}
30
31fn create_print_fn() -> Value {
32    Value::NativeFunction(Rc::new(|args: &[Value]| {
33        for (i, arg) in args.iter().enumerate() {
34            if i > 0 {
35                print!("\t");
36            }
37
38            print!("{}", arg);
39        }
40
41        Ok(NativeCallResult::Return(Value::Nil))
42    }))
43}
44
45fn create_println_fn() -> Value {
46    Value::NativeFunction(Rc::new(|args: &[Value]| {
47        for (i, arg) in args.iter().enumerate() {
48            if i > 0 {
49                print!("\t");
50            }
51
52            print!("{}", arg);
53        }
54
55        println!();
56        Ok(NativeCallResult::Return(Value::Nil))
57    }))
58}
59
60fn create_type_fn() -> Value {
61    Value::NativeFunction(Rc::new(|args: &[Value]| {
62        if args.is_empty() {
63            return Err("type() requires at least one argument".to_string());
64        }
65
66        let type_name = match &args[0] {
67            Value::Nil => "nil",
68            Value::Bool(_) => "bool",
69            Value::Int(_) => "int",
70            Value::Float(_) => "float",
71            Value::String(_) => "string",
72            Value::Array(_) => "array",
73            Value::Tuple(_) => "tuple",
74            Value::Map(_) => "map",
75            Value::Table(_) => "table",
76            Value::Struct { .. } | Value::WeakStruct(_) => "struct",
77            Value::Enum { .. } => "enum",
78            Value::Function(_) => "function",
79            Value::NativeFunction(_) => "function",
80            Value::Closure { .. } => "function",
81            Value::Iterator(_) => "iterator",
82            Value::Task(_) => "task",
83        };
84        Ok(NativeCallResult::Return(Value::string(type_name)))
85    }))
86}
87
88fn create_tostring_fn() -> Value {
89    Value::NativeFunction(Rc::new(|args: &[Value]| {
90        if args.is_empty() {
91            return Err("tostring() requires at least one argument".to_string());
92        }
93
94        Ok(NativeCallResult::Return(Value::string(format!(
95            "{}",
96            args[0]
97        ))))
98    }))
99}
100
101fn task_state_to_status_value(state: &TaskState) -> Value {
102    match state {
103        TaskState::Ready => Value::enum_unit("TaskStatus", "Ready"),
104        TaskState::Running => Value::enum_unit("TaskStatus", "Running"),
105        TaskState::Yielded => Value::enum_unit("TaskStatus", "Yielded"),
106        TaskState::Completed => Value::enum_unit("TaskStatus", "Completed"),
107        TaskState::Failed => Value::enum_unit("TaskStatus", "Failed"),
108        TaskState::Stopped => Value::enum_unit("TaskStatus", "Stopped"),
109    }
110}
111
112fn create_task_module() -> Value {
113    let mut entries: HashMap<ValueKey, Value> = HashMap::new();
114    entries.insert(string_key("run"), create_task_run_fn());
115    entries.insert(string_key("create"), create_task_create_fn());
116    entries.insert(string_key("status"), create_task_status_fn());
117    entries.insert(string_key("info"), create_task_info_fn());
118    entries.insert(string_key("resume"), create_task_resume_fn());
119    entries.insert(string_key("yield"), create_task_yield_fn());
120    entries.insert(string_key("stop"), create_task_stop_fn());
121    entries.insert(string_key("restart"), create_task_restart_fn());
122    entries.insert(string_key("current"), create_task_current_fn());
123    Value::table(entries)
124}
125
126fn create_task_run_fn() -> Value {
127    Value::NativeFunction(Rc::new(|args: &[Value]| {
128        if args.is_empty() {
129            return Err("task.run() requires a function".to_string());
130        }
131
132        let func = args[0].clone();
133        let rest: Vec<Value> = args.iter().skip(1).cloned().collect();
134        VM::with_current(move |vm| {
135            vm.spawn_task_value(func, rest)
136                .map(|handle| NativeCallResult::Return(Value::task(handle)))
137                .map_err(|e| e.to_string())
138        })
139    }))
140}
141
142fn create_task_create_fn() -> Value {
143    Value::NativeFunction(Rc::new(|args: &[Value]| {
144        if args.is_empty() {
145            return Err("task.create() requires a function".to_string());
146        }
147
148        let func = args[0].clone();
149        let rest: Vec<Value> = args.iter().skip(1).cloned().collect();
150        VM::with_current(move |vm| {
151            vm.create_task_value(func, rest)
152                .map(|handle| NativeCallResult::Return(Value::task(handle)))
153                .map_err(|e| e.to_string())
154        })
155    }))
156}
157
158fn create_task_status_fn() -> Value {
159    Value::NativeFunction(Rc::new(|args: &[Value]| {
160        if args.len() != 1 {
161            return Err("task.status() requires a task handle".to_string());
162        }
163
164        let handle = args[0]
165            .as_task_handle()
166            .ok_or_else(|| "task.status() requires a task handle value".to_string())?;
167        VM::with_current(move |vm| {
168            let task = vm.get_task_instance(handle).map_err(|e| e.to_string())?;
169            Ok(NativeCallResult::Return(task_state_to_status_value(
170                &task.state,
171            )))
172        })
173    }))
174}
175
176fn create_task_info_fn() -> Value {
177    Value::NativeFunction(Rc::new(|args: &[Value]| {
178        if args.len() != 1 {
179            return Err("task.info() requires a task handle".to_string());
180        }
181
182        let handle = args[0]
183            .as_task_handle()
184            .ok_or_else(|| "task.info() requires a task handle value".to_string())?;
185        VM::with_current(move |vm| {
186            let task = vm.get_task_instance(handle).map_err(|e| e.to_string())?;
187            let info = build_task_info_value(vm, task).map_err(|e| e.to_string())?;
188            Ok(NativeCallResult::Return(info))
189        })
190    }))
191}
192
193fn create_task_resume_fn() -> Value {
194    Value::NativeFunction(Rc::new(|args: &[Value]| {
195        if args.is_empty() {
196            return Err("task.resume() requires a task handle".to_string());
197        }
198
199        let handle = args[0]
200            .as_task_handle()
201            .ok_or_else(|| "task.resume() requires a task handle value".to_string())?;
202        let resume_value = args.get(1).cloned();
203        VM::with_current(move |vm| {
204            vm.resume_task_handle(handle, resume_value.clone())
205                .map_err(|e| e.to_string())?;
206            let task = vm.get_task_instance(handle).map_err(|e| e.to_string())?;
207            let info = build_task_info_value(vm, task).map_err(|e| e.to_string())?;
208            Ok(NativeCallResult::Return(info))
209        })
210    }))
211}
212
213fn create_task_yield_fn() -> Value {
214    Value::NativeFunction(Rc::new(|args: &[Value]| {
215        let value = args.get(0).cloned().unwrap_or(Value::Nil);
216        Ok(NativeCallResult::Yield(value))
217    }))
218}
219
220fn create_task_stop_fn() -> Value {
221    Value::NativeFunction(Rc::new(|args: &[Value]| {
222        if args.len() != 1 {
223            return Err("task.stop() requires a task handle".to_string());
224        }
225
226        let handle = args[0]
227            .as_task_handle()
228            .ok_or_else(|| "task.stop() requires a task handle value".to_string())?;
229        VM::with_current(move |vm| {
230            let before = vm
231                .get_task_instance(handle)
232                .map(|task| task.state.clone())
233                .map_err(|e| e.to_string())?;
234            vm.stop_task_handle(handle).map_err(|e| e.to_string())?;
235            let after = vm
236                .get_task_instance(handle)
237                .map(|task| task.state.clone())
238                .map_err(|e| e.to_string())?;
239            let changed = after == TaskState::Stopped && before != TaskState::Stopped;
240            Ok(NativeCallResult::Return(Value::Bool(changed)))
241        })
242    }))
243}
244
245fn create_task_restart_fn() -> Value {
246    Value::NativeFunction(Rc::new(|args: &[Value]| {
247        if args.len() != 1 {
248            return Err("task.restart() requires a task handle".to_string());
249        }
250
251        let handle = args[0]
252            .as_task_handle()
253            .ok_or_else(|| "task.restart() requires a task handle value".to_string())?;
254        VM::with_current(move |vm| {
255            vm.restart_task_handle(handle).map_err(|e| e.to_string())?;
256            let task = vm.get_task_instance(handle).map_err(|e| e.to_string())?;
257            let info = build_task_info_value(vm, task).map_err(|e| e.to_string())?;
258            Ok(NativeCallResult::Return(info))
259        })
260    }))
261}
262
263fn create_task_current_fn() -> Value {
264    Value::NativeFunction(Rc::new(|args: &[Value]| {
265        if !args.is_empty() {
266            return Err("task.current() takes no arguments".to_string());
267        }
268
269        VM::with_current(|vm| {
270            let value = match vm.current_task_handle() {
271                Some(handle) => Value::some(Value::task(handle)),
272                None => Value::none(),
273            };
274            Ok(NativeCallResult::Return(value))
275        })
276    }))
277}
278
279fn create_io_module() -> Value {
280    let mut entries: HashMap<ValueKey, Value> = HashMap::new();
281    entries.insert(string_key("read_file"), create_io_read_file_fn());
282    entries.insert(
283        string_key("read_file_bytes"),
284        create_io_read_file_bytes_fn(),
285    );
286    entries.insert(string_key("write_file"), create_io_write_file_fn());
287    entries.insert(string_key("read_stdin"), create_io_read_stdin_fn());
288    entries.insert(string_key("read_line"), create_io_read_line_fn());
289    entries.insert(string_key("write_stdout"), create_io_write_stdout_fn());
290    Value::table(entries)
291}
292
293fn create_io_read_file_fn() -> Value {
294    Value::NativeFunction(Rc::new(|args: &[Value]| {
295        if args.len() != 1 {
296            return Ok(NativeCallResult::Return(Value::err(Value::string(
297                "io.read_file(path) requires a single string path",
298            ))));
299        }
300
301        let path = match args[0].as_string() {
302            Some(p) => p,
303            None => {
304                return Ok(NativeCallResult::Return(Value::err(Value::string(
305                    "io.read_file(path) requires a string path",
306                ))))
307            }
308        };
309        match fs::read_to_string(path) {
310            Ok(contents) => Ok(NativeCallResult::Return(Value::ok(Value::string(contents)))),
311            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
312                err.to_string(),
313            )))),
314        }
315    }))
316}
317
318fn create_io_read_file_bytes_fn() -> Value {
319    Value::NativeFunction(Rc::new(|args: &[Value]| {
320        if args.len() != 1 {
321            return Ok(NativeCallResult::Return(Value::err(Value::string(
322                "io.read_file_bytes(path) requires a single string path",
323            ))));
324        }
325
326        let path = match args[0].as_string() {
327            Some(p) => p,
328            None => {
329                return Ok(NativeCallResult::Return(Value::err(Value::string(
330                    "io.read_file_bytes(path) requires a string path",
331                ))))
332            }
333        };
334
335        match fs::read(path) {
336            Ok(bytes) => {
337                let values: Vec<Value> = bytes.into_iter().map(|b| Value::Int(b as i64)).collect();
338                Ok(NativeCallResult::Return(Value::ok(Value::array(values))))
339            }
340
341            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
342                err.to_string(),
343            )))),
344        }
345    }))
346}
347
348fn create_io_write_file_fn() -> Value {
349    Value::NativeFunction(Rc::new(|args: &[Value]| {
350        if args.len() < 2 {
351            return Ok(NativeCallResult::Return(Value::err(Value::string(
352                "io.write_file(path, contents) requires a path and value",
353            ))));
354        }
355
356        let path = match args[0].as_string() {
357            Some(p) => p,
358            None => {
359                return Ok(NativeCallResult::Return(Value::err(Value::string(
360                    "io.write_file(path, contents) requires a string path",
361                ))))
362            }
363        };
364        let contents = if let Some(s) = args[1].as_string() {
365            s.to_string()
366        } else {
367            format!("{}", args[1])
368        };
369        match fs::write(path, contents) {
370            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
371            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
372                err.to_string(),
373            )))),
374        }
375    }))
376}
377
378fn create_io_read_stdin_fn() -> Value {
379    Value::NativeFunction(Rc::new(|args: &[Value]| {
380        if !args.is_empty() {
381            return Ok(NativeCallResult::Return(Value::err(Value::string(
382                "io.read_stdin() takes no arguments",
383            ))));
384        }
385
386        let mut buffer = String::new();
387        match io::stdin().read_to_string(&mut buffer) {
388            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::string(buffer)))),
389            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
390                err.to_string(),
391            )))),
392        }
393    }))
394}
395
396fn create_io_read_line_fn() -> Value {
397    Value::NativeFunction(Rc::new(|args: &[Value]| {
398        if !args.is_empty() {
399            return Ok(NativeCallResult::Return(Value::err(Value::string(
400                "io.read_line() takes no arguments",
401            ))));
402        }
403
404        let mut line = String::new();
405        match io::stdin().read_line(&mut line) {
406            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::string(line)))),
407            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
408                err.to_string(),
409            )))),
410        }
411    }))
412}
413
414fn create_io_write_stdout_fn() -> Value {
415    Value::NativeFunction(Rc::new(|args: &[Value]| {
416        let mut stdout = io::stdout();
417        for arg in args {
418            if let Err(err) = write!(stdout, "{}", arg) {
419                return Ok(NativeCallResult::Return(Value::err(Value::string(
420                    err.to_string(),
421                ))));
422            }
423        }
424
425        if let Err(err) = stdout.flush() {
426            return Ok(NativeCallResult::Return(Value::err(Value::string(
427                err.to_string(),
428            ))));
429        }
430
431        Ok(NativeCallResult::Return(Value::ok(Value::Nil)))
432    }))
433}
434
435fn create_os_module() -> Value {
436    let mut entries: HashMap<ValueKey, Value> = HashMap::new();
437    entries.insert(string_key("create_file"), create_os_create_file_fn());
438    entries.insert(string_key("create_dir"), create_os_create_dir_fn());
439    entries.insert(string_key("remove_file"), create_os_remove_file_fn());
440    entries.insert(string_key("remove_dir"), create_os_remove_dir_fn());
441    entries.insert(string_key("rename"), create_os_rename_fn());
442    Value::table(entries)
443}
444
445fn create_os_create_file_fn() -> Value {
446    Value::NativeFunction(Rc::new(|args: &[Value]| {
447        if args.len() != 1 {
448            return Ok(NativeCallResult::Return(Value::err(Value::string(
449                "os.create_file(path) requires a single string path",
450            ))));
451        }
452
453        let path = match args[0].as_string() {
454            Some(p) => p,
455            None => {
456                return Ok(NativeCallResult::Return(Value::err(Value::string(
457                    "os.create_file(path) requires a string path",
458                ))))
459            }
460        };
461        match fs::OpenOptions::new().write(true).create(true).open(path) {
462            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
463            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
464                err.to_string(),
465            )))),
466        }
467    }))
468}
469
470fn create_os_create_dir_fn() -> Value {
471    Value::NativeFunction(Rc::new(|args: &[Value]| {
472        if args.len() != 1 {
473            return Ok(NativeCallResult::Return(Value::err(Value::string(
474                "os.create_dir(path) requires a single string path",
475            ))));
476        }
477
478        let path = match args[0].as_string() {
479            Some(p) => p,
480            None => {
481                return Ok(NativeCallResult::Return(Value::err(Value::string(
482                    "os.create_dir(path) requires a string path",
483                ))))
484            }
485        };
486        match fs::create_dir_all(path) {
487            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
488            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
489                err.to_string(),
490            )))),
491        }
492    }))
493}
494
495fn create_os_remove_file_fn() -> Value {
496    Value::NativeFunction(Rc::new(|args: &[Value]| {
497        if args.len() != 1 {
498            return Ok(NativeCallResult::Return(Value::err(Value::string(
499                "os.remove_file(path) requires a single string path",
500            ))));
501        }
502
503        let path = match args[0].as_string() {
504            Some(p) => p,
505            None => {
506                return Ok(NativeCallResult::Return(Value::err(Value::string(
507                    "os.remove_file(path) requires a string path",
508                ))))
509            }
510        };
511        match fs::remove_file(path) {
512            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
513            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
514                err.to_string(),
515            )))),
516        }
517    }))
518}
519
520fn create_os_remove_dir_fn() -> Value {
521    Value::NativeFunction(Rc::new(|args: &[Value]| {
522        if args.len() != 1 {
523            return Ok(NativeCallResult::Return(Value::err(Value::string(
524                "os.remove_dir(path) requires a single string path",
525            ))));
526        }
527
528        let path = match args[0].as_string() {
529            Some(p) => p,
530            None => {
531                return Ok(NativeCallResult::Return(Value::err(Value::string(
532                    "os.remove_dir(path) requires a string path",
533                ))))
534            }
535        };
536        match fs::remove_dir_all(path) {
537            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
538            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
539                err.to_string(),
540            )))),
541        }
542    }))
543}
544
545fn create_os_rename_fn() -> Value {
546    Value::NativeFunction(Rc::new(|args: &[Value]| {
547        if args.len() != 2 {
548            return Ok(NativeCallResult::Return(Value::err(Value::string(
549                "os.rename(from, to) requires two string paths",
550            ))));
551        }
552
553        let from = match args[0].as_string() {
554            Some(f) => f,
555            None => {
556                return Ok(NativeCallResult::Return(Value::err(Value::string(
557                    "os.rename(from, to) requires string paths",
558                ))))
559            }
560        };
561        let to = match args[1].as_string() {
562            Some(t) => t,
563            None => {
564                return Ok(NativeCallResult::Return(Value::err(Value::string(
565                    "os.rename(from, to) requires string paths",
566                ))))
567            }
568        };
569        match fs::rename(from, to) {
570            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
571            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
572                err.to_string(),
573            )))),
574        }
575    }))
576}
577
578fn build_task_info_value(vm: &VM, task: &TaskInstance) -> Result<Value, LustError> {
579    let last_yield = match &task.last_yield {
580        Some(value) => Value::some(value.clone()),
581        None => Value::none(),
582    };
583    let last_result = match &task.last_result {
584        Some(value) => Value::some(value.clone()),
585        None => Value::none(),
586    };
587    let error = match task.error.as_ref() {
588        Some(err) => Value::some(Value::string(err.to_string())),
589        None => Value::none(),
590    };
591    vm.instantiate_struct(
592        "TaskInfo",
593        vec![
594            (
595                Rc::new("state".to_string()),
596                task_state_to_status_value(&task.state),
597            ),
598            (Rc::new("last_yield".to_string()), last_yield),
599            (Rc::new("last_result".to_string()), last_result),
600            (Rc::new("error".to_string()), error),
601        ],
602    )
603}
604
605fn string_key(name: &str) -> ValueKey {
606    ValueKey::String(Rc::new(name.to_string()))
607}
608
609#[cfg(test)]
610mod tests {
611    use super::*;
612    #[test]
613    fn stdlib_defaults_without_optional_modules() {
614        let stdlib = create_stdlib(&LustConfig::default());
615        assert!(!stdlib.iter().any(|(name, _)| *name == "io"));
616        assert!(!stdlib.iter().any(|(name, _)| *name == "os"));
617    }
618
619    #[test]
620    fn stdlib_includes_optional_modules_when_configured() {
621        let cfg =
622            LustConfig::from_toml_str("\"enabled modules\" = [\"io\", \"os\"]").expect("parse");
623        let stdlib = create_stdlib(&cfg);
624        assert!(stdlib.iter().any(|(name, _)| *name == "io"));
625        assert!(stdlib.iter().any(|(name, _)| *name == "os"));
626    }
627}