talc/talc-std/src/collection.rs

use std::cmp::Ordering;

use talc_lang::{exception::{exception, Result}, symbol::SYM_TYPE_ERROR, throw, value::{function::NativeFunc, Value}, vmcall, Vm};
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());
}

#[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_TYPE_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<i64> {
	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 {
			..=-1 => {
				vals.swap(eq, lt);
				lt += 1;
				eq += 1;
			},
			1.. => {
				vals.swap(eq, gt);
				gt -= 1;
			},
			0 => {
				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 <= 0 {
				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::Int(1)),
			Some(Ordering::Equal) => Ok(Value::Int(0)),
			Some(Ordering::Less) => Ok(Value::Int(-1)),
			None => throw!(*SYM_TYPE_ERROR, "values returned from sort key were incomparable"),
		}
	};
	let nf = NativeFunc::new(Box::new(f), 2).into();
	sort_inner(&mut list.borrow_mut(), Some(&nf), vm)?;
	Ok(list.into())
}