talc/talc-std/src/iter.rs

use std::{cell::RefCell, collections::HashMap, rc::Rc};

use talc_lang::{symbol::SYM_TYPE_ERROR, exception::Result, throw, value::{function::NativeFunc, range::RangeType, HashValue, Value}, vmcall, vmcalliter, Vm};
use talc_macros::native_func;

use crate::{unpack_args, unpack_varargs};

pub fn load(vm: &mut Vm) {
	// begin
	vm.set_global_name("iter", iter().into());
	vm.set_global_name("pairs", pairs().into());
	vm.set_global_name("once", once().into());
	vm.set_global_name("forever", forever().into());

	// modify
	vm.set_global_name("map", map().into());
	vm.set_global_name("scan", scan().into());
	vm.set_global_name("tee", tee().into());
	vm.set_global_name("filter", filter().into());
	vm.set_global_name("take", take().into());
	vm.set_global_name("skip", skip().into());
	vm.set_global_name("enumerate", enumerate().into());
	vm.set_global_name("cycle", cycle().into());
	vm.set_global_name("step", step().into());
	vm.set_global_name("rev", rev().into());

	// join
	vm.set_global_name("zip", zip().into());
	vm.set_global_name("alternate", alternate().into());
	vm.set_global_name("intersperse", intersperse().into());
	vm.set_global_name("cartprod", cartprod().into());
	vm.set_global_name("chain", chain().into());

	// end
	vm.set_global_name("list", list().into());
	vm.set_global_name("table", table().into());
	vm.set_global_name("len", len().into());
	vm.set_global_name("fold", fold().into());
	vm.set_global_name("foldi", foldi().into());
	vm.set_global_name("sum", sum().into());
	vm.set_global_name("prod", prod().into());
	vm.set_global_name("any", any().into());
	vm.set_global_name("all", all().into());
	vm.set_global_name("last", last().into());
	vm.set_global_name("nth", nth().into());
	vm.set_global_name("contains", contains().into());
	vm.set_global_name("index_of", index_of().into());
	vm.set_global_name("index_if", index_if().into());
	vm.set_global_name("find", find().into());
	vm.set_global_name("count", count().into());
}

//
//  begin iteration
//

#[native_func(1)]
pub fn iter(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, v] = unpack_args!(args);

	v.to_iter_function()
}

#[native_func(1)]
pub fn pairs(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, table] = unpack_args!(args);
	let Value::Table(table) = table else {
		throw!(*SYM_TYPE_ERROR, "pairs expected table, found {table:#}")
	};

	let keys: Vec<HashValue> = table.borrow().keys().cloned().collect();
	let keys = RefCell::new(keys.into_iter());

	let f = move |_: &mut Vm, _: Vec<Value>| -> Result<Value> {
		if let Some(k) = keys.borrow_mut().next() {
			let v = table.borrow().get(&k).cloned().unwrap_or(Value::Nil);
			Ok(Value::from(vec![k.into_inner(), v]))
		} else {
			Ok(Value::iter_pack(None))
		}
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(1)]
pub fn once(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, val] = unpack_args!(args);

	let v = RefCell::new(Some(val));
	let f = move |_: &mut Vm, _| {
		Ok(Value::iter_pack(v.borrow_mut().take()))
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(1)]
pub fn forever(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, val] = unpack_args!(args);

	let f = move |_: &mut Vm, _| {
		Ok(val.clone())
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

//
//  chain iteration
//



#[native_func(2)]
pub fn map(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, map, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let f = move |vm: &mut Vm, _| {
		match vmcalliter!(vm; iter.clone())? {
			Some(val) => Ok(vmcall!(vm; map.clone(), val)?),
			None => Ok(Value::iter_pack(None)),
		}
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(2)]
pub fn tee(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, tee, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let f = move |vm: &mut Vm, _| {
		match vmcalliter!(vm; iter.clone())? {
			Some(val) => {
				vmcall!(vm; tee.clone(), val.clone())?;
				Ok(val.into())
			}
			None => Ok(Value::iter_pack(None)),
		}
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(3)]
pub fn scan(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, init, func, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let result = RefCell::new(init);
	let f = move |vm: &mut Vm, _| {
		let Some(val) = vmcalliter!(vm; iter.clone())? else {
			result.take();
			return Ok(Value::iter_pack(None))
		};
		let r = vmcall!(vm; func.clone(), result.take(), val)?;
		*result.borrow_mut() = r.clone();
		Ok(r)
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(2)]
pub fn filter(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, filter, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let f = move |vm: &mut Vm, _| {
		loop {
			let Some(next) = vmcalliter!(vm; iter.clone())? else {
				return Ok(Value::iter_pack(None))
			};
			let res = vmcall!(vm; filter.clone(), next.clone())?;
			if res.truthy() {
				return Ok(next.into())
			}
		}
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(2)]
pub fn take(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, count, iter] = unpack_args!(args);
	let Value::Int(count) = count else {
		throw!(*SYM_TYPE_ERROR, "take expected integer")
	};
	let Ok(count) = count.try_into() else {
		throw!(*SYM_TYPE_ERROR, "take expected nonnegative integer")
	};
	let iter = iter.to_iter_function()?;

	let taken = RefCell::new(0);
	let f = move |vm: &mut Vm, _| {
		if *taken.borrow() >= count {
			return Ok(Value::iter_pack(None))
		}
		let Some(next) = vmcalliter!(vm; iter.clone())? else {
			*taken.borrow_mut() = count;
			return Ok(Value::iter_pack(None))
		};
		*taken.borrow_mut() += 1;
		Ok(next.into())
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(2)]
pub fn skip(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, count, iter] = unpack_args!(args);
	let Value::Int(count) = count else {
		throw!(*SYM_TYPE_ERROR, "count expected integer")
	};
	let Ok(count) = count.try_into() else {
		throw!(*SYM_TYPE_ERROR, "count expected nonnegative integer")
	};
	let iter = iter.to_iter_function()?;

	let skipped = RefCell::new(false);
	let f = move |vm: &mut Vm, _| {
		if *skipped.borrow() {
			return vmcall!(vm; iter.clone())
		}
		*skipped.borrow_mut() = true;
		for _ in 0..count {
			vmcall!(vm; iter.clone())?;
		}
		vmcall!(vm; iter.clone())
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(1)]
pub fn enumerate(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let counter = RefCell::new(0);
	let f = move |vm: &mut Vm, _| {
		let Some(next) = vmcalliter!(vm; iter.clone())? else {
			return Ok(Value::iter_pack(None))
		};
		let mut c = counter.borrow_mut();
		let n = *c;
		*c += 1;
		Ok(Value::from(vec![n.into(), next]))
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}


#[native_func(1)]
pub fn cycle(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let record = RefCell::new((Vec::<Value>::new(), false, 0));
	let f = move |vm: &mut Vm, _| {
		if record.borrow().1 {
			let mut r = record.borrow_mut();
			if r.0.is_empty() {
				return Ok(Value::iter_pack(None))
			}
			let val = r.0[r.2].clone();
			r.2 = (r.2 + 1) % r.0.len();
			return Ok(val.into())
		}
		match vmcalliter!(vm; iter.clone())? {
			Some(v) => {
				record.borrow_mut().0.push(v.clone());
				Ok(v.into())
			},
			None => {
				let mut r = record.borrow_mut();
				r.1 = true;
				if r.0.is_empty() {
					Ok(Value::iter_pack(None))
				} else {
					r.2 = (r.2 + 1) % r.0.len();
					Ok(r.0[0].clone().into())
				}
			}
		}
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[derive(Clone, Copy, Default)]
enum Step {
	First, Going, #[default] End,
}

#[native_func(2)]
pub fn step(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, by, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;
	let Value::Int(by) = by else {
		throw!(*SYM_TYPE_ERROR, "step expected integer")
	};
	if by <= 0 {
		throw!(*SYM_TYPE_ERROR, "step expected positive integer")
	}

	let state = RefCell::new(Step::First);
	let f = move |vm: &mut Vm, _| match state.take() {
		Step::First => {
			let res = vmcalliter!(vm; iter.clone())?;
			if res.is_some() {
				*state.borrow_mut() = Step::Going;
			}
			Ok(Value::iter_pack(res))
		},
		Step::Going => {
			for _ in 0..(by-1) {
				if vmcall!(vm; iter.clone())? == Value::Nil {
					return Ok(Value::iter_pack(None))
				}
			}
			let Some(x) = vmcalliter!(vm; iter.clone())? else {
				return Ok(Value::iter_pack(None))
			};
			*state.borrow_mut() = Step::Going;
			Ok(x)
		},
		Step::End => Ok(Value::Nil),
	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(1)]
pub fn rev(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let lst = RefCell::new(None);
	let f = move |vm: &mut Vm, _| {
		if lst.borrow().is_none() {
			let mut l = Vec::new();
			while let Some(x) = vmcalliter!(vm; iter.clone())? {
				l.push(x)
			}
			*lst.borrow_mut() = Some(l);
		}
		Ok(Value::iter_pack(lst.borrow_mut().as_mut().unwrap().pop()))

	};
	Ok(NativeFunc::new(Box::new(f), 0).into())
}


//
//  join
//


#[native_func(2..)]
pub fn zip(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let ([_, i1, i2], rest) = unpack_varargs!(args);
	let mut iters = Vec::with_capacity(2 + rest.len());
	for i in [i1, i2].into_iter().chain(rest.into_iter()) {
		iters.push(i.to_iter_function()?);
	}

	let f = move |vm: &mut Vm, _| {
		let mut res = Vec::with_capacity(iters.len());
		for i in iters.iter() {
			match vmcalliter!(vm; i.clone())? {
				Some(v) => res.push(v),
				None => return Ok(Value::iter_pack(None)),
			};
		}
		Ok(Value::from(res))
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[native_func(2..)]
pub fn alternate(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let ([_, i1, i2], rest) = unpack_varargs!(args);
	let mut iters = Vec::with_capacity(2 + rest.len());
	for i in [i1, i2].into_iter().chain(rest.into_iter()) {
		iters.push(i.to_iter_function()?);
	}

	let state = RefCell::new((0, false));
	let f = move |vm: &mut Vm, _| {
		let mut s = state.borrow_mut();
		if s.1 {
			return Ok(Value::iter_pack(None));
		}
		let n = s.0;
		s.0 = (s.0 + 1) % iters.len();
		drop(s);
		match vmcalliter!(vm; iters[n].clone())? {
			Some(v) => Ok(v),
			None => {
				state.borrow_mut().1 = true;
				Ok(Value::iter_pack(None))
			}
		}
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[derive(Default)]
enum Intersperse {
	Init, Waiting, HasNext(Value), #[default] End
}

#[native_func(2)]
pub fn intersperse(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, val, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let state = RefCell::new(Intersperse::Init);
	let f = move |vm: &mut Vm, _| {
		match state.take() {
			Intersperse::Init => match vmcalliter!(vm; iter.clone())? {
				Some(v) => {
					*state.borrow_mut() = Intersperse::Waiting;
					Ok(v)
				},
				None => {
					*state.borrow_mut() = Intersperse::End;
					Ok(Value::iter_pack(None))
				},
			},
			Intersperse::Waiting => match vmcalliter!(vm; iter.clone())? {
				Some(v) => {
					*state.borrow_mut() = Intersperse::HasNext(v);
					Ok(val.clone())
				},
				None => {
					*state.borrow_mut() = Intersperse::End;
					Ok(Value::iter_pack(None))
				},
			},
			Intersperse::HasNext(v) => {
				*state.borrow_mut() = Intersperse::Waiting;
				Ok(v)
			},
			Intersperse::End => Ok(Value::iter_pack(None)),
		}
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}


#[native_func(2)]
pub fn chain(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter1, iter2] = unpack_args!(args);
	let iter1 = iter1.to_iter_function()?;
	let iter2 = iter2.to_iter_function()?;

	let done_first = RefCell::new(false);
	let f = move |vm: &mut Vm, _| {
		if *done_first.borrow() {
			return vmcall!(vm; iter2.clone())
		}
		match vmcalliter!(vm; iter1.clone())? {
			Some(v) => Ok(v),
			None => {
				*done_first.borrow_mut() = true;
				vmcall!(vm; iter2.clone())
			}
		}
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}

#[derive(Default, Debug)]
struct CartProd {
	a_val: Value,
	b_data: Vec<Value>,
	b_idx: usize,
	b_done: bool,
	done: bool,
}

#[native_func(2)]
pub fn cartprod(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, a, b] = unpack_args!(args);
	let a = a.to_iter_function()?;
	let b = b.to_iter_function()?;

	let state = RefCell::new(CartProd::default());
	let f = move |vm: &mut Vm, _| {
		if state.borrow().done {
			return Ok(Value::Nil)
		}

		if state.borrow().b_idx >= state.borrow().b_data.len() {
			if !state.borrow().b_done {
				let v = vmcalliter!(vm; b.clone())?;
				let mut s = state.borrow_mut();
				match v {
					Some(x) => s.b_data.push(x),
					None => {
						s.b_done = true;
						if s.b_idx == 0 {
							s.done = true;
							return Ok(Value::iter_pack(None))
						}
						s.b_idx = 0;
					}
				};
			} else {
				if state.borrow().b_idx == 0 {
					state.borrow_mut().done = true;
					return Ok(Value::Nil)
				}
				state.borrow_mut().b_idx = 0;
			}
		}

		let b_res = state.borrow().b_data[state.borrow().b_idx].clone();

		if state.borrow().b_idx == 0 {
			match vmcalliter!(vm; a.clone())? {
				Some(v) => state.borrow_mut().a_val = v,
				None => {
					state.borrow_mut().done = true;
					return Ok(Value::iter_pack(None))
				}
			}
		}

		let a_res = state.borrow().a_val.clone();

		state.borrow_mut().b_idx += 1;

		Ok(Value::from(vec![a_res, b_res]))
	};

	Ok(NativeFunc::new(Box::new(f), 0).into())
}



//
//  end iteration
//



#[native_func(1)]
pub fn list(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;
	let mut result = Vec::new();
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		result.push(value);
	};
	Ok(result.into())
}

#[native_func(1)]
pub fn table(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;
	let mut result = HashMap::new();
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		let Value::List(l) = value else {
			throw!(*SYM_TYPE_ERROR, "table expected iterator to yield list")
		};
		let l = Rc::unwrap_or_clone(l).take();
		let Ok([k, v]): std::result::Result<[Value; 2], Vec<Value>> = l.try_into() else {
			throw!(*SYM_TYPE_ERROR, "table expected iterator to yield list of length 2")
		};
		result.insert(k.try_into()?, v);
	};
	Ok(result.into())
}

#[native_func(1)]
pub fn len(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, value] = unpack_args!(args);
	match value {
		Value::String(s) => return Ok((s.len() as i64).into()),
		Value::List(l) => return Ok((l.borrow().len() as i64).into()),
		Value::Table(t) => return Ok((t.borrow().len() as i64).into()),
		Value::Range(r) if r.ty != RangeType::Endless
			=> return Ok((r.len().unwrap() as i64).into()),
		_ => (),
	}
	let iter = value.to_iter_function()?;
	let mut len = 0;
	while vmcalliter!(vm; iter.clone())?.is_some() {
		len += 1;
	};
	Ok(len.into())
}

#[native_func(2)]
pub fn fold(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, func, iter] = unpack_args!(args);
	let iter: Value = iter.to_iter_function()?;

	let mut result = match vmcalliter!(vm; iter.clone())? {
		Some(v) => v,
		None => return Ok(Value::iter_pack(None)),
	};
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		result = vmcall!(vm; func.clone(), result, value)?;
	}
	Ok(result)
}

#[native_func(3)]
pub fn foldi(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, init, func, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let mut result = init;
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		result = vmcall!(vm; func.clone(), result, value)?;
	}
	Ok(result)
}

#[native_func(1)]
pub fn sum(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let mut result = Value::Int(0);
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		result = (result + value)?
	}
	Ok(result)
}

#[native_func(1)]
pub fn prod(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let mut result = Value::Int(1);
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		result = (result * value)?
	}
	Ok(result)
}

#[native_func(1)]
pub fn any(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		if value.truthy() {
			return Ok(true.into())
		}
	}
	Ok(false.into())
}

#[native_func(1)]
pub fn all(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		if !value.truthy() {
			return Ok(false.into())
		}
	}
	Ok(true.into())
}

#[native_func(1)]
pub fn last(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let mut last = Value::Nil;
	while let Some(value) = vmcalliter!(vm; iter.clone())? {
		last = value;
	}
	Ok(last)
}

#[native_func(2)]
pub fn nth(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, n, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;
	let Value::Int(n) = n else {
		throw!(*SYM_TYPE_ERROR, "nth expected integer")
	};

	for _ in 0..n {
		if vmcalliter!(vm; iter.clone())?.is_none() {
			return Ok(Value::Nil)
		}
	}
	match vmcalliter!(vm; iter)? {
		Some(v) => Ok(v),
		None => Ok(Value::Nil),
	}
}

#[native_func(2)]
pub fn contains(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, val, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	for _ in 0.. {
		match vmcalliter!(vm; iter.clone())? {
			None => return Ok(Value::Nil),
			Some(v) => if v == val {
				return Ok(true.into())
			}
		}
	}
	Ok(false.into())
}

#[native_func(2)]
pub fn index_of(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, val, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	for i in 0.. {
		match vmcalliter!(vm; iter.clone())? {
			None => return Ok(Value::Nil),
			Some(v) => if v == val {
				return Ok(i.into())
			}
		}
	}
	Ok(Value::Nil)
}

#[native_func(2)]
pub fn index_if(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, func, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	for i in 0.. {
		match vmcalliter!(vm; iter.clone())? {
			None => return Ok(Value::Nil),
			Some(v) => if vmcall!(vm; func.clone(), v)?.truthy() {
				return Ok(i.into())
			}
		}
	}
	Ok(Value::Nil)
}

#[native_func(2)]
pub fn find(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, func, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	loop {
		match vmcalliter!(vm; iter.clone())? {
			None => return Ok(Value::Nil),
			Some(v) => {
				if vmcall!(vm; func.clone(), v.clone())?.truthy() {
					return Ok(v)
				}
			}
		}
	}
}

#[native_func(1)]
pub fn count(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, iter] = unpack_args!(args);
	let iter = iter.to_iter_function()?;

	let mut map = HashMap::new();

	while let Some(v) = vmcalliter!(vm; iter.clone())? {
		let hv = v.try_into()?;
		map.entry(hv)
			.and_modify(|v: &mut Value| *v = (v.clone() + Value::Int(1)).unwrap())
			.or_insert(Value::Int(1));
	}

	Ok(map.into())
}