Skip to main content

tatara_lisp_script/stdlib/
list_ext.rs

1//! Extended list operators beyond tatara-lisp-eval's primitives.
2//!
3//!   (take N XS)       → first N elements
4//!   (drop N XS)       → all elements after the first N
5//!   (range LO HI)     → integers [LO, HI)
6//!   (range LO HI STEP) → integers [LO, HI) stepping by STEP
7//!   (distinct XS)     → same order, duplicates removed
8//!   (flatten XS)      → one-level flatten
9//!   (zip XS YS)       → list of 2-lists
10//!   (concat XS YS …)  → append any number of lists
11//!   (nth N XS)        → 0-based; nil if out of range
12//!   (last XS)         → last element or nil
13
14use 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            // Value doesn't impl PartialEq directly (it holds Arc<dyn Any>
72            // for Foreign); compare via Debug-formatted string. Cheap +
73            // good-enough for script-level "remove duplicates by structure".
74            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}