summaryrefslogtreecommitdiff
path: root/matrix-stdlib/src/iter.rs
diff options
context:
space:
mode:
Diffstat (limited to 'matrix-stdlib/src/iter.rs')
-rw-r--r--matrix-stdlib/src/iter.rs544
1 files changed, 544 insertions, 0 deletions
diff --git a/matrix-stdlib/src/iter.rs b/matrix-stdlib/src/iter.rs
new file mode 100644
index 0000000..74056ce
--- /dev/null
+++ b/matrix-stdlib/src/iter.rs
@@ -0,0 +1,544 @@
+use std::{cell::RefCell, rc::Rc};
+use matrix::{iter, vm::Vm, value::Value, Result, unpack_varargs, unpack_args};
+use matrix_macros::native_func;
+use crate::{error, VmArgs};
+
+use Value as V;
+
+macro_rules! next {
+ ($vm:expr, $frame:expr, $iter:expr) => {
+ $vm.run_fn($frame, $iter.clone(), vec![])
+ };
+}
+
+#[native_func(1)]
+fn len(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [value] = unpack_args!(args);
+ let len = match value {
+ V::List(l) => l.len(),
+ V::Matrix(m) => m.values.len(),
+ V::String(s) => s.len(),
+ V::Table(t) => t.len(),
+ _ => 1
+ };
+ Ok(V::Int(len as i64))
+}
+
+#[native_func(1)]
+fn sum((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut res = next!(vm, frame, iter)?;
+ if res == Value::Nil {
+ return Ok(res)
+ };
+ loop {
+ let next = next!(vm, frame, iter)?;
+ if next == Value::Nil { break }
+ res = (res + next)?;
+ }
+ Ok(res)
+}
+
+#[native_func(0..)]
+fn range(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let ([], varargs) = unpack_varargs!(args);
+ let (min, max, step) = match varargs.as_slice() {
+ [V::Int(max)] => {
+ (0, *max, 1)
+ },
+ [V::Int(min), V::Int(max)] => {
+ (*min, *max, 1)
+ },
+ [V::Int(min), V::Int(max), V::Int(step)] => {
+ (*min, *max, *step)
+ },
+ _ => return error!("range takes 1 to 3 [Int]'s")
+ };
+
+ let i = RefCell::new(min);
+ Ok(iter!(move |_,_| {
+ let curr = *(i.borrow());
+ *(i.borrow_mut()) = curr + step;
+ if curr >= max {
+ return Ok(V::Nil)
+ } else {
+ return Ok(V::Int(curr))
+ }
+ }))
+}
+
+#[native_func(1)]
+fn iter(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [value] = unpack_args!(args);
+ Ok(value.into_iter()?)
+}
+
+#[native_func(1)]
+fn once(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [value] = unpack_args!(args);
+ let first = RefCell::new(false);
+ Ok(iter!(move |_,_| {
+ if *first.borrow() == false {
+ *first.borrow_mut() = true;
+ Ok(value.clone())
+ } else {
+ Ok(Value::Nil)
+ }
+ }))
+}
+
+#[native_func(1)]
+fn list((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut list = Vec::new();
+ loop {
+ let res = next!(vm, frame, iter)?;
+ if res == Value::Nil { break }
+ list.push(res);
+ }
+ Ok(Value::List(list.into()))
+}
+
+#[native_func(2)]
+fn map(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [fun, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Function(fun) = fun else {
+ return error!("map 1st arg must be a function")
+ };
+ Ok(iter!(move |(vm,frame),_| {
+ let input = next!(vm, frame, iter)?;
+ if input == Value::Nil {
+ return Ok(input)
+ }
+ vm.run_fn(frame, fun.clone(), vec![input])
+ }))
+}
+
+#[native_func(2)]
+fn fold((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [fun, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Function(fun) = fun else {
+ return error!("fold 1st arg must be a function")
+ };
+ let mut res = next!(vm, frame, iter)?;
+ if res == Value::Nil {
+ return Ok(Value::Nil)
+ };
+ loop {
+ let next = next!(vm, frame, iter)?;
+ if next == Value::Nil {
+ return Ok(res)
+ };
+ res = vm.run_fn(frame, fun.clone(), vec![res, next])?;
+ }
+}
+
+#[native_func(3)]
+fn foldi((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [init, fun, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Function(fun) = fun else {
+ return error!("foldi 1st arg must be a function")
+ };
+ let mut res = init;
+ loop {
+ let next = next!(vm, frame, iter)?;
+ if next == Value::Nil {
+ return Ok(res)
+ };
+ res = vm.run_fn(frame, fun.clone(), vec![res, next])?;
+ }
+}
+
+#[native_func(1)]
+fn count((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut len = 0;
+ loop {
+ let res = vm.run_fn(frame, iter.clone(), vec![])?;
+ if res == Value::Nil { break };
+ len += 1;
+ }
+ Ok(Value::Int(len))
+}
+
+#[native_func(3)]
+fn scan(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [init, fun, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Function(fun) = fun else {
+ return error!("scan 2nd arg must be a function")
+ };
+ let res = RefCell::new(init);
+ Ok(iter!(move |(vm,frame),_| {
+ let next = next!(vm, frame, iter)?;
+ if next == Value::Nil {
+ return Ok(Value::Nil)
+ };
+ let res_next = vm.run_fn(frame, fun.clone(), vec![res.borrow().clone(), next])?;
+ *res.borrow_mut() = res_next;
+ Ok(res.borrow().clone())
+ }))
+}
+
+#[native_func(2)]
+fn filter(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [fun, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Function(fun) = fun else {
+ return error!("filter 1st arg must be a function")
+ };
+ Ok(iter!(move |(vm,frame),_| {
+ loop {
+ let next = next!(vm, frame, iter)?;
+ if next == Value::Nil {
+ return Ok(Value::Nil)
+ }
+ let res = vm.run_fn(frame, fun.clone(), vec![next.clone()])?;
+ if !!res {
+ return Ok(next)
+ }
+ }
+ }))
+}
+
+#[native_func(0..)]
+fn chain(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let ([], iters) = unpack_varargs!(args);
+
+ let mut chaind = Vec::new();
+ for iter in iters {
+ chaind.push(iter.into_iter_fn()?);
+ }
+
+ chaind.reverse();
+ let chaind = RefCell::new(chaind);
+ Ok(iter!(move |(vm,frame), _| {
+ loop {
+ let curr = match chaind.borrow_mut().last() {
+ Some(iter) => iter.clone(),
+ None => return Ok(Value::Nil)
+ };
+ match vm.run_fn(frame, curr.clone(), vec![]) {
+ Ok(Value::Nil) => {
+ chaind.borrow_mut().pop();
+ continue;
+ },
+ v => return v
+ };
+ }
+ }))
+}
+
+#[native_func(1)]
+fn lines(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [str] = unpack_args!(args);
+ let Value::String(str) = str else {
+ return error!("lines arg must be a string")
+ };
+ let lines: Vec<Rc<str>> = str.split_inclusive("\n").map(|s| Rc::from(s)).collect();
+ let res = RefCell::new(lines.into_iter());
+ Ok(iter!(move |_,_| {
+ match res.borrow_mut().next() {
+ Some(line) => Ok(Value::String(line)),
+ None => Ok(Value::Nil)
+ }
+ }))
+}
+
+#[native_func(1)]
+fn skip((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [count, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Int(i) = count else {
+ return error!("skip count requires a int")
+ };
+ for _ in 0..i {
+ next!(vm, frame, iter)?;
+ }
+ Ok(Value::Iter(iter))
+}
+
+#[native_func(2)]
+fn alternate(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [il, ir] = unpack_args!(args);
+ let il = il.into_iter_fn()?;
+ let ir = ir.into_iter_fn()?;
+ let flag = RefCell::new(Some(0));
+ Ok(iter!(move |(vm, frame),_| {
+ let f = *flag.borrow();
+ match f {
+ Some(0) => {
+ let val = next!(vm, frame, il)?;
+ if val == Value::Nil {
+ *flag.borrow_mut() = None;
+ } else {
+ *flag.borrow_mut() = Some(1);
+ }
+ Ok(val)
+ },
+ Some(1) => {
+ let val = next!(vm, frame, ir)?;
+ if val == Value::Nil {
+ *flag.borrow_mut() = None;
+ } else {
+ *flag.borrow_mut() = Some(0);
+ }
+ Ok(val)
+ },
+ _ => Ok(Value::Nil)
+ }
+ }))
+}
+
+#[native_func(2)]
+fn intersperse(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [value, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let flag = RefCell::new(Some(0));
+ Ok(iter!(move |(vm, frame),_| {
+ let f = *flag.borrow();
+ match f {
+ Some(0) => {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil {
+ *flag.borrow_mut() = None;
+ } else {
+ *flag.borrow_mut() = Some(1);
+ }
+ Ok(val)
+ },
+ Some(1) => {
+ let val = value.clone();
+ *flag.borrow_mut() = Some(0);
+ Ok(val)
+ },
+ _ => Ok(Value::Nil)
+ }
+ }))
+}
+
+#[native_func(2)]
+fn zip(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [il, ir] = unpack_args!(args);
+ let il = il.into_iter_fn()?;
+ let ir = ir.into_iter_fn()?;
+ let flag = RefCell::new(true);
+ Ok(iter!(move |(vm, frame),_| {
+ let f = *flag.borrow();
+ match f {
+ true => {
+ let vl = next!(vm, frame, il)?;
+ let vr = next!(vm, frame, ir)?;
+ if vl == Value::Nil || vr == Value::Nil {
+ *flag.borrow_mut() = false;
+ }
+ Ok(Value::List(vec![vl, vr].into()))
+ },
+ _ => Ok(Value::Nil)
+ }
+ }))
+}
+
+#[native_func(1)]
+fn unzip((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut ll = Vec::new();
+ let mut lr = Vec::new();
+ loop {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil {
+ break;
+ }
+ let Value::List(vals) = val else {
+ return error!("unzip only works over a iterator of pairs");
+ };
+ let vals = vals.into_inner();
+ if vals.len() != 2 {
+ return error!("unzip only works over a iterator of pairs");
+ }
+ let [l, r] = vals.try_into().unwrap();
+ ll.push(l);
+ lr.push(r);
+ }
+ let ll = Value::List(ll.into());
+ let lr = Value::List(lr.into());
+ Ok(Value::List(vec![ll, lr].into()))
+}
+
+#[native_func(1)]
+fn cycle((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut values = Vec::new();
+ loop {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil { break };
+ values.push(val);
+ }
+ let idx = RefCell::new(0_usize);
+ Ok(iter!(move |_,_| {
+ let i = *idx.borrow();
+ *idx.borrow_mut() += 1;
+ *idx.borrow_mut() %= values.len();
+ Ok(values[i].clone())
+ }))
+}
+
+#[native_func(2)]
+fn take((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [count, iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let Value::Int(count) = count else {
+ return error!("take requires an int amount to collect")
+ };
+ let mut values = Vec::new();
+ for _ in 0..count {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil { break };
+ values.push(val);
+ }
+ Ok(Value::List(values.into()))
+}
+
+#[native_func(2)]
+fn last((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut last = Value::Nil;
+ loop {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil { break };
+ last = val;
+ }
+ Ok(last)
+}
+
+#[native_func(2)]
+fn next((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let val = next!(vm, frame, iter)?;
+ Ok(val)
+}
+
+#[native_func(1)]
+fn min((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut min = Value::Nil;
+ loop {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil { break };
+ if min == Value::Nil || val < min {
+ min = val;
+ }
+ }
+ Ok(min)
+}
+
+#[native_func(1)]
+fn max((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut max = Value::Nil;
+ loop {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil { break };
+ if max == Value::Nil || val > max {
+ max = val;
+ }
+ }
+ Ok(max)
+}
+
+#[native_func(1)]
+fn rev((vm, frame): VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+ let mut values = Vec::new();
+ loop {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil { break };
+ values.push(val);
+ }
+ let values = RefCell::new(values);
+ Ok(iter!(move |_,_| {
+ match values.borrow_mut().pop() {
+ Some(v) => Ok(v),
+ None => Ok(Value::Nil)
+ }
+ }))
+}
+
+#[native_func(1)]
+fn enumerate(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [iter] = unpack_args!(args);
+ let iter = iter.into_iter_fn()?;
+
+ let idx = RefCell::new(0_i64);
+ Ok(iter!(move |(vm, frame),_| {
+ let val = next!(vm, frame, iter)?;
+ if val == Value::Nil {
+ return Ok(Value::Nil)
+ };
+ let curr = *idx.borrow();
+ *idx.borrow_mut() += 1;
+ Ok(Value::List(vec![Value::Int(curr), val].into()))
+ }))
+}
+
+#[native_func(1)]
+fn iterable(_: VmArgs, args: Vec<Value>) -> Result<Value> {
+ let [value] = unpack_args!(args);
+ use Value as V;
+ match value {
+ V::Function(_) |
+ V::List(_) |
+ V::Range(_) |
+ V::Iter(_)
+ => Ok(V::Bool(true)),
+ _ => Ok(V::Bool(false))
+ }
+}
+
+pub fn load(vm: &mut Vm) {
+ vm.load_global_fn(take(), "take");
+ vm.load_global_fn(unzip(), "unzip");
+ vm.load_global_fn(count(), "count");
+ vm.load_global_fn(len(), "len");
+ vm.load_global_fn(sum(), "sum");
+ vm.load_global_fn(min(), "min");
+ vm.load_global_fn(max(), "max");
+ vm.load_global_fn(next(), "next");
+ vm.load_global_fn(last(), "last");
+
+ vm.load_global_fn(list(), "list");
+ vm.load_global_fn(fold(), "fold");
+ vm.load_global_fn(foldi(), "foldi");
+ vm.load_global_fn(scan(), "scan");
+ vm.load_global_fn(chain(), "chain");
+ vm.load_global_fn(lines(), "lines");
+ vm.load_global_fn(skip(), "skip");
+
+ vm.load_global_fn(once(), "once");
+ vm.load_global_fn(iter(), "iter");
+ vm.load_global_fn(range(), "range");
+ vm.load_global_fn(map(), "map");
+ vm.load_global_fn(filter(), "filter");
+ vm.load_global_fn(skip(), "skip");
+ vm.load_global_fn(zip(), "zip");
+ vm.load_global_fn(cycle(), "cycle");
+ vm.load_global_fn(alternate(), "alternate");
+ vm.load_global_fn(intersperse(), "intersperse");
+ vm.load_global_fn(rev(), "rev");
+ vm.load_global_fn(enumerate(), "enumerate");
+
+ vm.load_global_fn(iterable(), "iterable");
+}
generated by cgit v1.2.3-freya (git 2.51.0) at 2025-10-15 14:43:12 +0000