Skip to main content

ion_core/
stdlib.rs

1//! Built-in standard library modules.
2//!
3//! These modules are automatically registered in every Engine instance
4//! and provide namespaced access to common functions and constants.
5//! The same functions remain available as top-level builtins for
6//! backwards compatibility.
7
8use std::sync::Arc;
9
10use crate::module::Module;
11use crate::value::Value;
12
13/// Output stream requested by Ion's `io` stdlib module.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum OutputStream {
16    Stdout,
17    Stderr,
18}
19
20/// Host-side output handler for Ion's `io::print*` functions.
21///
22/// Embedders install an implementation with `Engine::set_output` or
23/// `Engine::with_output`. The core runtime never writes directly to process
24/// stdout/stderr through `io::print*`.
25pub trait OutputHandler: Send + Sync {
26    fn write(&self, stream: OutputStream, text: &str) -> Result<(), String>;
27}
28
29/// Output handler that writes to the process stdout/stderr streams.
30#[derive(Debug, Default, Clone, Copy)]
31pub struct StdOutput;
32
33impl OutputHandler for StdOutput {
34    fn write(&self, stream: OutputStream, text: &str) -> Result<(), String> {
35        use std::io::Write;
36
37        match stream {
38            OutputStream::Stdout => {
39                let mut stdout = std::io::stdout().lock();
40                stdout.write_all(text.as_bytes()).map_err(|e| e.to_string())
41            }
42            OutputStream::Stderr => {
43                let mut stderr = std::io::stderr().lock();
44                stderr.write_all(text.as_bytes()).map_err(|e| e.to_string())
45            }
46        }
47    }
48}
49
50struct MissingOutputHandler;
51
52impl OutputHandler for MissingOutputHandler {
53    fn write(&self, _stream: OutputStream, _text: &str) -> Result<(), String> {
54        Err(ion_str!(
55            "io output handler is not configured; call Engine::set_output"
56        ))
57    }
58}
59
60pub(crate) fn missing_output_handler() -> Arc<dyn OutputHandler> {
61    Arc::new(MissingOutputHandler)
62}
63
64/// Build the `math` stdlib module.
65///
66/// Functions: abs, min, max, floor, ceil, round, sqrt, pow, clamp, log, log2, log10, sin, cos, tan, atan2
67/// Constants: PI, E, INF, NAN, TAU
68pub fn math_module() -> Module {
69    let mut m = Module::new("math");
70
71    // Constants
72    m.set("PI", Value::Float(std::f64::consts::PI));
73    m.set("E", Value::Float(std::f64::consts::E));
74    m.set("TAU", Value::Float(std::f64::consts::TAU));
75    m.set("INF", Value::Float(f64::INFINITY));
76    m.set("NAN", Value::Float(f64::NAN));
77
78    m.register_fn("abs", |args: &[Value]| {
79        if args.len() != 1 {
80            return Err(ion_str!("math::abs takes 1 argument"));
81        }
82        match &args[0] {
83            Value::Int(n) => Ok(Value::Int(n.abs())),
84            Value::Float(n) => Ok(Value::Float(n.abs())),
85            _ => Err(format!(
86                "{}{}",
87                ion_str!("math::abs not supported for "),
88                args[0].type_name()
89            )),
90        }
91    });
92
93    m.register_fn("min", |args: &[Value]| {
94        if args.len() < 2 {
95            return Err(ion_str!("math::min requires at least 2 arguments"));
96        }
97        let mut best = args[0].clone();
98        for arg in &args[1..] {
99            match (&best, arg) {
100                (Value::Int(a), Value::Int(b)) if b < a => best = arg.clone(),
101                (Value::Float(a), Value::Float(b)) if b < a => best = arg.clone(),
102                (Value::Int(a), Value::Float(b)) if *b < (*a as f64) => best = arg.clone(),
103                (Value::Float(a), Value::Int(b)) if (*b as f64) < *a => best = arg.clone(),
104                (Value::Int(_), Value::Int(_))
105                | (Value::Float(_), Value::Float(_))
106                | (Value::Int(_), Value::Float(_))
107                | (Value::Float(_), Value::Int(_)) => {}
108                _ => return Err(ion_str!("math::min requires numeric arguments")),
109            }
110        }
111        Ok(best)
112    });
113
114    m.register_fn("max", |args: &[Value]| {
115        if args.len() < 2 {
116            return Err(ion_str!("math::max requires at least 2 arguments"));
117        }
118        let mut best = args[0].clone();
119        for arg in &args[1..] {
120            match (&best, arg) {
121                (Value::Int(a), Value::Int(b)) if b > a => best = arg.clone(),
122                (Value::Float(a), Value::Float(b)) if b > a => best = arg.clone(),
123                (Value::Int(a), Value::Float(b)) if *b > (*a as f64) => best = arg.clone(),
124                (Value::Float(a), Value::Int(b)) if (*b as f64) > *a => best = arg.clone(),
125                (Value::Int(_), Value::Int(_))
126                | (Value::Float(_), Value::Float(_))
127                | (Value::Int(_), Value::Float(_))
128                | (Value::Float(_), Value::Int(_)) => {}
129                _ => return Err(ion_str!("math::max requires numeric arguments")),
130            }
131        }
132        Ok(best)
133    });
134
135    m.register_fn("floor", |args: &[Value]| match &args[0] {
136        Value::Float(n) => Ok(Value::Float(n.floor())),
137        Value::Int(n) => Ok(Value::Int(*n)),
138        _ => Err(format!(
139            "{}{}",
140            ion_str!("math::floor not supported for "),
141            args[0].type_name()
142        )),
143    });
144
145    m.register_fn("ceil", |args: &[Value]| match &args[0] {
146        Value::Float(n) => Ok(Value::Float(n.ceil())),
147        Value::Int(n) => Ok(Value::Int(*n)),
148        _ => Err(format!(
149            "{}{}",
150            ion_str!("math::ceil not supported for "),
151            args[0].type_name()
152        )),
153    });
154
155    m.register_fn("round", |args: &[Value]| match &args[0] {
156        Value::Float(n) => Ok(Value::Float(n.round())),
157        Value::Int(n) => Ok(Value::Int(*n)),
158        _ => Err(format!(
159            "{}{}",
160            ion_str!("math::round not supported for "),
161            args[0].type_name()
162        )),
163    });
164
165    m.register_fn("sqrt", |args: &[Value]| {
166        let n = args[0]
167            .as_float()
168            .ok_or(ion_str!("math::sqrt requires a number"))?;
169        Ok(Value::Float(n.sqrt()))
170    });
171
172    m.register_fn("pow", |args: &[Value]| {
173        if args.len() != 2 {
174            return Err(ion_str!("math::pow takes 2 arguments"));
175        }
176        match (&args[0], &args[1]) {
177            (Value::Int(base), Value::Int(exp)) => {
178                if *exp >= 0 {
179                    Ok(Value::Int(base.pow(*exp as u32)))
180                } else {
181                    Ok(Value::Float((*base as f64).powi(*exp as i32)))
182                }
183            }
184            _ => {
185                let b = args[0]
186                    .as_float()
187                    .ok_or(ion_str!("math::pow requires numeric arguments"))?;
188                let e = args[1]
189                    .as_float()
190                    .ok_or(ion_str!("math::pow requires numeric arguments"))?;
191                Ok(Value::Float(b.powf(e)))
192            }
193        }
194    });
195
196    m.register_fn("clamp", |args: &[Value]| {
197        if args.len() != 3 {
198            return Err(ion_str!(
199                "math::clamp requires 3 arguments: value, min, max"
200            ));
201        }
202        match (&args[0], &args[1], &args[2]) {
203            (Value::Int(v), Value::Int(lo), Value::Int(hi)) => Ok(Value::Int(*v.max(lo).min(hi))),
204            (Value::Float(v), Value::Float(lo), Value::Float(hi)) => {
205                Ok(Value::Float(v.max(*lo).min(*hi)))
206            }
207            _ => {
208                let v = args[0]
209                    .as_float()
210                    .ok_or(ion_str!("math::clamp requires numeric arguments"))?;
211                let lo = args[1]
212                    .as_float()
213                    .ok_or(ion_str!("math::clamp requires numeric arguments"))?;
214                let hi = args[2]
215                    .as_float()
216                    .ok_or(ion_str!("math::clamp requires numeric arguments"))?;
217                Ok(Value::Float(v.max(lo).min(hi)))
218            }
219        }
220    });
221
222    // Trigonometry
223    m.register_fn("sin", |args: &[Value]| {
224        let n = args[0]
225            .as_float()
226            .ok_or(ion_str!("math::sin requires a number"))?;
227        Ok(Value::Float(n.sin()))
228    });
229
230    m.register_fn("cos", |args: &[Value]| {
231        let n = args[0]
232            .as_float()
233            .ok_or(ion_str!("math::cos requires a number"))?;
234        Ok(Value::Float(n.cos()))
235    });
236
237    m.register_fn("tan", |args: &[Value]| {
238        let n = args[0]
239            .as_float()
240            .ok_or(ion_str!("math::tan requires a number"))?;
241        Ok(Value::Float(n.tan()))
242    });
243
244    m.register_fn("atan2", |args: &[Value]| {
245        if args.len() != 2 {
246            return Err(ion_str!("math::atan2 takes 2 arguments"));
247        }
248        let y = args[0]
249            .as_float()
250            .ok_or(ion_str!("math::atan2 requires numeric arguments"))?;
251        let x = args[1]
252            .as_float()
253            .ok_or(ion_str!("math::atan2 requires numeric arguments"))?;
254        Ok(Value::Float(y.atan2(x)))
255    });
256
257    // Logarithms
258    m.register_fn("log", |args: &[Value]| {
259        let n = args[0]
260            .as_float()
261            .ok_or(ion_str!("math::log requires a number"))?;
262        Ok(Value::Float(n.ln()))
263    });
264
265    m.register_fn("log2", |args: &[Value]| {
266        let n = args[0]
267            .as_float()
268            .ok_or(ion_str!("math::log2 requires a number"))?;
269        Ok(Value::Float(n.log2()))
270    });
271
272    m.register_fn("log10", |args: &[Value]| {
273        let n = args[0]
274            .as_float()
275            .ok_or(ion_str!("math::log10 requires a number"))?;
276        Ok(Value::Float(n.log10()))
277    });
278
279    // Rounding/check
280    m.register_fn("is_nan", |args: &[Value]| match &args[0] {
281        Value::Float(n) => Ok(Value::Bool(n.is_nan())),
282        Value::Int(_) => Ok(Value::Bool(false)),
283        _ => Err(format!(
284            "{}{}",
285            ion_str!("math::is_nan not supported for "),
286            args[0].type_name()
287        )),
288    });
289
290    m.register_fn("is_inf", |args: &[Value]| match &args[0] {
291        Value::Float(n) => Ok(Value::Bool(n.is_infinite())),
292        Value::Int(_) => Ok(Value::Bool(false)),
293        _ => Err(format!(
294            "{}{}",
295            ion_str!("math::is_inf not supported for "),
296            args[0].type_name()
297        )),
298    });
299
300    m
301}
302
303/// Build the `json` stdlib module.
304///
305/// Functions: encode, decode, pretty
306pub fn json_module() -> Module {
307    let mut m = Module::new("json");
308
309    m.register_fn("encode", |args: &[Value]| {
310        if args.len() != 1 {
311            return Err(ion_str!("json::encode takes 1 argument"));
312        }
313        let json = args[0].to_json();
314        Ok(Value::Str(json.to_string()))
315    });
316
317    m.register_fn("decode", |args: &[Value]| {
318        if args.len() != 1 {
319            return Err(ion_str!("json::decode takes 1 argument"));
320        }
321        let s = args[0]
322            .as_str()
323            .ok_or_else(|| ion_str!("json::decode requires a string"))?;
324        let json: serde_json::Value = serde_json::from_str(s)
325            .map_err(|e| format!("{}{}", ion_str!("json::decode error: "), e))?;
326        Ok(Value::from_json(json))
327    });
328
329    m.register_fn("pretty", |args: &[Value]| {
330        if args.len() != 1 {
331            return Err(ion_str!("json::pretty takes 1 argument"));
332        }
333        let json = args[0].to_json();
334        serde_json::to_string_pretty(&json)
335            .map(Value::Str)
336            .map_err(|e| format!("{}{}", ion_str!("json::pretty error: "), e))
337    });
338
339    #[cfg(feature = "msgpack")]
340    m.register_fn("msgpack_encode", |args: &[Value]| {
341        if args.len() != 1 {
342            return Err(ion_str!("json::msgpack_encode takes 1 argument"));
343        }
344        args[0].to_msgpack().map(Value::Bytes)
345    });
346
347    #[cfg(feature = "msgpack")]
348    m.register_fn("msgpack_decode", |args: &[Value]| {
349        if args.len() != 1 {
350            return Err(ion_str!("json::msgpack_decode takes 1 argument"));
351        }
352        let data = match &args[0] {
353            Value::Bytes(b) => b,
354            _ => return Err(ion_str!("json::msgpack_decode requires bytes")),
355        };
356        Value::from_msgpack(data)
357    });
358
359    m
360}
361
362fn format_output_args(args: &[Value]) -> String {
363    args.iter()
364        .map(|arg| arg.to_string())
365        .collect::<Vec<_>>()
366        .join(" ")
367}
368
369/// Build the `io` stdlib module.
370///
371/// Functions: print, println, eprintln
372pub fn io_module() -> Module {
373    io_module_with_output(missing_output_handler())
374}
375
376/// Build the `io` stdlib module with a host-provided output handler.
377pub fn io_module_with_output(output: Arc<dyn OutputHandler>) -> Module {
378    let mut m = Module::new("io");
379
380    let stdout = Arc::clone(&output);
381    m.register_closure("print", move |args: &[Value]| {
382        stdout.write(OutputStream::Stdout, &format_output_args(args))?;
383        Ok(Value::Unit)
384    });
385
386    let stdout = Arc::clone(&output);
387    m.register_closure("println", move |args: &[Value]| {
388        let mut text = format_output_args(args);
389        text.push('\n');
390        stdout.write(OutputStream::Stdout, &text)?;
391        Ok(Value::Unit)
392    });
393
394    m.register_closure("eprintln", move |args: &[Value]| {
395        let mut text = format_output_args(args);
396        text.push('\n');
397        output.write(OutputStream::Stderr, &text)?;
398        Ok(Value::Unit)
399    });
400
401    m
402}
403
404/// Build the `str` stdlib module.
405///
406/// Functions: join
407pub fn string_module() -> Module {
408    let mut m = Module::new("string");
409
410    m.register_fn("join", |args: &[Value]| {
411        if args.is_empty() || args.len() > 2 {
412            return Err(ion_str!(
413                "string::join requires 1-2 arguments: list, [separator]"
414            ));
415        }
416        let items = match &args[0] {
417            Value::List(items) => items,
418            _ => return Err(ion_str!("string::join requires a list as first argument")),
419        };
420        let sep = if args.len() > 1 {
421            args[1].as_str().unwrap_or("").to_string()
422        } else {
423            String::new()
424        };
425        let parts: Vec<String> = items.iter().map(|v| format!("{}", v)).collect();
426        Ok(Value::Str(parts.join(&sep)))
427    });
428
429    m
430}
431
432/// Register all stdlib modules in the given environment.
433pub fn register_stdlib(env: &mut crate::env::Env) {
434    register_stdlib_with_output(env, missing_output_handler());
435}
436
437/// Register all stdlib modules with a host-provided output handler.
438pub fn register_stdlib_with_output(env: &mut crate::env::Env, output: Arc<dyn OutputHandler>) {
439    let math = math_module();
440    env.define(math.name.clone(), math.to_value(), false);
441
442    let json = json_module();
443    env.define(json.name.clone(), json.to_value(), false);
444
445    let io = io_module_with_output(output);
446    env.define(io.name.clone(), io.to_value(), false);
447
448    let string_mod = string_module();
449    env.define(string_mod.name.clone(), string_mod.to_value(), false);
450}