talc/talc-std/src/collection.rs

use std::cmp::Ordering;

use talc_lang::{
	exception::{exception, Result},
	number::Int,
	prelude::*,
	symbol::{symbol, SYM_TYPE_ERROR, SYM_VALUE_ERROR},
	throw,
	value::{function::NativeFunc, Value},
	vm::Vm,
	vmcall,
};
use talc_macros::native_func;

use crate::unpack_args;

pub fn load(vm: &mut Vm) {
	vm.set_global_name("push", push().into());
	vm.set_global_name("pop", pop().into());
	vm.set_global_name("reverse", reverse().into());
	vm.set_global_name("clear", clear().into());
	vm.set_global_name("sort", sort().into());
	vm.set_global_name("sort_by", sort_by().into());
	vm.set_global_name("sort_key", sort_key().into());
	vm.set_global_name("pair", pair().into());
	vm.set_global_name("fst", fst().into());
	vm.set_global_name("snd", snd().into());
}

#[native_func(2)]
pub fn push(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, list, item] = unpack_args!(args);
	let Value::List(list) = list else {
		throw!(*SYM_TYPE_ERROR, "push expected list, found {list:#}")
	};
	list.borrow_mut().push(item);
	Ok(Value::Nil)
}

#[native_func(1)]
pub fn pop(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, list] = unpack_args!(args);
	let Value::List(list) = list else {
		throw!(*SYM_TYPE_ERROR, "pop expected list, found {list:#}")
	};
	let v = list.borrow_mut().pop();
	v.ok_or_else(|| exception!(*SYM_VALUE_ERROR, "attempt to pop empty list"))
}

#[native_func(1)]
pub fn reverse(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, list] = unpack_args!(args);
	let Value::List(list) = list else {
		throw!(*SYM_TYPE_ERROR, "reversed expected list, found {list:#}")
	};
	list.borrow_mut().reverse();
	Ok(Value::List(list))
}

#[native_func(1)]
pub fn clear(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, col] = unpack_args!(args);
	match &col {
		Value::List(list) => list.borrow_mut().clear(),
		Value::Table(table) => table.borrow_mut().clear(),
		_ => throw!(
			*SYM_TYPE_ERROR,
			"clear expected list or table, found {col:#}"
		),
	}
	Ok(col)
}

fn call_comparison(vm: &mut Vm, by: Value, l: Value, r: Value) -> Result<Int> {
	let ord = vmcall!(vm; by, l, r)?;
	let Value::Int(ord) = ord else {
		throw!(
			*SYM_TYPE_ERROR,
			"comparison function should return an integer"
		)
	};
	Ok(ord)
}

fn partition(vals: &mut [Value]) -> (usize, usize) {
	let pivot = vals[vals.len() / 2].clone();

	let mut lt = 0;
	let mut eq = 0;
	let mut gt = vals.len() - 1;

	while eq <= gt {
		let ord = vals[eq].partial_cmp(&pivot);
		match ord {
			Some(Ordering::Less) => {
				vals.swap(eq, lt);
				lt += 1;
				eq += 1;
			}
			Some(Ordering::Greater) => {
				vals.swap(eq, gt);
				gt -= 1;
			}
			Some(Ordering::Equal) | None => {
				eq += 1;
			}
		}
	}

	(lt, gt)
}

fn partition_by(vals: &mut [Value], by: &Value, vm: &mut Vm) -> Result<(usize, usize)> {
	let pivot = vals[vals.len() / 2].clone();

	let mut lt = 0;
	let mut eq = 0;
	let mut gt = vals.len() - 1;

	while eq <= gt {
		let ord = call_comparison(vm, by.clone(), vals[eq].clone(), pivot.clone())?;
		match ord.cmp(&Int::ZERO) {
			Ordering::Less => {
				vals.swap(eq, lt);
				lt += 1;
				eq += 1;
			}
			Ordering::Greater => {
				vals.swap(eq, gt);
				gt -= 1;
			}
			Ordering::Equal => {
				eq += 1;
			}
		}
	}

	Ok((lt, gt))
}

fn insertion(vals: &mut [Value]) {
	for i in 0..vals.len() {
		let mut j = i;
		while j > 0 && vals[j - 1] > vals[j] {
			vals.swap(j, j - 1);
			j -= 1;
		}
	}
}

fn insertion_by(vals: &mut [Value], by: &Value, vm: &mut Vm) -> Result<()> {
	for i in 0..vals.len() {
		let mut j = i;
		while j > 0 {
			let ord =
				call_comparison(vm, by.clone(), vals[j - 1].clone(), vals[j].clone())?;
			if !ord.is_positive() {
				break
			}
			vals.swap(j, j - 1);
			j -= 1;
		}
	}
	Ok(())
}

fn sort_inner(vals: &mut [Value], by: Option<&Value>, vm: &mut Vm) -> Result<()> {
	if vals.len() <= 1 {
		return Ok(())
	}
	if vals.len() <= 8 {
		match by {
			Some(by) => insertion_by(vals, by, vm)?,
			None => insertion(vals),
		}
		return Ok(())
	}
	let (lt, gt) = match by {
		Some(by) => partition_by(vals, by, vm)?,
		None => partition(vals),
	};
	sort_inner(&mut vals[..lt], by, vm)?;
	sort_inner(&mut vals[(gt + 1)..], by, vm)
}

#[native_func(1)]
pub fn sort(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, list] = unpack_args!(args);
	let Value::List(list) = list else {
		throw!(*SYM_TYPE_ERROR, "sort expected list, found {list:#}")
	};
	sort_inner(&mut list.borrow_mut(), None, vm)?;
	Ok(list.into())
}

#[native_func(2)]
pub fn sort_by(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, by, list] = unpack_args!(args);
	let Value::List(list) = list else {
		throw!(*SYM_TYPE_ERROR, "sort expected list, found {list:#}")
	};
	sort_inner(&mut list.borrow_mut(), Some(&by), vm)?;
	Ok(list.into())
}

#[native_func(2)]
pub fn sort_key(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, key, list] = unpack_args!(args);
	let Value::List(list) = list else {
		throw!(*SYM_TYPE_ERROR, "sort expected list, found {list:#}")
	};
	let f = move |vm: &mut Vm, args: Vec<Value>| {
		let [_, a, b] = unpack_args!(args);
		let a = vmcall!(vm; key.clone(), a)?;
		let b = vmcall!(vm; key.clone(), b)?;
		match a.partial_cmp(&b) {
			Some(Ordering::Greater) => Ok(Value::from(1)),
			Some(Ordering::Equal) => Ok(Value::from(0)),
			Some(Ordering::Less) => Ok(Value::from(-1)),
			None => throw!(
				*SYM_VALUE_ERROR,
				"values returned from sort key were incomparable"
			),
		}
	};
	let nf = NativeFunc::new(Box::new(f), 2, symbol!("inner[sort_key]")).into();
	sort_inner(&mut list.borrow_mut(), Some(&nf), vm)?;
	Ok(list.into())
}

#[native_func(2)]
pub fn pair(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, x, y] = unpack_args!(args);
	Ok(vec![x, y].into())
}

#[native_func(1)]
pub fn fst(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, l] = unpack_args!(args);
	let Value::List(l) = l else {
		throw!(*SYM_TYPE_ERROR, "fst: expected list")
	};
	let l = l.borrow();
	let [x, _] = l.as_slice() else {
		throw!(*SYM_VALUE_ERROR, "fst: list must have length 2")
	};
	Ok(x.clone())
}

#[native_func(1)]
pub fn snd(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, l] = unpack_args!(args);
	let Value::List(l) = l else {
		throw!(*SYM_TYPE_ERROR, "snd: expected list")
	};
	let l = l.borrow();
	let [_, y] = l.as_slice() else {
		throw!(*SYM_VALUE_ERROR, "snd: list must have length 2")
	};
	Ok(y.clone())
}