1use std::sync::Arc;
9
10use crate::module::Module;
11use crate::value::Value;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum OutputStream {
16 Stdout,
17 Stderr,
18}
19
20pub trait OutputHandler: Send + Sync {
26 fn write(&self, stream: OutputStream, text: &str) -> Result<(), String>;
27}
28
29#[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
64pub fn math_module() -> Module {
69 let mut m = Module::new("math");
70
71 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 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 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 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
303pub 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
369pub fn io_module() -> Module {
373 io_module_with_output(missing_output_handler())
374}
375
376pub 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
404pub 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
432pub fn register_stdlib(env: &mut crate::env::Env) {
434 register_stdlib_with_output(env, missing_output_handler());
435}
436
437pub 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}