84 lines
2.1 KiB
Rust
84 lines
2.1 KiB
Rust
use num_bigint::RandBigInt;
|
|
use rand::{seq::SliceRandom, Rng};
|
|
use talc_lang::{
|
|
exception::Result,
|
|
number::Int,
|
|
symbol::{SYM_TYPE_ERROR, SYM_VALUE_ERROR},
|
|
throw,
|
|
value::Value,
|
|
vm::Vm,
|
|
vmcalliter,
|
|
};
|
|
use talc_macros::native_func;
|
|
|
|
use crate::unpack_args;
|
|
|
|
pub fn load(vm: &mut Vm) {
|
|
vm.set_global_name("rand", rand().into());
|
|
vm.set_global_name("rand_in", rand_in().into());
|
|
vm.set_global_name("shuf", shuf().into());
|
|
}
|
|
|
|
#[native_func(0)]
|
|
pub fn rand(_: &mut Vm, _: Vec<Value>) -> Result<Value> {
|
|
Ok(Value::Float(rand::thread_rng().gen()))
|
|
}
|
|
|
|
#[native_func(1)]
|
|
pub fn rand_in(vm: &mut Vm, args: Vec<Value>) -> Result<Value> {
|
|
let [_, col] = unpack_args!(args);
|
|
match col {
|
|
Value::List(l) => {
|
|
let l = l.borrow();
|
|
let Some(v) = l.choose(&mut rand::thread_rng()) else {
|
|
throw!(*SYM_VALUE_ERROR, "rand_in: empty list")
|
|
};
|
|
Ok(v.clone())
|
|
}
|
|
Value::Table(t) => {
|
|
let t = t.borrow();
|
|
if t.is_empty() {
|
|
throw!(*SYM_VALUE_ERROR, "rand_in: empty table")
|
|
};
|
|
let i = rand::thread_rng().gen_range(0..t.len());
|
|
let key = t.keys().nth(i).unwrap();
|
|
Ok(key.clone().into_inner())
|
|
}
|
|
Value::Range(r) => {
|
|
if r.is_empty() {
|
|
throw!(*SYM_VALUE_ERROR, "rand_in: empty range")
|
|
}
|
|
let (Some(s), Some(e)) = (&r.start, &r.end) else {
|
|
throw!(*SYM_VALUE_ERROR, "rand_in: infinite range")
|
|
};
|
|
if r.inclusive {
|
|
Ok(Int::from(rand::thread_rng().gen_bigint_range(s, &(e + 1))).into())
|
|
} else {
|
|
Ok(Int::from(rand::thread_rng().gen_bigint_range(s, e)).into())
|
|
}
|
|
}
|
|
col => {
|
|
let iter = col.to_iter_function()?;
|
|
let mut values = Vec::new();
|
|
while let Some(v) = vmcalliter!(vm; iter.clone())? {
|
|
values.push(v);
|
|
}
|
|
if values.is_empty() {
|
|
throw!(*SYM_VALUE_ERROR, "rand_in: empty iterator")
|
|
}
|
|
let i = rand::thread_rng().gen_range(0..values.len());
|
|
let v = values.swap_remove(i);
|
|
Ok(v)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[native_func(1)]
|
|
pub fn shuf(_: &mut Vm, args: Vec<Value>) -> Result<Value> {
|
|
let [_, list] = unpack_args!(args);
|
|
let Value::List(list) = list else {
|
|
throw!(*SYM_TYPE_ERROR, "shuf expected list, got {list:#}")
|
|
};
|
|
list.borrow_mut().shuffle(&mut rand::thread_rng());
|
|
Ok(Value::List(list))
|
|
}
|