lust/vm/
stdlib.rs

1use super::corelib::string_key;
2use crate::bytecode::{NativeCallResult, Value, ValueKey};
3use crate::config::LustConfig;
4use crate::LustInt;
5use hashbrown::HashMap;
6use std::fs;
7use std::io::{self, Read, Write};
8use std::rc::Rc;
9pub fn create_stdlib(config: &LustConfig) -> Vec<(&'static str, Value)> {
10    let mut stdlib = vec![
11        ("print", create_print_fn()),
12        ("println", create_println_fn()),
13        ("type", create_type_fn()),
14    ];
15    if config.is_module_enabled("io") {
16        stdlib.push(("io", create_io_module()));
17    }
18
19    if config.is_module_enabled("os") {
20        stdlib.push(("os", create_os_module()));
21    }
22
23    stdlib
24}
25
26fn create_print_fn() -> Value {
27    Value::NativeFunction(Rc::new(|args: &[Value]| {
28        for (i, arg) in args.iter().enumerate() {
29            if i > 0 {
30                print!("\t");
31            }
32
33            print!("{}", arg);
34        }
35
36        Ok(NativeCallResult::Return(Value::Nil))
37    }))
38}
39
40fn create_println_fn() -> Value {
41    Value::NativeFunction(Rc::new(|args: &[Value]| {
42        for (i, arg) in args.iter().enumerate() {
43            if i > 0 {
44                print!("\t");
45            }
46
47            print!("{}", arg);
48        }
49
50        println!();
51        Ok(NativeCallResult::Return(Value::Nil))
52    }))
53}
54
55fn create_type_fn() -> Value {
56    Value::NativeFunction(Rc::new(|args: &[Value]| {
57        if args.is_empty() {
58            return Err("type() requires at least one argument".to_string());
59        }
60
61        let type_name = match &args[0] {
62            Value::Nil => "nil",
63            Value::Bool(_) => "bool",
64            Value::Int(_) => "int",
65            Value::Float(_) => "float",
66            Value::String(_) => "string",
67            Value::Array(_) => "array",
68            Value::Tuple(_) => "tuple",
69            Value::Map(_) => "map",
70            Value::Struct { .. } | Value::WeakStruct(_) => "struct",
71            Value::Enum { .. } => "enum",
72            Value::Function(_) => "function",
73            Value::NativeFunction(_) => "function",
74            Value::Closure { .. } => "function",
75            Value::Iterator(_) => "iterator",
76            Value::Task(_) => "task",
77        };
78        Ok(NativeCallResult::Return(Value::string(type_name)))
79    }))
80}
81
82fn create_io_module() -> Value {
83    let mut entries: HashMap<ValueKey, Value> = HashMap::new();
84    entries.insert(string_key("read_file"), create_io_read_file_fn());
85    entries.insert(
86        string_key("read_file_bytes"),
87        create_io_read_file_bytes_fn(),
88    );
89    entries.insert(string_key("write_file"), create_io_write_file_fn());
90    entries.insert(string_key("read_stdin"), create_io_read_stdin_fn());
91    entries.insert(string_key("read_line"), create_io_read_line_fn());
92    entries.insert(string_key("write_stdout"), create_io_write_stdout_fn());
93    Value::map(entries)
94}
95
96fn create_io_read_file_fn() -> Value {
97    Value::NativeFunction(Rc::new(|args: &[Value]| {
98        if args.len() != 1 {
99            return Ok(NativeCallResult::Return(Value::err(Value::string(
100                "io.read_file(path) requires a single string path",
101            ))));
102        }
103
104        let path = match args[0].as_string() {
105            Some(p) => p,
106            None => {
107                return Ok(NativeCallResult::Return(Value::err(Value::string(
108                    "io.read_file(path) requires a string path",
109                ))))
110            }
111        };
112        match fs::read_to_string(path) {
113            Ok(contents) => Ok(NativeCallResult::Return(Value::ok(Value::string(contents)))),
114            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
115                err.to_string(),
116            )))),
117        }
118    }))
119}
120
121fn create_io_read_file_bytes_fn() -> Value {
122    Value::NativeFunction(Rc::new(|args: &[Value]| {
123        if args.len() != 1 {
124            return Ok(NativeCallResult::Return(Value::err(Value::string(
125                "io.read_file_bytes(path) requires a single string path",
126            ))));
127        }
128
129        let path = match args[0].as_string() {
130            Some(p) => p,
131            None => {
132                return Ok(NativeCallResult::Return(Value::err(Value::string(
133                    "io.read_file_bytes(path) requires a string path",
134                ))))
135            }
136        };
137
138        match fs::read(path) {
139            Ok(bytes) => {
140                let values: Vec<Value> = bytes
141                    .into_iter()
142                    .map(|b| Value::Int(b as LustInt))
143                    .collect();
144                Ok(NativeCallResult::Return(Value::ok(Value::array(values))))
145            }
146
147            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
148                err.to_string(),
149            )))),
150        }
151    }))
152}
153
154fn create_io_write_file_fn() -> Value {
155    Value::NativeFunction(Rc::new(|args: &[Value]| {
156        if args.len() < 2 {
157            return Ok(NativeCallResult::Return(Value::err(Value::string(
158                "io.write_file(path, contents) requires a path and value",
159            ))));
160        }
161
162        let path = match args[0].as_string() {
163            Some(p) => p,
164            None => {
165                return Ok(NativeCallResult::Return(Value::err(Value::string(
166                    "io.write_file(path, contents) requires a string path",
167                ))))
168            }
169        };
170        let contents = if let Some(s) = args[1].as_string() {
171            s.to_string()
172        } else {
173            format!("{}", args[1])
174        };
175        match fs::write(path, contents) {
176            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
177            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
178                err.to_string(),
179            )))),
180        }
181    }))
182}
183
184fn create_io_read_stdin_fn() -> Value {
185    Value::NativeFunction(Rc::new(|args: &[Value]| {
186        if !args.is_empty() {
187            return Ok(NativeCallResult::Return(Value::err(Value::string(
188                "io.read_stdin() takes no arguments",
189            ))));
190        }
191
192        let mut buffer = String::new();
193        match io::stdin().read_to_string(&mut buffer) {
194            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::string(buffer)))),
195            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
196                err.to_string(),
197            )))),
198        }
199    }))
200}
201
202fn create_io_read_line_fn() -> Value {
203    Value::NativeFunction(Rc::new(|args: &[Value]| {
204        if !args.is_empty() {
205            return Ok(NativeCallResult::Return(Value::err(Value::string(
206                "io.read_line() takes no arguments",
207            ))));
208        }
209
210        let mut line = String::new();
211        match io::stdin().read_line(&mut line) {
212            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::string(line)))),
213            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
214                err.to_string(),
215            )))),
216        }
217    }))
218}
219
220fn create_io_write_stdout_fn() -> Value {
221    Value::NativeFunction(Rc::new(|args: &[Value]| {
222        let mut stdout = io::stdout();
223        for arg in args {
224            if let Err(err) = write!(stdout, "{}", arg) {
225                return Ok(NativeCallResult::Return(Value::err(Value::string(
226                    err.to_string(),
227                ))));
228            }
229        }
230
231        if let Err(err) = stdout.flush() {
232            return Ok(NativeCallResult::Return(Value::err(Value::string(
233                err.to_string(),
234            ))));
235        }
236
237        Ok(NativeCallResult::Return(Value::ok(Value::Nil)))
238    }))
239}
240
241fn create_os_module() -> Value {
242    let mut entries: HashMap<ValueKey, Value> = HashMap::new();
243    entries.insert(string_key("create_file"), create_os_create_file_fn());
244    entries.insert(string_key("create_dir"), create_os_create_dir_fn());
245    entries.insert(string_key("remove_file"), create_os_remove_file_fn());
246    entries.insert(string_key("remove_dir"), create_os_remove_dir_fn());
247    entries.insert(string_key("rename"), create_os_rename_fn());
248    Value::map(entries)
249}
250
251fn create_os_create_file_fn() -> Value {
252    Value::NativeFunction(Rc::new(|args: &[Value]| {
253        if args.len() != 1 {
254            return Ok(NativeCallResult::Return(Value::err(Value::string(
255                "os.create_file(path) requires a single string path",
256            ))));
257        }
258
259        let path = match args[0].as_string() {
260            Some(p) => p,
261            None => {
262                return Ok(NativeCallResult::Return(Value::err(Value::string(
263                    "os.create_file(path) requires a string path",
264                ))))
265            }
266        };
267        match fs::OpenOptions::new().write(true).create(true).open(path) {
268            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
269            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
270                err.to_string(),
271            )))),
272        }
273    }))
274}
275
276fn create_os_create_dir_fn() -> Value {
277    Value::NativeFunction(Rc::new(|args: &[Value]| {
278        if args.len() != 1 {
279            return Ok(NativeCallResult::Return(Value::err(Value::string(
280                "os.create_dir(path) requires a single string path",
281            ))));
282        }
283
284        let path = match args[0].as_string() {
285            Some(p) => p,
286            None => {
287                return Ok(NativeCallResult::Return(Value::err(Value::string(
288                    "os.create_dir(path) requires a string path",
289                ))))
290            }
291        };
292        match fs::create_dir_all(path) {
293            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
294            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
295                err.to_string(),
296            )))),
297        }
298    }))
299}
300
301fn create_os_remove_file_fn() -> Value {
302    Value::NativeFunction(Rc::new(|args: &[Value]| {
303        if args.len() != 1 {
304            return Ok(NativeCallResult::Return(Value::err(Value::string(
305                "os.remove_file(path) requires a single string path",
306            ))));
307        }
308
309        let path = match args[0].as_string() {
310            Some(p) => p,
311            None => {
312                return Ok(NativeCallResult::Return(Value::err(Value::string(
313                    "os.remove_file(path) requires a string path",
314                ))))
315            }
316        };
317        match fs::remove_file(path) {
318            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
319            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
320                err.to_string(),
321            )))),
322        }
323    }))
324}
325
326fn create_os_remove_dir_fn() -> Value {
327    Value::NativeFunction(Rc::new(|args: &[Value]| {
328        if args.len() != 1 {
329            return Ok(NativeCallResult::Return(Value::err(Value::string(
330                "os.remove_dir(path) requires a single string path",
331            ))));
332        }
333
334        let path = match args[0].as_string() {
335            Some(p) => p,
336            None => {
337                return Ok(NativeCallResult::Return(Value::err(Value::string(
338                    "os.remove_dir(path) requires a string path",
339                ))))
340            }
341        };
342        match fs::remove_dir_all(path) {
343            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
344            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
345                err.to_string(),
346            )))),
347        }
348    }))
349}
350
351fn create_os_rename_fn() -> Value {
352    Value::NativeFunction(Rc::new(|args: &[Value]| {
353        if args.len() != 2 {
354            return Ok(NativeCallResult::Return(Value::err(Value::string(
355                "os.rename(from, to) requires two string paths",
356            ))));
357        }
358
359        let from = match args[0].as_string() {
360            Some(f) => f,
361            None => {
362                return Ok(NativeCallResult::Return(Value::err(Value::string(
363                    "os.rename(from, to) requires string paths",
364                ))))
365            }
366        };
367        let to = match args[1].as_string() {
368            Some(t) => t,
369            None => {
370                return Ok(NativeCallResult::Return(Value::err(Value::string(
371                    "os.rename(from, to) requires string paths",
372                ))))
373            }
374        };
375        match fs::rename(from, to) {
376            Ok(_) => Ok(NativeCallResult::Return(Value::ok(Value::Nil))),
377            Err(err) => Ok(NativeCallResult::Return(Value::err(Value::string(
378                err.to_string(),
379            )))),
380        }
381    }))
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387    #[test]
388    fn stdlib_defaults_without_optional_modules() {
389        let stdlib = create_stdlib(&LustConfig::default());
390        assert!(!stdlib.iter().any(|(name, _)| *name == "io"));
391        assert!(!stdlib.iter().any(|(name, _)| *name == "os"));
392    }
393
394    #[test]
395    fn stdlib_includes_optional_modules_when_configured() {
396        let cfg =
397            LustConfig::from_toml_str("\"enabled modules\" = [\"io\", \"os\"]").expect("parse");
398        let stdlib = create_stdlib(&cfg);
399        assert!(stdlib.iter().any(|(name, _)| *name == "io"));
400        assert!(stdlib.iter().any(|(name, _)| *name == "os"));
401    }
402}