tatara_lisp_script/stdlib/
list_ext.rs1use std::sync::Arc;
15
16use tatara_lisp_eval::{Arity, EvalError, Interpreter, Value};
17
18use crate::script_ctx::ScriptCtx;
19use crate::stdlib::env::int_arg;
20
21pub fn install(interp: &mut Interpreter<ScriptCtx>) {
22 interp.register_fn(
23 "take",
24 Arity::Exact(2),
25 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
26 let n = int_arg(&args[0], "take", sp)?.max(0) as usize;
27 let xs = list_arg(&args[1], "take", sp)?;
28 Ok(Value::list(xs.iter().take(n).cloned().collect::<Vec<_>>()))
29 },
30 );
31
32 interp.register_fn(
33 "drop",
34 Arity::Exact(2),
35 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
36 let n = int_arg(&args[0], "drop", sp)?.max(0) as usize;
37 let xs = list_arg(&args[1], "drop", sp)?;
38 Ok(Value::list(xs.iter().skip(n).cloned().collect::<Vec<_>>()))
39 },
40 );
41
42 interp.register_fn(
43 "range",
44 Arity::Range(2, 3),
45 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
46 let lo = int_arg(&args[0], "range", sp)?;
47 let hi = int_arg(&args[1], "range", sp)?;
48 let step = args
49 .get(2)
50 .map(|v| int_arg(v, "range", sp))
51 .transpose()?
52 .unwrap_or(1);
53 if step == 0 {
54 return Err(EvalError::native_fn("range", "step cannot be 0", sp));
55 }
56 let mut out = Vec::new();
57 let mut i = lo;
58 while (step > 0 && i < hi) || (step < 0 && i > hi) {
59 out.push(Value::Int(i));
60 i += step;
61 }
62 Ok(Value::list(out))
63 },
64 );
65
66 interp.register_fn(
67 "distinct",
68 Arity::Exact(1),
69 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
70 let xs = list_arg(&args[0], "distinct", sp)?;
71 let mut seen_keys: Vec<String> = Vec::new();
75 let mut out: Vec<Value> = Vec::new();
76 for item in xs.iter() {
77 let key = format!("{item:?}");
78 if !seen_keys.contains(&key) {
79 seen_keys.push(key);
80 out.push(item.clone());
81 }
82 }
83 Ok(Value::list(out))
84 },
85 );
86
87 interp.register_fn(
88 "flatten",
89 Arity::Exact(1),
90 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
91 let xs = list_arg(&args[0], "flatten", sp)?;
92 let mut out = Vec::new();
93 for item in xs.iter() {
94 match item {
95 Value::List(inner) => out.extend(inner.iter().cloned()),
96 Value::Nil => {}
97 other => out.push(other.clone()),
98 }
99 }
100 Ok(Value::list(out))
101 },
102 );
103
104 interp.register_fn(
105 "zip",
106 Arity::Exact(2),
107 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
108 let xs = list_arg(&args[0], "zip", sp)?;
109 let ys = list_arg(&args[1], "zip", sp)?;
110 let out: Vec<Value> = xs
111 .iter()
112 .zip(ys.iter())
113 .map(|(a, b)| Value::list(vec![a.clone(), b.clone()]))
114 .collect();
115 Ok(Value::list(out))
116 },
117 );
118
119 interp.register_fn(
120 "concat",
121 Arity::AtLeast(0),
122 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
123 let mut out = Vec::new();
124 for v in args {
125 let xs = list_arg(v, "concat", sp)?;
126 out.extend(xs.iter().cloned());
127 }
128 Ok(Value::list(out))
129 },
130 );
131
132 interp.register_fn(
133 "nth",
134 Arity::Exact(2),
135 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
136 let n = int_arg(&args[0], "nth", sp)?;
137 let xs = list_arg(&args[1], "nth", sp)?;
138 if n < 0 {
139 return Ok(Value::Nil);
140 }
141 Ok(xs.get(n as usize).cloned().unwrap_or(Value::Nil))
142 },
143 );
144
145 interp.register_fn(
146 "last",
147 Arity::Exact(1),
148 |args: &[Value], _ctx: &mut ScriptCtx, sp| {
149 let xs = list_arg(&args[0], "last", sp)?;
150 Ok(xs.last().cloned().unwrap_or(Value::Nil))
151 },
152 );
153}
154
155fn list_arg(
156 v: &Value,
157 fname: &'static str,
158 sp: tatara_lisp::Span,
159) -> Result<Arc<Vec<Value>>, EvalError> {
160 match v {
161 Value::List(xs) => Ok(xs.clone()),
162 Value::Nil => Ok(Arc::new(Vec::new())),
163 other => Err(EvalError::native_fn(
164 fname,
165 format!("expected list, got {}", other.type_name()),
166 sp,
167 )),
168 }
169}