talc/talc-std/src/string.rs

use talc_lang::{exception::Result, lformat, lstring::LString, symbol::SYM_TYPE_ERROR, throw, value::Value, Vm};
use talc_macros::native_func;

use crate::unpack_args;

pub fn load(vm: &mut Vm) {
	vm.set_global_name("ord", ord().into());
	vm.set_global_name("chr", chr().into());
	vm.set_global_name("len_bytes", len_bytes().into());
	vm.set_global_name("trim", trim().into());
	vm.set_global_name("upper", upper().into());
	vm.set_global_name("lower", lower().into());
	vm.set_global_name("starts_with", starts_with().into());
	vm.set_global_name("ends_with", ends_with().into());
	vm.set_global_name("is_utf8", is_utf8().into());
	vm.set_global_name("to_utf8", to_utf8().into());
	vm.set_global_name("to_utf8_lossy", to_utf8_lossy().into());
	vm.set_global_name("str_to_bytes", str_to_bytes().into());
	vm.set_global_name("str_of_bytes", str_of_bytes().into());
	vm.set_global_name("format", _format().into());
}

#[native_func(1)]
pub fn ord(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "ord expected string argument, got {s:#}")
	};
	let mut chars = s.chars();
	let Some(c) = chars.next() else {
		throw!(*SYM_TYPE_ERROR, "argument to ord must have length 1")
	};
	if chars.next().is_some() {
		throw!(*SYM_TYPE_ERROR, "argument to ord must have length 1")
	};
	Ok(Value::Int(c as u32 as i64))
}

#[native_func(1)]
pub fn chr(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, i] = unpack_args!(args);
	let Value::Int(i) = i else {
		throw!(*SYM_TYPE_ERROR, "chr expected integer argument, got {i:#}")
	};
	let Ok(i) = u32::try_from(i) else {
		throw!(*SYM_TYPE_ERROR, "argument to chr is not a valid codepoint")
	};
	let Some(c) = char::from_u32(i) else {
		throw!(*SYM_TYPE_ERROR, "argument to chr is not a valid codepoint")
	};
	Ok(Value::String(LString::from(c).into()))
}

#[native_func(1)]
pub fn len_bytes(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "len_bytes expected string argument, got {s:#}")
	};
	Ok(Value::Int(s.len() as i64))
}

#[native_func(1)]
pub fn trim(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "trim expected string argument, got {s:#}")
	};
	Ok(s.trim().into())
}

#[native_func(1)]
pub fn upper(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "upper expected string argument, got {s:#}")
	};
	Ok(s.to_uppercase().into())
}

#[native_func(1)]
pub fn lower(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "lower expected string argument, got {s:#}")
	};
	Ok(s.to_lowercase().into())
}

#[native_func(2)]
pub fn starts_with(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, d, pre] = unpack_args!(args);
	let res = match (d, pre) {
		(Value::List(d), Value::List(pre)) => d.borrow().starts_with(&pre.borrow()),
		(Value::String(d), Value::String(pre)) => d.starts_with(&pre),
		(d, pre) => throw!(*SYM_TYPE_ERROR,
			"starts_with expected two lists or strings, got {d:#} and {pre:#}")
	};
	Ok(res.into())
}

#[native_func(2)]
pub fn ends_with(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, d, suf] = unpack_args!(args);
	let res = match (d, suf) {
		(Value::List(d), Value::List(suf)) => d.borrow().ends_with(&suf.borrow()),
		(Value::String(d), Value::String(suf)) => d.ends_with(&suf),
		(d, suf) => throw!(*SYM_TYPE_ERROR,
			"ends_with expected two lists or strings, got {d:#} and {suf:#}")
	};
	Ok(res.into())
}

#[native_func(1)]
pub fn is_utf8(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "is_utf8 expected string argument, got {s:#}")
	};
	Ok(s.is_utf8().into())
}

#[native_func(1)]
pub fn to_utf8(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "is_utf8 expected string argument, got {s:#}")
	};
	if s.is_utf8() {
		Ok(s.into())
	} else {
		Ok(Value::Nil)
	}
}

#[native_func(1)]
pub fn to_utf8_lossy(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "is_utf8 expected string argument, got {s:#}")
	};
	Ok(s.to_utf8_lossy().into())
}

#[native_func(1)]
pub fn str_to_bytes(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, s] = unpack_args!(args);
	let Value::String(s) = s else {
		throw!(*SYM_TYPE_ERROR, "str_to_bytes expected string argument, got {s:#}")
	};
	Ok(s.as_bytes()
	   .iter()
	   .map(|v| (*v as i64).into())
	   .collect::<Vec<Value>>()
	   .into())
}

#[native_func(1)]
pub fn str_of_bytes(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, b] = unpack_args!(args);
	let Value::List(b) = b else {
		throw!(*SYM_TYPE_ERROR, "str_of_bytes expected list argument, got {b:#}")
	};
	let bytes: Vec<u8> = b.borrow().iter()
		.map(|v| match v {
			Value::Int(i) if (0..=255).contains(i) => Ok(*i as u8),
			_ => throw!(*SYM_TYPE_ERROR, "str_of_bytes expected list of integers in 0..=255"),
		}).collect::<Result<Vec<u8>>>()?;
	Ok(LString::from(bytes).into())
}

#[native_func(2)]
pub fn _format(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
	let [_, fstr, fargs] = unpack_args!(args);
	let (Value::String(fstr), Value::List(fargs)) = (&fstr, &fargs) else {
		throw!(*SYM_TYPE_ERROR, "format expected string and list, got {fstr:#} and {fargs:#}")
	};
	let mut res = LString::new();
	let mut bytes = fstr.bytes();
	let mut faidx = 0;
	while let Some(b) = bytes.next() {
		if b == b'%' {
			match bytes.next() {
				Some(b'%') => res.push_byte(b'%'),
				Some(b @ (b'#' | b'?')) => {
					let fargs = fargs.borrow();
					let Some(a) = fargs.get(faidx) else {
						throw!(*SYM_TYPE_ERROR, "not enough args for format string")
					};
					faidx += 1;
					if b == b'?' {
						res += &lformat!("{a:#}");
					} else {
						res += &lformat!("{a}");
					}

				},
				_ => throw!(*SYM_TYPE_ERROR, "invalid format code")
			}
		} else {
			res.push_byte(b);
		}
	}
	Ok(res.into())
}